abstract_split_read.cpp

/*
 * Copyright 2024-present Alibaba Inc.
 *
 * Licensed 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/core/operation/abstract_split_read.h"

#include <cassert>
#include <cstddef>
#include <map>
#include <set>
#include <utility>

#include "arrow/type.h"
#include "fmt/format.h"
#include "paimon/common/data/blob_utils.h"
#include "paimon/common/data/shredding/map_shared_shredding_file_reader.h"
#include "paimon/common/data/shredding/map_shared_shredding_utils.h"
#include "paimon/common/reader/delegating_prefetch_reader.h"
#include "paimon/common/reader/predicate_batch_reader.h"
#include "paimon/common/reader/prefetch_file_batch_reader_impl.h"
#include "paimon/common/table/special_fields.h"
#include "paimon/common/types/data_field.h"
#include "paimon/common/utils/object_utils.h"
#include "paimon/common/utils/string_utils.h"
#include "paimon/core/io/chain_split_file_path_factory.h"
#include "paimon/core/io/complete_row_tracking_fields_reader.h"
#include "paimon/core/io/data_file_meta.h"
#include "paimon/core/io/data_file_path_factory.h"
#include "paimon/core/io/field_mapping_reader.h"
#include "paimon/core/operation/internal_read_context.h"
#include "paimon/core/partition/partition_info.h"
#include "paimon/core/schema/table_schema.h"
#include "paimon/core/table/source/chain_split_impl.h"
#include "paimon/core/table/source/data_split_impl.h"
#include "paimon/core/utils/branch_manager.h"
#include "paimon/core/utils/field_mapping.h"
#include "paimon/core/utils/nested_projection_utils.h"
#include "paimon/format/file_format.h"
#include "paimon/format/file_format_factory.h"
#include "paimon/fs/file_system.h"
#include "paimon/status.h"

namespace paimon {
class BinaryRow;
class Executor;
class FileStorePathFactory;
class MemoryPool;
class Predicate;

AbstractSplitRead::AbstractSplitRead(const std::shared_ptr<FileStorePathFactory>& path_factory,
                                     const std::shared_ptr<InternalReadContext>& context,
                                     std::unique_ptr<SchemaManager>&& schema_manager,
                                     const std::shared_ptr<MemoryPool>& memory_pool,
                                     const std::shared_ptr<Executor>& executor)
    : pool_(memory_pool),
      executor_(executor),
      path_factory_(path_factory),
      options_(context->GetCoreOptions()),
      raw_read_schema_(context->GetReadSchema()),
      context_(context),
      schema_manager_(std::move(schema_manager)) {}

Result<std::vector<std::unique_ptr<FileBatchReader>>> AbstractSplitRead::CreateRawFileReaders(
    const BinaryRow& partition, const std::vector<std::shared_ptr<DataFileMeta>>& data_files,
    const std::shared_ptr<arrow::Schema>& read_schema, const std::shared_ptr<Predicate>& predicate,
    DeletionVector::Factory dv_factory, const std::optional<std::vector<Range>>& row_ranges,
    const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
    if (data_files.empty()) {
        return std::vector<std::unique_ptr<FileBatchReader>>();
    }
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<FieldMappingBuilder> field_mapping_builder,
        FieldMappingBuilder::Create(read_schema, context_->GetPartitionKeys(), predicate));

    std::vector<std::unique_ptr<FileBatchReader>> raw_file_readers;
    raw_file_readers.reserve(data_files.size());
    for (const auto& file : data_files) {
        auto data_file_path = data_file_path_factory->ToPath(file);
        PAIMON_ASSIGN_OR_RAISE(std::string data_file_identifier, file->FileFormat());
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReaderBuilder> reader_builder,
                               PrepareReaderBuilder(data_file_identifier));
        PAIMON_ASSIGN_OR_RAISE(
            std::unique_ptr<FileBatchReader> file_reader,
            CreateFieldMappingReader(data_file_path, file, partition, reader_builder.get(),
                                     field_mapping_builder.get(), dv_factory, row_ranges,
                                     data_file_path_factory));
        if (file_reader) {
            raw_file_readers.push_back(std::move(file_reader));
        }
    }
    return std::move(raw_file_readers);
}

bool AbstractSplitRead::NeedCompleteRowTrackingFields(
    bool row_tracking_enabled, const std::shared_ptr<arrow::Schema>& read_schema) {
    if (row_tracking_enabled &&
        (read_schema->GetFieldIndex(SpecialFields::RowId().Name()) != -1 ||
         read_schema->GetFieldIndex(SpecialFields::SequenceNumber().Name()) != -1)) {
        return true;
    }
    return false;
}
Result<std::unique_ptr<BatchReader>> AbstractSplitRead::ApplyPredicateFilterIfNeeded(
    std::unique_ptr<BatchReader>&& reader, const std::shared_ptr<Predicate>& predicate) const {
    if (!context_->EnablePredicateFilter() || predicate == nullptr) {
        return std::move(reader);
    }
    return PredicateBatchReader::Create(std::move(reader), predicate, pool_);
}

Result<std::shared_ptr<DataFilePathFactory>> AbstractSplitRead::CreateDataFilePathFactory(
    const std::shared_ptr<DataSplitImpl>& data_split) const {
    auto chain_split = std::dynamic_pointer_cast<ChainSplitImpl>(data_split);
    if (chain_split) {
        return ChainSplitFilePathFactory::Create(chain_split->DataFiles(),
                                                 chain_split->FileBucketPathMapping(),
                                                 chain_split->FileBranchMapping());
    }

    PAIMON_ASSIGN_OR_RAISE(
        std::shared_ptr<DataFilePathFactory> base_factory,
        path_factory_->CreateDataFilePathFactory(data_split->Partition(), data_split->Bucket()));
    return base_factory;
}

Result<const SchemaManager*> AbstractSplitRead::GetSchemaManagerForBranch(
    const std::string& branch) const {
    std::string normalized_branch = BranchManager::NormalizeBranch(branch);
    std::string current_branch = BranchManager::NormalizeBranch(options_.GetBranch());
    if (StringUtils::ToLowerCase(normalized_branch) == StringUtils::ToLowerCase(current_branch)) {
        return schema_manager_.get();
    }

    auto it = branch_schema_managers_.find(normalized_branch);
    if (it != branch_schema_managers_.end()) {
        return it->second.get();
    }

    auto schema_manager = std::make_unique<SchemaManager>(options_.GetFileSystem(),
                                                          context_->GetPath(), normalized_branch);
    auto [inserted_it, _] =
        branch_schema_managers_.emplace(normalized_branch, std::move(schema_manager));
    return inserted_it->second.get();
}

Result<std::shared_ptr<TableSchema>> AbstractSplitRead::ReadDataSchema(
    const std::shared_ptr<DataFileMeta>& file_meta,
    const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
    const SchemaManager* schema_manager = schema_manager_.get();
    bool current_branch = true;

    auto chain_path_factory =
        std::dynamic_pointer_cast<ChainSplitFilePathFactory>(data_file_path_factory);
    if (chain_path_factory) {
        std::optional<std::string> branch = chain_path_factory->BranchForFile(file_meta->file_name);
        if (!branch) {
            return Status::Invalid(
                fmt::format("branch is missing for ChainSplit file {}", file_meta->file_name));
        }
        PAIMON_ASSIGN_OR_RAISE(schema_manager, GetSchemaManagerForBranch(branch.value()));
        current_branch = schema_manager == schema_manager_.get();
    }

    if (current_branch && file_meta->schema_id == context_->GetTableSchema()->Id()) {
        return context_->GetTableSchema();
    }
    return schema_manager->ReadSchema(file_meta->schema_id);
}

Result<std::unique_ptr<ReaderBuilder>> AbstractSplitRead::PrepareReaderBuilder(
    const std::string& format_identifier) const {
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileFormat> file_format,
                           FileFormatFactory::Get(format_identifier, options_.ToMap()));
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<ReaderBuilder> reader_builder,
                           file_format->CreateReaderBuilder(options_.GetReadBatchSize()));
    reader_builder->WithMemoryPool(pool_);
    reader_builder->WithCache(options_.GetCache());
    return reader_builder;
}

Result<std::unique_ptr<FileBatchReader>> AbstractSplitRead::CreateFileBatchReader(
    const std::string& file_format_identifier, const std::string& data_file_path,
    const ReaderBuilder* reader_builder) const {
    if (file_format_identifier == "lance") {
        // lance do not support stream build with input stream
        return reader_builder->Build(data_file_path);
    }
    if (context_->EnablePrefetch() && file_format_identifier != "blob" &&
        file_format_identifier != "avro") {
        PAIMON_ASSIGN_OR_RAISE(
            std::unique_ptr<PrefetchFileBatchReaderImpl> prefetch_reader,
            PrefetchFileBatchReaderImpl::Create(
                data_file_path, reader_builder, options_.GetFileSystem(),
                context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(),
                context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(),
                executor_,
                /*initialize_read_ranges=*/false, context_->GetPrefetchCacheMode(),
                context_->GetCacheConfig(), pool_));
        return std::make_unique<DelegatingPrefetchReader>(std::move(prefetch_reader));
    } else {
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InputStream> input_stream,
                               options_.GetFileSystem()->Open(data_file_path));
        return reader_builder->Build(input_stream);
    }
}

Result<std::unique_ptr<FileBatchReader>> AbstractSplitRead::CreateFieldMappingReader(
    const std::string& data_file_path, const std::shared_ptr<DataFileMeta>& file_meta,
    const BinaryRow& partition, const ReaderBuilder* reader_builder,
    const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory,
    const std::optional<std::vector<Range>>& row_ranges,
    const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<TableSchema> data_schema,
                           ReadDataSchema(file_meta, data_file_path_factory));
    PAIMON_ASSIGN_OR_RAISE(CoreOptions data_options,
                           CoreOptions::FromMap(data_schema->Options(), options_.GetFileSystem()));
    auto blob_inline_fields = data_options.GetBlobInlineFields();

    std::unique_ptr<FieldMapping> field_mapping;
    if (!data_schema->PrimaryKeys().empty()) {
        // for pk table, add special fields to file schema when field mapping
        std::vector<DataField> file_fields = {SpecialFields::SequenceNumber(),
                                              SpecialFields::ValueKind()};
        file_fields.insert(file_fields.end(), data_schema->Fields().begin(),
                           data_schema->Fields().end());
        PAIMON_ASSIGN_OR_RAISE(field_mapping,
                               field_mapping_builder->CreateFieldMapping(file_fields));
    } else {
        PAIMON_ASSIGN_OR_RAISE(
            std::vector<DataField> projected_data_fields,
            ProjectFieldsForRowTrackingAndDataEvolution(data_schema, file_meta->write_cols));
        auto converted_fields =
            BlobUtils::ConvertBlobInlineDataFields(projected_data_fields, blob_inline_fields);
        PAIMON_ASSIGN_OR_RAISE(field_mapping,
                               field_mapping_builder->CreateFieldMapping(converted_fields));
    }

    auto read_schema = DataField::ConvertDataFieldsToArrowSchema(
        field_mapping->non_partition_info.non_partition_data_schema);

    PAIMON_ASSIGN_OR_RAISE(std::string file_format_identifier, file_meta->FileFormat());
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<FileBatchReader> file_reader,
        CreateFileBatchReader(file_format_identifier, data_file_path, reader_builder));
    std::set<int32_t> skip_map_selected_keys_filter_field_ids;
    if (file_format_identifier != "blob") {
        std::pair<std::unique_ptr<FileBatchReader>, std::set<int32_t>> shared_shredding_result;
        PAIMON_ASSIGN_OR_RAISE(shared_shredding_result, ApplySharedShreddingReaderIfNeeded(
                                                            std::move(file_reader), read_schema));
        file_reader = std::move(shared_shredding_result.first);
        skip_map_selected_keys_filter_field_ids = std::move(shared_shredding_result.second);
    }
    if (NeedCompleteRowTrackingFields(options_.RowTrackingEnabled(), read_schema)) {
        file_reader = std::make_unique<CompleteRowTrackingFieldsBatchReader>(
            std::move(file_reader), file_meta->first_row_id, file_meta->max_sequence_number, pool_);
    }
    const auto& predicate = field_mapping->non_partition_info.non_partition_filter;
    auto all_data_schema = DataField::ConvertDataFieldsToArrowSchema(data_schema->Fields());
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileBatchReader> final_reader,
                           ApplyIndexAndDvReaderIfNeeded(
                               std::move(file_reader), file_meta, all_data_schema, read_schema,
                               predicate, dv_factory, row_ranges, data_file_path_factory));
    if (!final_reader) {
        // file is skipped by index or dv
        return std::unique_ptr<FileBatchReader>();
    }

    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<FieldMappingReader> mapping_reader,
        FieldMappingReader::Create(field_mapping_builder->GetReadFieldCount(),
                                   std::move(final_reader), partition, std::move(field_mapping),
                                   std::move(skip_map_selected_keys_filter_field_ids), pool_));
    return mapping_reader;
}

Result<std::pair<std::unique_ptr<FileBatchReader>, std::set<int32_t>>>
AbstractSplitRead::ApplySharedShreddingReaderIfNeeded(
    std::unique_ptr<FileBatchReader>&& file_reader,
    const std::shared_ptr<arrow::Schema>& read_schema) const {
    std::set<int32_t> handled_shared_shredding_field_ids;
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema,
                           file_reader->GetFileSchema());
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> file_arrow_schema,
                                      arrow::ImportSchema(file_schema.get()));
    std::map<std::string, MapSharedShreddingFileReader::SharedShreddingContext>
        shared_shredding_name_to_context;
    for (const auto& read_field : read_schema->fields()) {
        const auto& field_name = read_field->name();
        auto file_field = file_arrow_schema->GetFieldByName(field_name);
        if (!file_field) {
            // may exists field _ROW_ID in read schema
            continue;
        }
        std::shared_ptr<arrow::KeyValueMetadata> metadata =
            std::const_pointer_cast<arrow::KeyValueMetadata>(file_field->metadata());
        if (!MapSharedShreddingUtils::HasShreddingMetadata(metadata)) {
            // not a map shared shredding field
            continue;
        }
        // get meta
        PAIMON_ASSIGN_OR_RAISE(MapSharedShreddingFieldMeta meta,
                               MapSharedShreddingUtils::DeserializeMetadata(
                                   metadata, MapSharedShreddingDefine::kDefaultDictCompression));
        // get selected_keys
        PAIMON_ASSIGN_OR_RAISE(std::vector<std::string> selected_keys,
                               NestedProjectionUtils::GetMapSelectedKeys(read_field));
        if (selected_keys.empty()) {
            // select all keys
            selected_keys.reserve(meta.name_to_id.size());
            for (const auto& [key_name, _] : meta.name_to_id) {
                selected_keys.push_back(key_name);
            }
        }
        // get map type
        auto map_type = arrow::internal::checked_pointer_cast<arrow::MapType>(read_field->type());
        shared_shredding_name_to_context.emplace(
            field_name,
            MapSharedShreddingFileReader::SharedShreddingContext(meta, selected_keys, map_type));
        PAIMON_ASSIGN_OR_RAISE(int32_t field_id,
                               NestedProjectionUtils::GetPaimonFieldId(read_field));
        handled_shared_shredding_field_ids.insert(field_id);
    }
    if (!shared_shredding_name_to_context.empty()) {
        file_reader = std::make_unique<MapSharedShreddingFileReader>(
            std::move(file_reader), std::move(shared_shredding_name_to_context), pool_);
    }
    return std::make_pair(std::move(file_reader), std::move(handled_shared_shredding_field_ids));
}

Result<std::vector<DataField>> AbstractSplitRead::ProjectFieldsForRowTrackingAndDataEvolution(
    const std::shared_ptr<TableSchema>& data_schema,
    const std::optional<std::vector<std::string>>& write_cols) {
    std::vector<DataField> projected_fields;
    const std::vector<std::string>& partition_keys = data_schema->PartitionKeys();
    if (write_cols == std::nullopt) {
        projected_fields = data_schema->Fields();
    } else {
        if (write_cols.value().empty()) {
            return Status::Invalid("write cols cannot be empty");
        }
        for (const auto& write_col : write_cols.value()) {
            if (write_col == SpecialFields::RowId().Name() ||
                write_col == SpecialFields::SequenceNumber().Name()) {
                continue;
            }
            PAIMON_ASSIGN_OR_RAISE(DataField field, data_schema->GetField(write_col));
            projected_fields.push_back(field);
        }
        for (const auto& partition_key : partition_keys) {
            if (!ObjectUtils::Contains(write_cols.value(), partition_key)) {
                PAIMON_ASSIGN_OR_RAISE(DataField partition_field,
                                       data_schema->GetField(partition_key));
                projected_fields.push_back(partition_field);
            }
        }
    }

    projected_fields.push_back(SpecialFields::RowId());
    projected_fields.push_back(SpecialFields::SequenceNumber());
    return projected_fields;
}

}  // namespace paimon