src/paimon/common/utils/binary_row_partition_computer.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/common/utils/binary_row_partition_computer.h"

#include <algorithm>
#include <cstddef>

#include "arrow/type.h"
#include "fmt/format.h"
#include "fmt/ranges.h"
#include "paimon/common/data/binary_row.h"
#include "paimon/common/data/binary_row_writer.h"
#include "paimon/common/utils/string_utils.h"
#include "paimon/macros.h"
#include "paimon/result.h"
#include "paimon/status.h"

namespace paimon {
class MemoryPool;

BinaryRowPartitionComputer::BinaryRowPartitionComputer(
    const std::vector<std::string>& partition_keys, const std::shared_ptr<arrow::Schema>& schema,
    const std::string& default_part_value,
    const std::vector<PartitionConverter>& partition_converters,
    const std::shared_ptr<MemoryPool>& memory_pool)
    : memory_pool_(memory_pool),
      partition_keys_(partition_keys),
      schema_(schema),
      default_part_value_(default_part_value),
      partition_converters_(partition_converters) {}

Result<std::unique_ptr<BinaryRowPartitionComputer>> BinaryRowPartitionComputer::Create(
    const std::vector<std::string>& partition_keys, const std::shared_ptr<arrow::Schema>& schema,
    const std::string& default_part_value, bool legacy_partition_name_enabled,
    const std::shared_ptr<MemoryPool>& memory_pool) {
    if (PAIMON_UNLIKELY(schema == nullptr)) {
        return Status::Invalid(
            "create binary row partition computer failed, schema is null pointer.");
    }
    if (PAIMON_UNLIKELY(memory_pool == nullptr)) {
        return Status::Invalid(
            "create binary row partition computer failed, memory pool is null pointer.");
    }
    std::vector<PartitionConverter> partition_converters;
    for (const auto& partition_key : partition_keys) {
        PAIMON_ASSIGN_OR_RAISE(arrow::Type::type type_id,
                               GetTypeFromArrowSchema(schema, partition_key));
        PAIMON_ASSIGN_OR_RAISE(
            DataConverterUtils::StrToBinaryRowConverter converter,
            DataConverterUtils::CreateDataToBinaryRowConverter(type_id, memory_pool.get()));
        PAIMON_ASSIGN_OR_RAISE(DataConverterUtils::BinaryRowFieldToStrConverter reconverter,
                               DataConverterUtils::CreateBinaryRowFieldToStringConverter(
                                   type_id, legacy_partition_name_enabled));
        partition_converters.emplace_back(partition_key, std::move(converter),
                                          std::move(reconverter));
    }
    return std::unique_ptr<BinaryRowPartitionComputer>(new BinaryRowPartitionComputer(
        partition_keys, schema, default_part_value, partition_converters, memory_pool));
}

Result<BinaryRow> BinaryRowPartitionComputer::ToBinaryRow(
    const std::map<std::string, std::string>& partition) const {
    return ConvertToBinaryRow(partition, /*included_fields=*/nullptr);
}

Result<BinaryRow> BinaryRowPartitionComputer::ConvertToBinaryRow(
    const std::map<std::string, std::string>& partition, std::vector<bool>* included_fields) const {
    BinaryRow binary_row(partition_converters_.size());
    BinaryRowWriter writer(&binary_row, /*initial_size=*/0, memory_pool_.get());
    if (included_fields != nullptr) {
        included_fields->assign(partition_converters_.size(), false);
    }
    for (size_t field_idx = 0; field_idx < partition_converters_.size(); ++field_idx) {
        const PartitionConverter& partition_converter = partition_converters_[field_idx];
        auto input_iter = partition.find(partition_converter.partition_key);
        if (input_iter == partition.end()) {
            if (included_fields != nullptr) {
                writer.SetNullAt(field_idx);
                continue;
            }
            return Status::Invalid(
                fmt::format("can not find partition key '{}' in input partition '{}'",
                            partition_converter.partition_key, partition));
        }
        if (included_fields != nullptr) {
            (*included_fields)[field_idx] = true;
        }
        if (input_iter->second == default_part_value_) {
            // TODO(yonghao.fyh): when support decimal/ timestamp in partition, use
            // WriteTimestamp(null) for non compact precision
            writer.SetNullAt(field_idx);
        } else {
            PAIMON_RETURN_NOT_OK(
                partition_converter.converter(input_iter->second, field_idx, &writer));
        }
    }
    writer.Complete();
    return binary_row;
}

Status BinaryRowPartitionComputer::CheckPartitionKeys(
    const std::map<std::string, std::string>& partition) const {
    for (const auto& [partition_key, _] : partition) {
        if (std::find(partition_keys_.begin(), partition_keys_.end(), partition_key) ==
            partition_keys_.end()) {
            return Status::Invalid(
                fmt::format("field {} does not exist in partition keys", partition_key));
        }
    }
    return Status::OK();
}

Result<BinaryRow> BinaryRowPartitionComputer::ToPartialBinaryRow(
    const std::map<std::string, std::string>& partition) const {
    PAIMON_RETURN_NOT_OK(CheckPartitionKeys(partition));
    std::vector<bool> included_fields;
    return ConvertToBinaryRow(partition, &included_fields);
}

Result<std::map<std::string, std::string>> BinaryRowPartitionComputer::NormalizePartitionSpec(
    const std::map<std::string, std::string>& partition) const {
    PAIMON_RETURN_NOT_OK(CheckPartitionKeys(partition));

    std::vector<bool> included_fields;
    PAIMON_ASSIGN_OR_RAISE(BinaryRow binary_row, ConvertToBinaryRow(partition, &included_fields));

    std::vector<std::pair<std::string, std::string>> normalized_values;
    PAIMON_ASSIGN_OR_RAISE(normalized_values, GeneratePartitionVector(binary_row));
    std::map<std::string, std::string> normalized_partition;
    for (size_t field_idx = 0; field_idx < normalized_values.size(); ++field_idx) {
        if (included_fields[field_idx]) {
            normalized_partition.insert(normalized_values[field_idx]);
        }
    }
    return normalized_partition;
}

Result<std::vector<std::pair<std::string, std::string>>>
BinaryRowPartitionComputer::GeneratePartitionVector(const BinaryRow& partition) const {
    if (static_cast<size_t>(partition.GetFieldCount()) != partition_converters_.size()) {
        return Status::Invalid(fmt::format(
            "partition binary row field count {} not match with partition converter size {}",
            partition.GetFieldCount(), partition_converters_.size()));
    }
    std::vector<std::pair<std::string, std::string>> result;
    for (size_t field_idx = 0; field_idx < partition_converters_.size(); field_idx++) {
        const auto& partition_extractor = partition_converters_[field_idx];
        const auto& partition_key = partition_extractor.partition_key;
        const auto& to_str = partition_extractor.reconverter;
        if (partition.IsNullAt(field_idx)) {
            result.emplace_back(partition_key, default_part_value_);
        } else {
            PAIMON_ASSIGN_OR_RAISE(std::string partition_field_str, to_str(partition, field_idx));
            if (StringUtils::IsNullOrWhitespaceOnly(partition_field_str)) {
                partition_field_str = default_part_value_;
            }
            result.emplace_back(partition_key, partition_field_str);
        }
    }
    return result;
}

Result<arrow::Type::type> BinaryRowPartitionComputer::GetTypeFromArrowSchema(
    const std::shared_ptr<arrow::Schema>& schema, const std::string& field_name) {
    auto field = schema->GetFieldByName(field_name);
    if (field == nullptr) {
        return Status::Invalid(
            fmt::format("field {} not in schema {}", field_name, schema->ToString()));
    }
    return field->type()->id();
}

Result<std::string> BinaryRowPartitionComputer::PartToSimpleString(
    const std::shared_ptr<arrow::Schema>& partition_type, const BinaryRow& partition,
    const std::string& delimiter, int32_t max_length, bool legacy_partition_name_enabled) {
    std::vector<DataConverterUtils::BinaryRowFieldToStrConverter> partition_converters;
    partition_converters.reserve(partition_type->num_fields());
    for (const auto& field : partition_type->fields()) {
        PAIMON_ASSIGN_OR_RAISE(DataConverterUtils::BinaryRowFieldToStrConverter converter,
                               DataConverterUtils::CreateBinaryRowFieldToStringConverter(
                                   field->type()->id(), legacy_partition_name_enabled));
        partition_converters.emplace_back(converter);
    }
    std::vector<std::string> partition_vec;
    partition_vec.reserve(partition_converters.size());
    for (size_t field_idx = 0; field_idx < partition_converters.size(); field_idx++) {
        const auto& to_str = partition_converters[field_idx];
        if (partition.IsNullAt(field_idx)) {
            partition_vec.push_back("null");
        } else {
            PAIMON_ASSIGN_OR_RAISE(std::string partition_field_str, to_str(partition, field_idx));
            partition_vec.push_back(partition_field_str);
        }
    }
    return fmt::format("{}", fmt::join(partition_vec, delimiter)).substr(0, max_length);
}
}  // namespace paimon