src/paimon/common/predicate/literal_converter.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/predicate/literal_converter.h"

#include <optional>

#include "arrow/array/array_base.h"
#include "arrow/array/array_binary.h"
#include "arrow/array/array_decimal.h"
#include "arrow/array/array_dict.h"
#include "arrow/array/array_primitive.h"
#include "arrow/array/builder_binary.h"
#include "arrow/array/builder_decimal.h"
#include "arrow/array/builder_primitive.h"
#include "arrow/array/builder_time.h"
#include "arrow/memory_pool.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
#include "arrow/util/decimal.h"
#include "fmt/format.h"
#include "paimon/common/data/binary_string.h"
#include "paimon/common/data/internal_row.h"
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/checked_cast.h"
#include "paimon/common/utils/date_time_utils.h"
#include "paimon/common/utils/field_type_utils.h"
#include "paimon/common/utils/string_utils.h"
#include "paimon/data/decimal.h"
#include "paimon/data/timestamp.h"
#include "paimon/defs.h"
#include "paimon/memory/bytes.h"
#include "paimon/status.h"

namespace paimon {
Result<Literal> LiteralConverter::ConvertLiteralsFromString(const FieldType& type,
                                                            const std::string& value_str) {
    switch (type) {
        case FieldType::BOOLEAN: {
            auto value = StringUtils::StringToValue<bool>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(fmt::format("cannot convert {} to BOOLEAN type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::TINYINT: {
            auto value = StringUtils::StringToValue<int8_t>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(fmt::format("cannot convert {} to TINYINT type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::SMALLINT: {
            auto value = StringUtils::StringToValue<int16_t>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(
                    fmt::format("cannot convert {} to SMALLINT type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::INT: {
            auto value = StringUtils::StringToValue<int32_t>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(fmt::format("cannot convert {} to INT type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::DATE: {
            PAIMON_ASSIGN_OR_RAISE(int32_t value, StringUtils::StringToDate(value_str));
            return Literal(FieldType::DATE, value);
        }
        case FieldType::BIGINT: {
            auto value = StringUtils::StringToValue<int64_t>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(fmt::format("cannot convert {} to BIGINT type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::FLOAT: {
            auto value = StringUtils::StringToValue<float>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(fmt::format("cannot convert {} to FLOAT type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::DOUBLE: {
            auto value = StringUtils::StringToValue<double>(value_str);
            if (value == std::nullopt) {
                return Status::Invalid(fmt::format("cannot convert {} to DOUBLE type", value_str));
            }
            return Literal(value.value());
        }
        case FieldType::STRING:
        case FieldType::BINARY:
            return Literal(type, value_str.data(), value_str.size());
        default:
            return Status::Invalid(
                fmt::format("Do not support type {} in ConvertLiteralsFromString",
                            FieldTypeUtils::FieldTypeToString(type)));
    }
}

Result<Literal> LiteralConverter::ConvertLiteralsFromRow(
    const std::shared_ptr<arrow::Schema>& schema, const InternalRow& row, int32_t field_idx,
    const FieldType& type) {
    if (row.IsNullAt(field_idx)) {
        return Literal(type);
    }
    switch (type) {
        case FieldType::BOOLEAN:
            return Literal(row.GetBoolean(field_idx));
        case FieldType::TINYINT:
            return Literal(static_cast<int8_t>(row.GetByte(field_idx)));
        case FieldType::SMALLINT:
            return Literal(row.GetShort(field_idx));
        case FieldType::INT:
            return Literal(row.GetInt(field_idx));
        case FieldType::BIGINT:
            return Literal(row.GetLong(field_idx));
        case FieldType::FLOAT:
            return Literal(row.GetFloat(field_idx));
        case FieldType::DOUBLE:
            return Literal(row.GetDouble(field_idx));
        case FieldType::STRING: {
            std::string field = row.GetString(field_idx).ToString();
            return Literal(type, field.data(), field.size());
        }
        case FieldType::BINARY: {
            auto field = row.GetBinary(field_idx);
            return Literal(type, field->data(), field->size());
        }
        case FieldType::TIMESTAMP: {
            const std::shared_ptr<arrow::DataType>& field_type = schema->field(field_idx)->type();
            if (!field_type || field_type->id() != arrow::Type::TIMESTAMP) {
                return Status::Invalid(
                    fmt::format("Convert literal from row not valid for schema {}, field_idx {}",
                                schema->ToString(), field_idx));
            }
            auto timestamp_type = checked_pointer_cast<arrow::TimestampType>(field_type);
            int32_t precision = DateTimeUtils::GetPrecisionFromType(timestamp_type);
            Timestamp field = row.GetTimestamp(field_idx, precision);
            return Literal(field);
        }
        case FieldType::DECIMAL: {
            const std::shared_ptr<arrow::DataType>& field_type = schema->field(field_idx)->type();
            if (!field_type || field_type->id() != arrow::Type::DECIMAL128) {
                return Status::Invalid(
                    fmt::format("Convert literal from row not valid for schema {}, field_idx {}",
                                schema->ToString(), field_idx));
            }
            auto* decimal_type = checked_cast<arrow::Decimal128Type*>(field_type.get());
            auto precision = decimal_type->precision();
            auto scale = decimal_type->scale();
            Decimal field = row.GetDecimal(field_idx, precision, scale);
            return Literal(field);
        }
        case FieldType::DATE:
            return Literal(FieldType::DATE, row.GetInt(field_idx));
        case FieldType::ARRAY:
        case FieldType::VECTOR:
        case FieldType::MAP:
        case FieldType::STRUCT:
        default:
            return Status::Invalid(fmt::format("Convert literal from row not valid for {}",
                                               FieldTypeUtils::FieldTypeToString(type)));
    }
}

Result<std::vector<Literal>> LiteralConverter::ConvertLiteralsFromArray(const arrow::Array& array,
                                                                        bool own_data) {
    const auto kind = array.type_id();
    switch (kind) {
        case arrow::Type::type::BOOL:
            return GetLiteralFromGenericArray<arrow::BooleanType>(array, FieldType::BOOLEAN);
        case arrow::Type::type::INT8:
            return GetLiteralFromGenericArray<arrow::Int8Type>(array, FieldType::TINYINT);
        case arrow::Type::type::INT16:
            return GetLiteralFromGenericArray<arrow::Int16Type>(array, FieldType::SMALLINT);
        case arrow::Type::type::INT32:
            return GetLiteralFromGenericArray<arrow::Int32Type>(array, FieldType::INT);
        case arrow::Type::type::INT64:
            return GetLiteralFromGenericArray<arrow::Int64Type>(array, FieldType::BIGINT);
        case arrow::Type::type::FLOAT:
            return GetLiteralFromGenericArray<arrow::FloatType>(array, FieldType::FLOAT);
        case arrow::Type::type::DOUBLE:
            return GetLiteralFromGenericArray<arrow::DoubleType>(array, FieldType::DOUBLE);
        case arrow::Type::type::STRING:
            return GetLiteralFromStringArray<arrow::StringType>(array, FieldType::STRING, own_data);
        case arrow::Type::type::BINARY:
            return GetLiteralFromStringArray<arrow::BinaryType>(array, FieldType::BINARY, own_data);
        case arrow::Type::type::TIMESTAMP:
            return GetLiteralFromTimestampArray(array);
        case arrow::Type::type::DECIMAL128:
            return GetLiteralFromDecimalArray(array);
        case arrow::Type::type::DATE32:
            return GetLiteralFromDateArray(array);
        case arrow::Type::type::DICTIONARY: {
            const auto& dict_array = checked_cast<const arrow::DictionaryArray&>(array);
            auto* dict_type = checked_cast<arrow::DictionaryType*>(dict_array.type().get());
            auto value_type_id = dict_type->value_type()->id();
            auto index_type_id = dict_type->index_type()->id();
            if ((value_type_id == arrow::Type::type::STRING ||
                 value_type_id == arrow::Type::type::BINARY) &&
                index_type_id == arrow::Type::type::INT32) {
                FieldType literal_type = value_type_id == arrow::Type::type::STRING
                                             ? FieldType::STRING
                                             : FieldType::BINARY;
                return GetLiteralFromDictionaryArray<arrow::BinaryArray, arrow::Int32Array>(
                    dict_array, literal_type, own_data);
            } else if ((value_type_id == arrow::Type::type::LARGE_STRING ||
                        value_type_id == arrow::Type::type::LARGE_BINARY) &&
                       index_type_id == arrow::Type::type::INT64) {
                FieldType literal_type = value_type_id == arrow::Type::type::LARGE_STRING
                                             ? FieldType::STRING
                                             : FieldType::BINARY;
                return GetLiteralFromDictionaryArray<arrow::LargeBinaryArray, arrow::Int64Array>(
                    dict_array, literal_type, own_data);
            } else {
                return Status::Invalid(
                    "only support [STRING|BINARY, INT32] or "
                    "[LARGE_STRING|LARGE_BINARY, INT64] for DictionaryArray");
            }
        }
        default:
            return Status::Invalid(
                fmt::format("Not support literal on arrow {} type", array.type()->ToString()));
    }
}

std::vector<Literal> LiteralConverter::GetLiteralFromDecimalArray(const arrow::Array& array) {
    using ArrayType = typename arrow::TypeTraits<arrow::Decimal128Type>::ArrayType;
    const auto& array_(checked_cast<const ArrayType&>(array));
    auto* arrow_type = checked_cast<arrow::Decimal128Type*>(array.type().get());
    int32_t precision = arrow_type->precision();
    int32_t scale = arrow_type->scale();
    std::vector<Literal> literals;
    literals.reserve(array_.length());
    for (int64_t i = 0; i < array_.length(); i++) {
        if (array_.IsNull(i)) {
            literals.emplace_back(FieldType::DECIMAL);
        } else {
            const arrow::Decimal128 decimal(array_.GetValue(i));
            auto value = static_cast<Decimal::int128_t>(
                static_cast<Decimal::uint128_t>(static_cast<uint64_t>(decimal.high_bits())) << 64 |
                decimal.low_bits());
            literals.emplace_back(Decimal(precision, scale, value));
        }
    }
    return literals;
}

std::vector<Literal> LiteralConverter::GetLiteralFromDateArray(const arrow::Array& array) {
    using ArrayType = typename arrow::TypeTraits<arrow::Date32Type>::ArrayType;
    const auto& array_(checked_cast<const ArrayType&>(array));
    std::vector<Literal> literals;
    literals.reserve(array_.length());
    for (int64_t i = 0; i < array_.length(); i++) {
        if (array_.IsNull(i)) {
            literals.emplace_back(FieldType::DATE);
        } else {
            literals.emplace_back(FieldType::DATE, array_.Value(i));
        }
    }
    return literals;
}

std::vector<Literal> LiteralConverter::GetLiteralFromTimestampArray(const arrow::Array& array) {
    using ArrayType = typename arrow::TypeTraits<arrow::TimestampType>::ArrayType;
    const auto& array_(checked_cast<const ArrayType&>(array));
    auto timestamp_type = checked_pointer_cast<arrow::TimestampType>(array_.type());
    DateTimeUtils::TimeType time_type = DateTimeUtils::GetTimeTypeFromArrowType(timestamp_type);
    std::vector<Literal> literals;
    literals.reserve(array_.length());
    for (int64_t i = 0; i < array_.length(); i++) {
        if (array_.IsNull(i)) {
            literals.emplace_back(FieldType::TIMESTAMP);
        } else {
            int64_t data = array_.Value(i);
            auto [milli, nano] = DateTimeUtils::TimestampConverter(
                data, time_type, DateTimeUtils::TimeType::MILLISECOND,
                DateTimeUtils::TimeType::NANOSECOND);
            literals.emplace_back(Timestamp(milli, nano));
        }
    }
    return literals;
}

arrow::TimeUnit::type LiteralConverter::MinRequiredTimeUnit(const std::vector<Literal>& literals) {
    bool needs_micro = false;
    for (const auto& literal : literals) {
        if (literal.IsNull()) {
            continue;
        }
        const int32_t nano = literal.GetValue<Timestamp>().GetNanoOfMillisecond();
        if (nano % 1000 != 0) {
            return arrow::TimeUnit::NANO;
        }
        needs_micro = needs_micro || nano != 0;
    }
    return needs_micro ? arrow::TimeUnit::MICRO : arrow::TimeUnit::MILLI;
}

namespace {
// Appends every literal, read by `extract`, to a `BuilderType`, keeping a null literal a null.
template <typename BuilderType, typename Extract>
Result<std::shared_ptr<arrow::Array>> BuildArray(const std::vector<Literal>& literals,
                                                 Extract extract, arrow::MemoryPool* pool) {
    BuilderType builder(pool);
    PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast<int64_t>(literals.size())));
    for (const auto& literal : literals) {
        if (literal.IsNull()) {
            PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull());
            continue;
        }
        PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(extract(literal)));
    }
    std::shared_ptr<arrow::Array> array;
    PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array));
    return array;
}

// Writes the literals to a decimal array of the precision and the scale they carry themselves. One
// array holds one arrow type, so a non null literal is needed to settle it, and every non null
// literal has to carry the same pair: rescaling a value to another scale loses digits or overflows,
// which is not something a conversion decides on its own.
Result<std::shared_ptr<arrow::Array>> BuildDecimalArray(const std::vector<Literal>& literals,
                                                        arrow::MemoryPool* pool) {
    std::optional<Decimal> typed_value;
    for (const auto& literal : literals) {
        if (!literal.IsNull()) {
            typed_value = literal.GetValue<Decimal>();
            break;
        }
    }
    if (typed_value == std::nullopt) {
        return Status::Invalid(
            "Not support converting literals of DECIMAL type to an arrow array without a non null "
            "literal to take the precision and the scale from");
    }
    int32_t precision = typed_value->Precision();
    int32_t scale = typed_value->Scale();
    // `arrow::decimal128` checks the precision fatally, `Make` reports it instead.
    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::DataType> type,
                                      arrow::Decimal128Type::Make(precision, scale));
    arrow::Decimal128Builder builder(type, pool);
    PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast<int64_t>(literals.size())));
    for (const auto& literal : literals) {
        if (literal.IsNull()) {
            PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull());
            continue;
        }
        auto value = literal.GetValue<Decimal>();
        if (value.Precision() != precision || value.Scale() != scale) {
            return Status::Invalid(fmt::format(
                "Not support converting literals of DECIMAL type to an arrow array, {} and {} do "
                "not share one precision and scale",
                typed_value->ToString(), value.ToString()));
        }
        PAIMON_RETURN_NOT_OK_FROM_ARROW(
            builder.Append(arrow::Decimal128(value.HighBits(), value.LowBits())));
    }
    std::shared_ptr<arrow::Array> array;
    PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array));
    return array;
}

// Writes the literals to a timestamp array of the finest time unit that keeps every value, one
// `is_in` can only compare against a column of that very unit. One array holds one arrow type, so
// a non null literal is needed to settle the unit.
Result<std::shared_ptr<arrow::Array>> BuildTimestampArray(const std::vector<Literal>& literals,
                                                          arrow::MemoryPool* pool) {
    bool has_value = false;
    for (const auto& literal : literals) {
        if (!literal.IsNull()) {
            has_value = true;
            break;
        }
    }
    if (!has_value) {
        return Status::Invalid(
            "Not support converting literals of TIMESTAMP type to an arrow array without a non "
            "null literal to take the time unit from");
    }
    arrow::TimestampBuilder builder(
        arrow::timestamp(LiteralConverter::MinRequiredTimeUnit(literals)), pool);
    const DateTimeUtils::TimeType time_type = DateTimeUtils::GetTimeTypeFromArrowType(
        checked_pointer_cast<arrow::TimestampType>(builder.type()));
    PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast<int64_t>(literals.size())));
    for (const auto& literal : literals) {
        if (literal.IsNull()) {
            PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull());
            continue;
        }
        PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(
            DateTimeUtils::TimestampToInteger(literal.GetValue<Timestamp>(), time_type)));
    }
    std::shared_ptr<arrow::Array> array;
    PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array));
    return array;
}
}  // namespace

Result<std::shared_ptr<arrow::Array>> LiteralConverter::ConvertLiteralsToArray(
    const FieldType& field_type, const std::vector<Literal>& literals, arrow::MemoryPool* pool) {
    switch (field_type) {
        case FieldType::BOOLEAN:
            return BuildArray<arrow::BooleanBuilder>(
                literals, [](const Literal& literal) { return literal.GetValue<bool>(); }, pool);
        case FieldType::TINYINT:
            return BuildArray<arrow::Int8Builder>(
                literals, [](const Literal& literal) { return literal.GetValue<int8_t>(); }, pool);
        case FieldType::SMALLINT:
            return BuildArray<arrow::Int16Builder>(
                literals, [](const Literal& literal) { return literal.GetValue<int16_t>(); }, pool);
        case FieldType::INT:
            return BuildArray<arrow::Int32Builder>(
                literals, [](const Literal& literal) { return literal.GetValue<int32_t>(); }, pool);
        case FieldType::BIGINT:
            return BuildArray<arrow::Int64Builder>(
                literals, [](const Literal& literal) { return literal.GetValue<int64_t>(); }, pool);
        case FieldType::FLOAT:
            return BuildArray<arrow::FloatBuilder>(
                literals, [](const Literal& literal) { return literal.GetValue<float>(); }, pool);
        case FieldType::DOUBLE:
            return BuildArray<arrow::DoubleBuilder>(
                literals, [](const Literal& literal) { return literal.GetValue<double>(); }, pool);
        case FieldType::DATE:
            return BuildArray<arrow::Date32Builder>(
                literals, [](const Literal& literal) { return literal.GetValue<int32_t>(); }, pool);
        case FieldType::STRING:
            return BuildArray<arrow::StringBuilder>(
                literals, [](const Literal& literal) { return literal.GetValue<std::string>(); },
                pool);
        case FieldType::BINARY:
            return BuildArray<arrow::BinaryBuilder>(
                literals, [](const Literal& literal) { return literal.GetValue<std::string>(); },
                pool);
        case FieldType::DECIMAL:
            return BuildDecimalArray(literals, pool);
        case FieldType::TIMESTAMP:
            return BuildTimestampArray(literals, pool);
        default:
            return Status::Invalid(
                fmt::format("Not support converting literals of {} type to an arrow array",
                            FieldTypeUtils::FieldTypeToString(field_type)));
    }
}
}  // namespace paimon