src/paimon/core/io/field_mapping_reader.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/core/io/field_mapping_reader.h"

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

#include "arrow/api.h"
#include "arrow/array/array_base.h"
#include "arrow/array/array_nested.h"
#include "arrow/array/util.h"
#include "arrow/c/abi.h"
#include "arrow/c/bridge.h"
#include "arrow/scalar.h"
#include "fmt/format.h"
#include "paimon/common/data/binary_string.h"
#include "paimon/common/types/data_field.h"
#include "paimon/common/utils/arrow/mem_utils.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/checked_cast.h"
#include "paimon/core/casting/cast_executor.h"
#include "paimon/core/casting/casting_utils.h"
#include "paimon/core/utils/field_mapping.h"
#include "paimon/core/utils/nested_projection_utils.h"
#include "paimon/memory/bytes.h"
#include "paimon/reader/batch_reader.h"

namespace paimon {
class MemoryPool;

Result<std::unique_ptr<FieldMappingReader>> FieldMappingReader::Create(
    int32_t field_count, std::unique_ptr<FileBatchReader>&& reader, const BinaryRow& partition,
    std::unique_ptr<FieldMapping>&& mapping,
    std::set<int32_t>&& skip_map_selected_keys_filter_field_ids,
    const std::shared_ptr<arrow::MemoryPool>& arrow_pool) {
    auto mapping_reader = std::unique_ptr<FieldMappingReader>(new FieldMappingReader(
        field_count, std::move(reader), partition, std::move(mapping), arrow_pool));

    mapping_reader->need_mapping_ = false;
    mapping_reader->need_casting_ = false;
    mapping_reader->skip_map_selected_keys_filter_field_ids_ =
        std::move(skip_map_selected_keys_filter_field_ids);

    if (mapping_reader->non_exist_field_info_ != std::nullopt ||
        mapping_reader->partition_info_ != std::nullopt) {
        mapping_reader->need_mapping_ = true;
    }

    for (int32_t i = 0;
         i <
         static_cast<int32_t>(mapping_reader->non_partition_info_.idx_in_target_read_schema.size());
         i++) {
        if (i != mapping_reader->non_partition_info_.idx_in_target_read_schema[i]) {
            mapping_reader->need_mapping_ = true;
        }
        if (mapping_reader->non_partition_info_.cast_executors[i] != nullptr) {
            mapping_reader->need_casting_ = true;
        }
        // A differing nested type needs the AlignArrayToReadType reshape below.
        if (!mapping_reader->non_partition_info_.non_partition_data_schema[i].Type()->Equals(
                *mapping_reader->non_partition_info_.non_partition_read_schema[i].Type())) {
            mapping_reader->need_casting_ = true;
        }
        // Field name change (RENAME COLUMN) also requires mapping: data schema
        // carries the file's physical name while read schema carries the
        // post-rename logical name. If we skipped mapping, the inner reader's
        // batch would be passed through with the old physical name and the
        // consumer's name-based lookup against the read schema would fail.
        if (mapping_reader->non_partition_info_.non_partition_data_schema[i].Name() !=
            mapping_reader->non_partition_info_.non_partition_read_schema[i].Name()) {
            mapping_reader->need_mapping_ = true;
        }
        // Map selected-keys metadata must be validated in Create() (fail-fast).
        // Non-empty selected-keys also requires mapping so that
        // FilterMapArrayBySelectedKeys can filter out unwanted entries.
        PAIMON_ASSIGN_OR_RAISE(
            bool has_map_selected_keys,
            NestedProjectionUtils::HasMapSelectedKeysRecursively(
                mapping_reader->non_partition_info_.non_partition_read_schema[i].ArrowField()));
        if (has_map_selected_keys &&
            mapping_reader->skip_map_selected_keys_filter_field_ids_.count(
                mapping_reader->non_partition_info_.non_partition_read_schema[i].Id()) == 0) {
            mapping_reader->need_mapping_ = true;
        }
    }

    return mapping_reader;
}

FieldMappingReader::FieldMappingReader(int32_t field_count,
                                       std::unique_ptr<FileBatchReader>&& reader,
                                       const BinaryRow& partition,
                                       std::unique_ptr<FieldMapping>&& mapping,
                                       const std::shared_ptr<arrow::MemoryPool>& arrow_pool)
    : field_count_(field_count),
      arrow_pool_(arrow_pool),
      reader_(std::move(reader)),
      partition_(partition),
      partition_info_(mapping->partition_info),
      non_partition_info_(mapping->non_partition_info),
      non_exist_field_info_(mapping->non_exist_field_info) {}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::CastNonPartitionArrayIfNeed(
    const std::shared_ptr<arrow::Array>& src_array) const {
    if (!need_casting_) {
        return src_array;
    }
    auto* struct_array = checked_cast<arrow::StructArray*>(src_array.get());
    int32_t field_count = struct_array->num_fields();
    assert(static_cast<size_t>(field_count) == non_partition_info_.cast_executors.size());
    arrow::ArrayVector casted_array;
    std::vector<std::string> casted_field_names;
    casted_array.reserve(field_count);
    casted_field_names.reserve(field_count);
    for (int32_t i = 0; i < field_count; i++) {
        std::shared_ptr<arrow::Array> column;
        if (non_partition_info_.cast_executors[i] != nullptr) {
            PAIMON_ASSIGN_OR_RAISE(
                std::shared_ptr<arrow::Array> single_column_array,
                CastingUtils::DecodeDictionary(struct_array->field(i), arrow_pool_.get()));
            PAIMON_ASSIGN_OR_RAISE(
                column,
                non_partition_info_.cast_executors[i]->Cast(
                    single_column_array, non_partition_info_.non_partition_read_schema[i].Type(),
                    arrow_pool_.get()));
        } else {
            // The reader may return a dictionary for an unchanged logical type, so reconstruct
            // the struct type from the actual column.
            column = struct_array->field(i);
        }
        // Null-fill nested fields added by schema evolution. Only when the data and
        // read types differ -- the reader may hand back a dictionary-encoded array
        // for an unchanged type, which is not a reshape target.
        if (!non_partition_info_.non_partition_data_schema[i].Type()->Equals(
                *non_partition_info_.non_partition_read_schema[i].Type())) {
            PAIMON_ASSIGN_OR_RAISE(
                column, NestedProjectionUtils::AlignArrayToReadType(
                            column, non_partition_info_.non_partition_read_schema[i].Type(),
                            arrow_pool_.get()));
        }
        casted_array.push_back(column);
        casted_field_names.push_back(non_partition_info_.non_partition_data_schema[i].Name());
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array,
                                      arrow::StructArray::Make(casted_array, casted_field_names));
    return arrow_array;
}

Result<BatchReader::ReadBatchWithBitmap> FieldMappingReader::NextBatchWithBitmap() {
    PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap non_partition_result_with_bitmap,
                           reader_->NextBatchWithBitmap());
    if (!need_mapping_ && !need_casting_) {
        return non_partition_result_with_bitmap;
    }
    if (BatchReader::IsEofBatch(non_partition_result_with_bitmap)) {
        // read finish
        partition_array_.reset();
        non_exist_array_.reset();
        return non_partition_result_with_bitmap;
    }
    auto& [non_partition_result, bitmap] = non_partition_result_with_bitmap;
    auto& [c_array, c_schema] = non_partition_result;
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> non_partition_array,
                                      arrow::ImportArray(c_array.get(), c_schema.get()));

    arrow::ArrayVector target_array(field_count_);
    std::vector<std::string> target_field_names(field_count_);
    // mapping non-partition array
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> casted_non_partition_array,
                           CastNonPartitionArrayIfNeed(non_partition_array));
    PAIMON_RETURN_NOT_OK(MappingFields(
        casted_non_partition_array, non_partition_info_.non_partition_read_schema,
        non_partition_info_.idx_in_target_read_schema, &target_array, &target_field_names));

    // mapping partition array
    if (partition_info_ != std::nullopt) {
        if (!partition_array_ || partition_array_->length() < non_partition_array->length()) {
            PAIMON_ASSIGN_OR_RAISE(partition_array_,
                                   GeneratePartitionArray(non_partition_array->length()));
        }
        auto trim_partition_array = partition_array_->Slice(0, non_partition_array->length());
        PAIMON_RETURN_NOT_OK(MappingFields(
            trim_partition_array, partition_info_.value().partition_read_schema,
            partition_info_.value().idx_in_target_read_schema, &target_array, &target_field_names));
    }
    // mapping non-exist array
    if (non_exist_field_info_ != std::nullopt) {
        if (!non_exist_array_ || non_exist_array_->length() < non_partition_array->length()) {
            PAIMON_ASSIGN_OR_RAISE(non_exist_array_,
                                   GenerateNonExistArray(non_partition_array->length()));
        }
        auto trim_non_exist_array = non_exist_array_->Slice(0, non_partition_array->length());
        PAIMON_RETURN_NOT_OK(MappingFields(trim_non_exist_array,
                                           non_exist_field_info_.value().non_exist_read_schema,
                                           non_exist_field_info_.value().idx_in_target_read_schema,
                                           &target_array, &target_field_names));
    }

    // construct target array
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array,
                                      arrow::StructArray::Make(target_array, target_field_names));
    std::unique_ptr<ArrowArray> target_c_arrow_array = std::make_unique<ArrowArray>();
    std::unique_ptr<ArrowSchema> target_c_schema = std::make_unique<ArrowSchema>();
    PAIMON_RETURN_NOT_OK_FROM_ARROW(
        arrow::ExportArray(*arrow_array, target_c_arrow_array.get(), target_c_schema.get()));
    PAIMON_RETURN_NOT_OK(
        AddArrowArrayLifetime(target_c_arrow_array.get(), target_c_schema.get(), arrow_pool_));
    auto target_batch = std::make_pair(std::move(target_c_arrow_array), std::move(target_c_schema));
    return std::make_pair(std::move(target_batch), std::move(bitmap));
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::GenerateSinglePartitionArray(
    int32_t idx, int32_t batch_size) const {
    const auto& type = partition_info_.value().partition_read_schema[idx].Type();
    if (partition_.IsNullAt(partition_info_.value().idx_in_partition[idx])) {
        // for null partition value
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            std::shared_ptr<arrow::Array> null_array,
            arrow::MakeArrayOfNull(type, batch_size, arrow_pool_.get()));
        return null_array;
    }
    auto type_id = type->id();
    std::shared_ptr<arrow::Scalar> scalar;
    switch (type_id) {
        case arrow::Type::type::BOOL: {
            bool value = partition_.GetBoolean(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::BooleanScalar>(value);
            break;
        }
        case arrow::Type::type::INT8: {
            int8_t value = partition_.GetByte(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int8Scalar>(value);
            break;
        }
        case arrow::Type::type::INT16: {
            int16_t value = partition_.GetShort(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int16Scalar>(value);
            break;
        }
        case arrow::Type::type::INT32: {
            int32_t value = partition_.GetInt(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int32Scalar>(value);
            break;
        }
        case arrow::Type::type::INT64: {
            int64_t value = partition_.GetLong(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Int64Scalar>(value);
            break;
        }
        case arrow::Type::type::STRING: {
            BinaryString value =
                partition_.GetString(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::StringScalar>(value.ToString());
            break;
        }
        case arrow::Type::type::BINARY: {
            auto value = partition_.GetBinary(partition_info_.value().idx_in_partition[idx]);
            std::string value_str(value->data(), value->size());
            scalar = std::make_shared<arrow::BinaryScalar>(value_str);
            break;
        }
        case arrow::Type::type::DATE32: {
            int32_t value = partition_.GetDate(partition_info_.value().idx_in_partition[idx]);
            scalar = std::make_shared<arrow::Date32Scalar>(value);
            break;
        }
        default:
            return Status::Invalid(
                fmt::format("Not support arrow type {} for partition",
                            partition_info_.value().partition_read_schema[idx].Type()->ToString()));
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
        std::shared_ptr<arrow::Array> arrow_array,
        arrow::MakeArrayFromScalar(*scalar, batch_size, arrow_pool_.get()));
    return arrow_array;
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::GeneratePartitionArray(
    int32_t batch_size) const {
    arrow::ArrayVector partition_array;
    std::vector<std::string> partition_field_names;
    partition_array.reserve(partition_info_.value().partition_read_schema.size());
    partition_field_names.reserve(partition_info_.value().partition_read_schema.size());
    for (size_t i = 0; i < partition_info_.value().partition_read_schema.size(); i++) {
        PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> single_partition_array,
                               GenerateSinglePartitionArray(i, batch_size));
        partition_array.push_back(single_partition_array);
        partition_field_names.push_back(partition_info_.value().partition_read_schema[i].Name());
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
        std::shared_ptr<arrow::Array> arrow_array,
        arrow::StructArray::Make(partition_array, partition_field_names));
    return arrow_array;
}

Result<std::shared_ptr<arrow::Array>> FieldMappingReader::GenerateNonExistArray(
    int32_t batch_size) const {
    arrow::ArrayVector non_exist_array;
    std::vector<std::string> non_exist_field_names;
    non_exist_array.reserve(non_exist_field_info_.value().non_exist_read_schema.size());
    non_exist_field_names.reserve(non_exist_field_info_.value().non_exist_read_schema.size());
    for (const auto& non_exist_field : non_exist_field_info_.value().non_exist_read_schema) {
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
            std::shared_ptr<arrow::Array> null_array,
            arrow::MakeArrayOfNull(non_exist_field.Type(), batch_size, arrow_pool_.get()));
        non_exist_array.push_back(null_array);
        non_exist_field_names.push_back(non_exist_field.Name());
    }
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
        std::shared_ptr<arrow::Array> arrow_array,
        arrow::StructArray::Make(non_exist_array, non_exist_field_names));
    return arrow_array;
}

Status FieldMappingReader::MappingFields(const std::shared_ptr<arrow::Array>& data_array,
                                         const std::vector<DataField>& read_fields_of_data_array,
                                         const std::vector<int32_t>& idx_in_target_schema,
                                         arrow::ArrayVector* target_array,
                                         std::vector<std::string>* target_field_names) {
    auto* struct_array = checked_cast<arrow::StructArray*>(data_array.get());
    assert(struct_array->fields().size() == idx_in_target_schema.size());
    for (size_t i = 0; i < idx_in_target_schema.size(); i++) {
        std::shared_ptr<arrow::Array> field_array = struct_array->field(i);
        const DataField& read_field = read_fields_of_data_array[i];

        // Filter map entries by selected keys recursively (supports MAP nested in STRUCT).
        if (skip_map_selected_keys_filter_field_ids_.count(read_field.Id()) == 0) {
            PAIMON_ASSIGN_OR_RAISE(field_array,
                                   NestedProjectionUtils::FilterMapArrayBySelectedKeysRecursively(
                                       field_array, read_field.ArrowField(), arrow_pool_.get()));
        }

        (*target_array)[idx_in_target_schema[i]] = std::move(field_array);
        (*target_field_names)[idx_in_target_schema[i]] = read_field.Name();
    }
    return Status::OK();
}

}  // namespace paimon