src/paimon/core/operation/read_context.cpp

/*
 * 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.
 */

#include "paimon/read_context.h"

#include <utility>

#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
#include "paimon/common/utils/path_util.h"
#include "paimon/core/utils/branch_manager.h"
#include "paimon/executor.h"
#include "paimon/memory/memory_pool.h"
#include "paimon/status.h"

namespace paimon {
class Predicate;

ReadContext::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, 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)
    : path_(path),
      branch_(branch),
      read_field_names_(read_field_names),
      read_field_ids_(read_field_ids),
      predicate_(predicate),
      enable_predicate_filter_(enable_predicate_filter),
      enable_prefetch_(enable_prefetch),
      enable_late_materializing_(enable_late_materializing),
      prefetch_batch_count_(prefetch_batch_count),
      prefetch_max_parallel_num_(prefetch_max_parallel_num),
      enable_multi_thread_row_to_batch_(enable_multi_thread_row_to_batch),
      row_to_batch_thread_number_(row_to_batch_thread_number),
      table_schema_(table_schema),
      memory_pool_(memory_pool),
      executor_(executor),
      specific_file_system_(specific_file_system),
      fs_scheme_to_identifier_map_(fs_scheme_to_identifier_map),
      realtime_context_(realtime_context),
      options_(options),
      read_ahead_cache_enabled_(read_ahead_cache_enabled),
      cache_config_(cache_config),
      cache_(cache) {}

ReadContext::~ReadContext() {
    if (read_schema_ && read_schema_->release) {
        read_schema_->release(read_schema_.get());
    }
}

void ReadContext::SetReadSchema(std::unique_ptr<ArrowSchema> schema) {
    if (schema && schema->release) {
        if (schema.get() == read_schema_.get()) {
            return;
        }
        if (read_schema_ && read_schema_->release) {
            read_schema_->release(read_schema_.get());
        }
        read_schema_ = std::move(schema);
    }
}

class ReadContextBuilder::Impl {
 public:
    friend class ReadContextBuilder;
    void Reset() {
        branch_ = BranchManager::DEFAULT_MAIN_BRANCH;
        read_field_names_.clear();
        read_field_ids_.clear();
        read_schema_.reset();
        fs_scheme_to_identifier_map_.clear();
        options_.clear();
        predicate_.reset();
        enable_predicate_filter_ = false;
        enable_prefetch_ = false;
        enable_late_materializing_ = false;
        read_ahead_cache_enabled_ = true;
        prefetch_batch_count_ = 600;
        prefetch_max_parallel_num_ = 3;
        enable_multi_thread_row_to_batch_ = false;
        row_to_batch_thread_number_ = 1;
        table_schema_ = std::nullopt;
        memory_pool_ = GetDefaultPool();
        executor_.reset();
        specific_file_system_.reset();
        realtime_context_.reset();
        cache_config_ = CacheConfig();
        cache_.reset();
    }

 private:
    std::string path_;
    std::string branch_ = BranchManager::DEFAULT_MAIN_BRANCH;
    std::vector<std::string> read_field_names_;
    std::vector<int32_t> read_field_ids_;
    std::unique_ptr<ArrowSchema> read_schema_;
    std::map<std::string, std::string> fs_scheme_to_identifier_map_;
    std::map<std::string, std::string> options_;
    std::shared_ptr<Predicate> predicate_;
    bool enable_predicate_filter_ = false;
    bool enable_prefetch_ = false;
    bool enable_late_materializing_ = false;
    uint32_t prefetch_batch_count_ = 600;
    uint32_t prefetch_max_parallel_num_ = 3;
    bool enable_multi_thread_row_to_batch_ = false;
    uint32_t row_to_batch_thread_number_ = 1;
    std::optional<std::string> table_schema_;
    std::shared_ptr<MemoryPool> memory_pool_ = GetDefaultPool();
    std::shared_ptr<Executor> executor_;
    std::shared_ptr<FileSystem> specific_file_system_;
    std::shared_ptr<RealtimeContext> realtime_context_;
    bool read_ahead_cache_enabled_ = true;
    CacheConfig cache_config_;
    std::shared_ptr<Cache> cache_;
};

ReadContextBuilder::ReadContextBuilder(const std::string& path)
    : impl_(std::make_unique<ReadContextBuilder::Impl>()) {
    impl_->path_ = path;
}

ReadContextBuilder::~ReadContextBuilder() = default;

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

ReadContextBuilder& ReadContextBuilder::AddOption(const std::string& key,
                                                  const std::string& value) {
    impl_->options_[key] = value;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetOptions(const std::map<std::string, std::string>& opts) {
    impl_->options_ = opts;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetReadFieldNames(
    const std::vector<std::string>& read_field_names) {
    impl_->read_field_names_ = read_field_names;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetReadFieldIds(
    const std::vector<int32_t>& read_field_ids) {
    impl_->read_field_ids_ = read_field_ids;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetReadSchema(std::unique_ptr<ArrowSchema> read_schema) {
    if (read_schema && read_schema->release) {
        impl_->read_schema_ = std::move(read_schema);
    }
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetPredicate(const std::shared_ptr<Predicate>& predicate) {
    impl_->predicate_ = predicate;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::EnablePredicateFilter(bool enabled) {
    impl_->enable_predicate_filter_ = enabled;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::EnablePrefetch(bool enabled) {
    impl_->enable_prefetch_ = enabled;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::EnableLateMaterializing(bool enabled) {
    impl_->enable_late_materializing_ = enabled;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetPrefetchBatchCount(uint32_t batch_count) {
    impl_->prefetch_batch_count_ = batch_count;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetPrefetchMaxParallelNum(uint32_t max_parallel_num) {
    impl_->prefetch_max_parallel_num_ = max_parallel_num;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::EnableMultiThreadRowToBatch(bool enabled) {
    impl_->enable_multi_thread_row_to_batch_ = enabled;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetRowToBatchThreadNumber(uint32_t thread_number) {
    impl_->row_to_batch_thread_number_ = thread_number;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithMemoryPool(
    const std::shared_ptr<MemoryPool>& memory_pool) {
    impl_->memory_pool_ = memory_pool;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithExecutor(const std::shared_ptr<Executor>& executor) {
    impl_->executor_ = executor;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithRealtimeContext(
    const std::shared_ptr<RealtimeContext>& realtime_context) {
    impl_->realtime_context_ = realtime_context;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetTableSchema(const std::string& table_schema) {
    impl_->table_schema_ = table_schema;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithBranch(const std::string& branch) {
    impl_->branch_ = branch;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithFileSystemSchemeToIdentifierMap(
    const std::map<std::string, std::string>& fs_scheme_to_identifier_map) {
    impl_->fs_scheme_to_identifier_map_ = fs_scheme_to_identifier_map;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithFileSystem(
    const std::shared_ptr<FileSystem>& file_system) {
    impl_->specific_file_system_ = file_system;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::SetReadAheadCacheEnabled(bool enabled) {
    impl_->read_ahead_cache_enabled_ = enabled;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithCacheConfig(const CacheConfig& cache_config) {
    impl_->cache_config_ = cache_config;
    return *this;
}

ReadContextBuilder& ReadContextBuilder::WithCache(const std::shared_ptr<Cache>& cache) {
    impl_->cache_ = cache;
    return *this;
}

Result<std::unique_ptr<ReadContext>> ReadContextBuilder::Finish() {
    PAIMON_ASSIGN_OR_RAISE(impl_->path_, PathUtil::NormalizePath(impl_->path_));
    if (impl_->path_.empty()) {
        return Status::Invalid("cannot read with empty table path");
    }
    // The branch names a directory under the table path, so it must stay a single path component.
    PAIMON_RETURN_NOT_OK(BranchManager::CheckValidBranch(impl_->branch_));
    if (impl_->enable_prefetch_ && impl_->prefetch_max_parallel_num_ == 0) {
        return Status::Invalid("prefetch max parallel num should be greater than 0");
    }
    if (impl_->enable_prefetch_ && impl_->prefetch_batch_count_ <= 0) {
        return Status::Invalid("prefetch batch count should be greater than 0");
    }
    if (impl_->enable_prefetch_ &&
        impl_->prefetch_batch_count_ < impl_->prefetch_max_parallel_num_) {
        return Status::Invalid(
            "prefetch batch count should be greater than or equal to prefetch max parallel num");
    }
    if (impl_->specific_file_system_ && !impl_->fs_scheme_to_identifier_map_.empty()) {
        return Status::Invalid(
            "WithFileSystem() and WithFileSystemSchemeToIdentifierMap() cannot be used together");
    }
    if (!impl_->executor_) {
        // If the user do not set executor, create default executor by prefetch batch count
        uint32_t thread_count = impl_->enable_prefetch_ ? impl_->prefetch_max_parallel_num_ : 1;
        PAIMON_ASSIGN_OR_RAISE(impl_->executor_, CreateDefaultExecutor(thread_count));
    }

    if (impl_->enable_multi_thread_row_to_batch_ && impl_->row_to_batch_thread_number_ <= 0) {
        return Status::Invalid("row to batch thread number should be greater than 0");
    }
    auto ctx = std::make_unique<ReadContext>(
        impl_->path_, impl_->branch_, impl_->read_field_names_, impl_->read_field_ids_,
        impl_->predicate_, impl_->enable_predicate_filter_, impl_->enable_prefetch_,
        impl_->enable_late_materializing_, impl_->prefetch_batch_count_,
        impl_->prefetch_max_parallel_num_, impl_->enable_multi_thread_row_to_batch_,
        impl_->row_to_batch_thread_number_, impl_->table_schema_, impl_->memory_pool_,
        impl_->executor_, impl_->specific_file_system_, impl_->fs_scheme_to_identifier_map_,
        impl_->realtime_context_, impl_->options_, impl_->read_ahead_cache_enabled_,
        impl_->cache_config_, impl_->cache_);
    if (impl_->read_schema_ && impl_->read_schema_->release) {
        ctx->SetReadSchema(std::move(impl_->read_schema_));
    }
    impl_->Reset();
    return ctx;
}

}  // namespace paimon