include/paimon/read_context.h

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <cstdint>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "arrow/c/abi.h"
#include "paimon/cache/cache.h"
#include "paimon/catalog/identifier.h"
#include "paimon/predicate/predicate.h"
#include "paimon/result.h"
#include "paimon/type_fwd.h"
#include "paimon/utils/prefetch_cache_config.h"
#include "paimon/visibility.h"

namespace paimon {
class Catalog;
class Executor;
struct FullTextSearch;
class FormatTable;
class MemoryPool;
class Predicate;
class FileSystem;
class RealtimeContext;
struct VectorSearch;

/// `ReadContext` is some configuration for read operations.
///
/// Please do not use this class directly, use `ReadContextBuilder` to build a `ReadContext` which
/// has input validation.
/// @see ReadContextBuilder
class PAIMON_EXPORT ReadContext {
 public:
    ReadContext(
        const std::string& path, const std::string& branch,
        const std::vector<std::string>& read_field_names,
        const std::vector<int32_t>& read_field_ids, const std::shared_ptr<Predicate>& predicate,
        const std::shared_ptr<VectorSearch>& vector_search,
        const std::shared_ptr<FullTextSearch>& full_text_search, bool enable_predicate_filter,
        bool enable_prefetch, bool enable_late_materializing, uint32_t prefetch_batch_count,
        uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch,
        uint32_t row_to_batch_thread_number, const std::optional<std::string>& table_schema,
        const std::shared_ptr<MemoryPool>& memory_pool, const std::shared_ptr<Executor>& executor,
        const std::shared_ptr<FileSystem>& specific_file_system,
        const std::map<std::string, std::string>& fs_scheme_to_identifier_map,
        const std::shared_ptr<RealtimeContext>& realtime_context,
        const std::map<std::string, std::string>& options, bool read_ahead_cache_enabled,
        const CacheConfig& cache_config, const std::shared_ptr<Cache>& cache,
        const std::shared_ptr<FormatTable>& format_table, WarmupLevel warmup_level,
        const std::shared_ptr<Catalog>& catalog, const std::optional<Identifier>& identifier);

    ~ReadContext();

    const std::string& GetPath() const {
        return path_;
    }

    const std::string& GetBranch() const {
        return branch_;
    }

    const std::map<std::string, std::string>& GetFileSystemSchemeToIdentifierMap() const {
        return fs_scheme_to_identifier_map_;
    }

    const std::map<std::string, std::string>& GetOptions() const {
        return options_;
    }

    const std::vector<std::string>& GetReadFieldNames() const {
        return read_field_names_;
    }

    const std::vector<int32_t>& GetReadFieldIds() const {
        return read_field_ids_;
    }

    const std::shared_ptr<Predicate>& GetPredicate() const {
        return predicate_;
    }

    const std::shared_ptr<VectorSearch>& GetVectorSearch() const {
        return vector_search_;
    }

    const std::shared_ptr<FullTextSearch>& GetFullTextSearch() const {
        return full_text_search_;
    }

    /// Whether a root TopK request is configured; full-text predicates are not TopK requests.
    bool HasVectorOrFullTextSearch() const {
        return vector_search_ != nullptr || full_text_search_ != nullptr;
    }

    bool EnablePredicateFilter() const {
        return enable_predicate_filter_;
    }
    bool EnablePrefetch() const {
        return enable_prefetch_;
    }
    bool EnableLateMaterializing() const {
        return enable_late_materializing_;
    }
    uint32_t GetPrefetchBatchCount() const {
        return prefetch_batch_count_;
    }
    uint32_t GetPrefetchMaxParallelNum() const {
        return prefetch_max_parallel_num_;
    }
    bool EnableMultiThreadRowToBatch() const {
        return enable_multi_thread_row_to_batch_;
    }
    uint32_t GetRowToBatchThreadNumber() const {
        return row_to_batch_thread_number_;
    }
    const std::optional<std::string>& GetSpecificTableSchema() const {
        return table_schema_;
    }
    std::shared_ptr<MemoryPool> GetMemoryPool() const {
        return memory_pool_;
    }
    std::shared_ptr<Executor> GetExecutor() const {
        return executor_;
    }
    std::shared_ptr<FileSystem> GetSpecificFileSystem() const {
        return specific_file_system_;
    }

    /// Returns the context used to resolve process-local real-time split tickets.
    std::shared_ptr<RealtimeContext> GetRealtimeContext() const {
        return realtime_context_;
    }

    bool ReadAheadCacheEnabled() const {
        return read_ahead_cache_enabled_;
    }

    const CacheConfig& GetCacheConfig() const {
        return cache_config_;
    }

    std::shared_ptr<Cache> GetCache() const {
        return cache_;
    }

    /// The format table this context was built from, or null when it names a table path and the
    /// schema under that path says what kind of table it is.
    const std::shared_ptr<FormatTable>& GetFormatTable() const {
        return format_table_;
    }

    WarmupLevel GetWarmupLevel() const {
        return warmup_level_;
    }

    /// Returns the catalog supplying table metadata, or null if unset.
    const std::shared_ptr<Catalog>& GetCatalog() const {
        return catalog_;
    }

    /// Returns the table identifier supplied with the catalog.
    const std::optional<Identifier>& GetIdentifier() const {
        return identifier_;
    }

    /// Whether a read schema (C ArrowSchema) for nested column pruning was provided.
    bool HasReadSchema() const {
        return read_schema_ != nullptr && read_schema_->release != nullptr;
    }

    /// Get the read schema as a mutable C ArrowSchema pointer.
    /// ImportSchema will consume (release) the schema content.
    ArrowSchema* GetReadSchema() {
        return read_schema_.get();
    }

    /// Set the read schema from a C ArrowSchema unique_ptr and take ownership of
    /// schema resources (released via ArrowSchema::release in destructor).
    /// Called internally by ReadContextBuilder.
    void SetReadSchema(std::unique_ptr<ArrowSchema> schema);

 private:
    std::string path_;
    std::string branch_;
    std::vector<std::string> read_field_names_;
    std::vector<int32_t> read_field_ids_;
    std::shared_ptr<Predicate> predicate_;
    std::shared_ptr<VectorSearch> vector_search_;
    std::shared_ptr<FullTextSearch> full_text_search_;
    bool enable_predicate_filter_;
    bool enable_prefetch_;
    bool enable_late_materializing_;
    uint32_t prefetch_batch_count_;
    uint32_t prefetch_max_parallel_num_;
    bool enable_multi_thread_row_to_batch_;
    uint32_t row_to_batch_thread_number_;
    std::optional<std::string> table_schema_;
    std::shared_ptr<MemoryPool> memory_pool_;
    std::shared_ptr<Executor> executor_;
    std::shared_ptr<FileSystem> specific_file_system_;
    std::map<std::string, std::string> fs_scheme_to_identifier_map_;
    std::shared_ptr<RealtimeContext> realtime_context_;
    std::map<std::string, std::string> options_;
    bool read_ahead_cache_enabled_;
    CacheConfig cache_config_;
    std::shared_ptr<Cache> cache_;
    std::shared_ptr<FormatTable> format_table_;
    WarmupLevel warmup_level_;
    std::shared_ptr<Catalog> catalog_;
    std::optional<Identifier> identifier_;
    // Owns schema resources and releases ArrowSchema::release in destructor.
    std::unique_ptr<ArrowSchema> read_schema_;
};

/// `ReadContextBuilder` used to build a `ReadContext`, has input validation.
class PAIMON_EXPORT ReadContextBuilder {
 public:
    /// Constructs a `ReadContextBuilder` with required parameters.
    /// @param path The root path of the table.
    explicit ReadContextBuilder(const std::string& path);

    /// Constructs a `ReadContextBuilder` for a format table that is already loaded: the only way
    /// to read one whose schema lives in a metastore rather than under its location, such as a
    /// table a REST catalog serves. The table carries what such a location does not say, so
    /// `SetTableSchema()`, `WithFileSystem()`, `WithFileSystemSchemeToIdentifierMap()` and a
    /// branch are refused here rather than ignored.
    ///
    /// @param table The format table to read, as `Catalog::GetFormatTable()` hands it back.
    explicit ReadContextBuilder(const std::shared_ptr<FormatTable>& table);

    ~ReadContextBuilder();

    ReadContextBuilder(ReadContextBuilder&&) noexcept;
    ReadContextBuilder& operator=(ReadContextBuilder&&) noexcept;

    /// Set the schema fields to read from the table.
    ///
    /// If not set, all fields from the table schema will be read. This is useful for
    /// projection pushdown to reduce I/O and improve performance by reading only
    /// the required columns.
    ///
    /// @param read_field_names Vector of field names to read from the table.
    /// @return Reference to this builder for method chaining.
    /// @note Currently supports top-level field selection. For nested field selection
    ///       use SetReadSchema(std::unique_ptr<ArrowSchema>) instead.
    ReadContextBuilder& SetReadFieldNames(const std::vector<std::string>& read_field_names);
    /// Set the schema fields to read from the table.
    ///
    /// If not set, all fields from the table schema will be read. This is useful for
    /// projection pushdown to reduce I/O and improve performance by reading only
    /// the required columns.
    ///
    /// @param read_field_ids Vector of field ids to read from the table.
    /// @return Reference to this builder for method chaining.
    /// @note Currently supports top-level field selection.
    /// @note SetReadFieldIds() and SetReadFieldNames() are mutually exclusive.
    ///       Calling both will ignore the read schema set by SetReadFieldNames().
    ReadContextBuilder& SetReadFieldIds(const std::vector<int32_t>& read_field_ids);

    /// Set the read Arrow Schema for nested column pruning.
    ///
    /// The read schema is an Arrow C Data Interface schema where STRUCT types
    /// may contain only a subset of the original sub-fields, enabling nested column
    /// pruning to reduce I/O. Field matching is based on field name: the system
    /// looks up each field by name in the table schema and rebuilds the aligned
    /// schema using the table schema's type and metadata. Metadata propagation
    /// from the user-provided schema is whitelist-based: currently only
    /// "paimon.map.selected-keys" is preserved and merged into the final aligned
    /// schema.
    ///
    /// To prune map entries by key, attach metadata "paimon.map.selected-keys"
    /// to the target map field in read schema. The value is a comma-separated
    /// key list, for example: "k1,k2". Only map fields with string key type
    /// (Arrow utf8) are supported.
    ///
    /// Attaching this metadata to a MAP field only filters the returned MAP after
    /// reading. To push down selected keys from a shared-shredding MAP and return
    /// them as STRUCT children, build the field with
    /// `MapSharedShreddingAccessBuilder`. To read selected paths from a VARIANT
    /// field, build the field with `VariantAccessBuilder`.
    ///
    /// Example:
    /// @code{.cpp}
    /// auto map_field = arrow::field("m", arrow::map(arrow::utf8(), arrow::int32()));
    /// auto map_meta = arrow::KeyValueMetadata::Make(
    ///     {"paimon.map.selected-keys"}, {"k1,k2"});
    /// auto projected_schema = arrow::schema({
    ///     arrow::field("id", arrow::int64()),
    ///     map_field->WithMetadata(map_meta),
    /// });
    ///
    /// auto c_schema = std::make_unique<ArrowSchema>();
    /// arrow::Status export_status = arrow::ExportSchema(*projected_schema, c_schema.get());
    /// if (!export_status.ok()) {
    ///     return paimon::Status::Invalid(export_status.message());
    /// }
    ///
    /// ReadContextBuilder builder("/path/to/table");
    /// builder.SetReadSchema(std::move(c_schema));
    /// @endcode
    ///
    /// @param read_schema Arrow C Schema. Ownership of schema resources is transferred
    ///        to the built ReadContext.
    /// @return Reference to this builder for method chaining.
    /// @note Priority: read_schema > read_field_ids > read_field_names.
    ///       When set, read_field_ids and read_field_names are ignored.
    ReadContextBuilder& SetReadSchema(std::unique_ptr<ArrowSchema> read_schema);

    /// Set a configuration options map to set some option entries which are not defined in the
    /// table schema or whose values you want to overwrite.
    /// @note The options map will clear the options added by `AddOption()` before.
    /// @param options The configuration options map.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetOptions(const std::map<std::string, std::string>& options);

    /// Add a single configuration option which is not defined in the table schema or whose value
    /// you want to overwrite.
    ///
    /// If you want to add multiple options, call `AddOption()` multiple times or use `SetOptions()`
    /// instead.
    /// @param key The option key.
    /// @param value The option value.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& AddOption(const std::string& key, const std::string& value);

    /// Set a predicate for filtering data during reading.
    ///
    /// The predicate is used for both partition pruning and data filtering.
    /// It can significantly improve performance by reducing the amount of data
    /// that needs to be read and processed.
    ///
    /// All referenced fields, including full-text match fields, must be projected in the read
    /// schema.
    /// TODO(lisizhuo.lsz): The behavior here will be aligned with Java Paimon in a follow-up
    /// change: fields referenced in predicates will no longer be required to appear in the
    /// projection. Full-text matches can be combined with scalar predicates using
    /// PredicateBuilder::And/Or. Full-text predicate execution is not implemented yet; reading such
    /// a predicate returns NotImplemented.
    ///
    /// @param predicate Shared pointer to the predicate for data filtering.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetPredicate(const std::shared_ptr<Predicate>& predicate);

    /// Configure the root vector TopK request, logically applied after the predicate filter.
    /// The search field need not be projected. SetVectorSearch is mutually exclusive with
    /// SetFullTextSearch and must have a positive limit.
    /// Scores are required internally for TopK. Project `_<field>_INDEX_SCORE` to include them
    /// in the result. Per-field score output is not implemented yet; reading such a projection
    /// returns NotImplemented.
    /// TODO(lisizhuo.lsz): Execute combined read predicates and merge TopK across files.
    /// @param vector_search Vector search request, or nullptr to clear it.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetVectorSearch(const std::shared_ptr<VectorSearch>& vector_search);

    /// Configure the root full-text TopK request, logically applied after the predicate filter.
    /// The search field need not be projected. Use PredicateBuilder::FullTextMatch for a
    /// full-text filter without TopK truncation. SetFullTextSearch is mutually exclusive with
    /// SetVectorSearch and must have a positive limit.
    /// Scores are required internally for TopK. Project `_<field>_INDEX_SCORE` to include them
    /// in the result. Per-field score output is not implemented yet; reading such a projection
    /// returns NotImplemented.
    /// TODO(lisizhuo.lsz): Execute combined read predicates and merge TopK across files.
    /// @param full_text_search Full-text search request, or nullptr to clear it.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetFullTextSearch(const std::shared_ptr<FullTextSearch>& full_text_search);

    /// Whether to perform precise filtering according to predicates for data read from format
    /// reader.
    /// @param enabled Whether to enable precise filtering (default: false)
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& EnablePredicateFilter(bool enabled);

    /// Enable or disable prefetching of data batches from individual files.
    ///
    /// When enabled, the reader will prefetch multiple batches in parallel to
    /// improve throughput by overlapping I/O with computation. This is particularly
    /// beneficial for high-latency storage systems.
    ///
    /// @param enabled Whether to enable prefetching (default: false)
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& EnablePrefetch(bool enabled);

    /// Enable or disable late materialization (probe/payload two-phase reads). When enabled,
    /// each parallel reader under the prefetch layer performs a probe read of predicate
    /// columns first and only materializes payload columns for matched rows.
    /// @param enabled Whether to enable late materialization (default: false)
    /// @return Reference to this builder for method chaining.
    /// @note Without a pushed-down predicate the late-materializing reader degrades to a
    /// plain passthrough.
    ReadContextBuilder& EnableLateMaterializing(bool enabled);

    /// Enable or disable the read-ahead cache for read operations.
    ///
    /// A read-ahead cache is used to prebuffer data ranges before they are needed,
    /// which can improve read performance by reducing redundant I/O operations.
    /// @param enabled Whether to enable the read-ahead cache (default: true)
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetReadAheadCacheEnabled(bool enabled);

    /// Set the cache configuration for prefetch read operations.
    ///
    /// @param config The cache configuration to use.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& WithCacheConfig(const CacheConfig& config);

    /// Set how far the reader prepares the next file before it is read.
    ///
    /// Warmup overlaps remote-storage latency with the read of the current file. Higher levels hide
    /// more latency but use more memory, and may warm files that a query never reads (for example
    /// when a LIMIT stops the scan early).
    /// @param level The warmup level to use (default: WarmupLevel::RAW).
    /// @return Reference to this builder for method chaining.
    /// @note WarmupLevel::RAW warms the read-ahead cache, so it has no effect and behaves like
    /// WarmupLevel::NONE when the cache is off (see SetReadAheadCacheEnabled()).
    /// WarmupLevel::DECODED does not depend on the cache.
    /// @see WarmupLevel
    ReadContextBuilder& SetWarmupLevel(WarmupLevel level);

    /// Set the total number of batches to prefetch across all files.
    ///
    /// This controls the memory usage and parallelism of the prefetching mechanism.
    /// Higher values can improve throughput but consume more memory.
    ///
    /// @param batch_count Total number of batches to prefetch (default: 600)
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetPrefetchBatchCount(uint32_t batch_count);

    /// Set the maximum number of parallel prefetch operations.
    ///
    /// This limits the number of concurrent I/O operations to prevent overwhelming
    /// the storage system or consuming excessive system resources.
    ///
    /// @param parallel_num Maximum parallel prefetch operations (default: 3)
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& SetPrefetchMaxParallelNum(uint32_t parallel_num);

    /// Enable or disable multi-threaded row-to-batch conversion in merge-on-read scenarios.
    ///
    /// When enabled, multiple threads are used to convert row data to batch format
    /// during merge operations, which can improve performance for CPU-intensive
    /// merge operations.
    ///
    /// @param enabled Whether to enable multi-threaded conversion (default: false)
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& EnableMultiThreadRowToBatch(bool enabled);

    /// Set the number of threads for row-to-batch conversion in merge-on-read scenarios.
    ///
    /// This controls the parallelism of row-to-batch conversion during merge operations.
    /// Higher values can improve performance but may affect result ordering.
    ///
    /// @param thread_number Number of conversion threads (default: 1)
    /// @return Reference to this builder for method chaining.
    /// @note If thread_number > 1, Arrow batches from the reader may not be in primary key order.
    ReadContextBuilder& SetRowToBatchThreadNumber(uint32_t thread_number);

    /// Set custom memory pool for memory management.
    /// @param memory_pool The memory pool to use.
    /// @return Reference to this builder for method chaining.
    /// @note If not set, the default system memory pool will be used.
    ReadContextBuilder& WithMemoryPool(const std::shared_ptr<MemoryPool>& memory_pool);

    /// Set custom executor for task execution.
    /// @param executor The executor to use.
    /// @return Reference to this builder for method chaining.
    /// @note If not set, the default system executor will be used.
    ReadContextBuilder& WithExecutor(const std::shared_ptr<Executor>& executor);

    /// Sets the shared context used to resolve process-local real-time split tickets.
    /// @param realtime_context Context shared with the scan that created the real-time splits.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& WithRealtimeContext(
        const std::shared_ptr<RealtimeContext>& realtime_context);

    /// Set the table schema as a string to avoid schema loading I/O operations.
    ///
    /// This optimization allows the reader to use a pre-loaded schema instead of
    /// reading it from the table metadata, which can improve performance especially
    /// in scenarios with many small read operations.
    ///
    /// @param table_schema String representation of the table schema.
    /// @return Reference to this builder for method chaining.
    /// @note The user must ensure that the schema string is valid and matches the table.
    /// @note If not set, the schema will be loaded from the table path.
    ReadContextBuilder& SetTableSchema(const std::string& table_schema);

    /// Set the specific branch to read from in a versioned table.
    ///
    /// Paimon supports branching for data versioning and time travel queries.
    /// This method allows reading from a specific branch instead of the main branch.
    ///
    /// The `branch` option names the branch too, as it does for a scan of one branch; naming two
    /// different branches is refused rather than silently resolved.
    ///
    /// @param branch Name of the branch to read from.
    /// @return Reference to this builder for method chaining.
    /// @note Default branch is "main" if not specified. An empty name is the main branch.
    ReadContextBuilder& WithBranch(const std::string& branch);

    /// Sets a mapping from URI schemes (e.g., "file", "oss") to registered file system
    /// identifiers. This allows selecting different pre-registered file system implementations
    /// based on the URI scheme at runtime.
    ///
    /// @param fs_scheme_to_identifier_map Map from URI scheme (like "oss") to the corresponding
    /// file system identifier.
    /// @return Reference to this builder for method chaining.
    /// @note
    ///   - This method is intended for environments where multiple file systems are pre-registered.
    ///   - The specified identifiers must correspond to file systems that have been registered at
    ///   compile time or initialization.
    ///   - Cannot be used together with `WithFileSystem()`.
    ///   - If not set, use default file system (configured in `Options::FILE_SYSTEM`).
    /// Example:
    ///   builder.WithFileSystemSchemeToIdentifierMap({{"oss", "jindo"}, {"file", "local"}});
    ///
    ReadContextBuilder& WithFileSystemSchemeToIdentifierMap(
        const std::map<std::string, std::string>& fs_scheme_to_identifier_map);

    /// Sets a custom file system instance to be used for all file operations in this read context.
    /// This bypasses the global file system registry and uses the provided implementation directly.
    ///
    /// @param file_system The file system to use.
    /// @return Reference to this builder for method chaining.
    /// @note If not set, use default file system (configured in `Options::FILE_SYSTEM`)
    ReadContextBuilder& WithFileSystem(const std::shared_ptr<FileSystem>& file_system);

    /// Reads a native table through its own file system - including the per-table temporary
    /// credentials a catalog that issues them hands out through `Catalog::GetTableFileSystem`.
    /// This is a shorthand for `WithFileSystem(catalog->GetTableFileSystem(identifier))`; an
    /// explicit `WithFileSystem()` takes precedence, so the catalog is not asked.
    /// @param catalog Non-null catalog, read when `Finish()` builds the context.
    /// @param identifier The native table to read.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& WithCatalog(const std::shared_ptr<Catalog>& catalog,
                                    const Identifier& identifier);

    /// Inject a cache for read operations. Passing nullptr disables cache.
    /// @return Reference to this builder for method chaining.
    ReadContextBuilder& WithCache(const std::shared_ptr<Cache>& cache);

    /// Build and return a `ReadContext` instance with input validation.
    /// @return Result containing the constructed `ReadContext` or an error status.
    Result<std::unique_ptr<ReadContext>> Finish();

 private:
    class Impl;

    std::unique_ptr<Impl> impl_;
};
}  // namespace paimon