perf(predicate): evaluate IN and NOT IN with arrow::compute::is_in (#265)

Author: lucasfang

src/paimon/CMakeLists.txt

@@ -125,6 +125,7 @@ set(PAIMON_COMMON_SRCS
     common/predicate/like.cpp
     common/predicate/literal_converter.cpp
     common/predicate/literal.cpp
+    common/predicate/multi_literals_leaf_function.cpp
     common/predicate/not_equal.cpp
     common/predicate/not_in.cpp
     common/predicate/or.cpp
@@ -609,6 +610,7 @@ if(PAIMON_BUILD_TESTS)
                     common/options/time_duration_test.cpp
                     common/predicate/literal_converter_test.cpp
                     common/predicate/literal_test.cpp
+                    common/predicate/multi_literals_leaf_function_test.cpp
                     common/predicate/predicate_test.cpp
                     common/predicate/predicate_utils_test.cpp
                     common/predicate/predicate_validator_test.cpp

src/paimon/common/predicate/and.h

@@ -49,17 +49,17 @@ class And : public CompoundFunction {
         return instance;
     }
 
-    Result<std::vector<char>> Test(
-        const arrow::Array& array,
-        const std::vector<std::shared_ptr<Predicate>>& children) const override {
+    Result<std::vector<char>> Test(const arrow::Array& array,
+                                   const std::vector<std::shared_ptr<Predicate>>& children,
+                                   arrow::MemoryPool* pool) const override {
         std::vector<char> is_valid(array.length(), true);
         for (const auto& child : children) {
             auto child_filter = std::dynamic_pointer_cast<PredicateFilter>(child);
             if (!child_filter) {
                 return Status::Invalid(
                     fmt::format("child filter {} does not support Test", child->ToString()));
             }
-            PAIMON_ASSIGN_OR_RAISE(std::vector<char> child_valid, child_filter->Test(array));
+            PAIMON_ASSIGN_OR_RAISE(std::vector<char> child_valid, child_filter->Test(array, pool));
             for (size_t i = 0; i < is_valid.size(); i++) {
                 is_valid[i] = (is_valid[i] & child_valid[i]);
             }

src/paimon/common/predicate/compound_function.h

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "arrow/array/array_base.h"
+#include "arrow/type_fwd.h"
 #include "paimon/common/data/internal_array.h"
 #include "paimon/common/data/internal_row.h"
 #include "paimon/predicate/function.h"
@@ -32,9 +33,10 @@ namespace paimon {
 class CompoundFunction : public Function {
  public:
     // input array is the struct array of all fields
-    virtual Result<std::vector<char>> Test(
-        const arrow::Array& array,
-        const std::vector<std::shared_ptr<Predicate>>& children) const = 0;
+    // `pool` is where any arrow buffer the evaluation allocates comes from, it must not be null.
+    virtual Result<std::vector<char>> Test(const arrow::Array& array,
+                                           const std::vector<std::shared_ptr<Predicate>>& children,
+                                           arrow::MemoryPool* pool) const = 0;
 
     virtual Result<bool> Test(const std::shared_ptr<arrow::Schema>& schema, const InternalRow& row,
                               const std::vector<std::shared_ptr<Predicate>>& children) const = 0;

src/paimon/common/predicate/compound_predicate_impl.h

@@ -32,8 +32,9 @@ class CompoundPredicateImpl : public CompoundPredicate, public PredicateFilter {
                           const std::vector<std::shared_ptr<Predicate>>& children)
         : CompoundPredicate(compound_function, children) {}
 
-    Result<std::vector<char>> Test(const arrow::Array& array) const override {
-        return compound_function_.Test(array, children_);
+    Result<std::vector<char>> Test(const arrow::Array& array,
+                                   arrow::MemoryPool* pool) const override {
+        return compound_function_.Test(array, children_, pool);
     }
 
     Result<bool> Test(const std::shared_ptr<arrow::Schema>& schema,

src/paimon/common/predicate/leaf_function.h

@@ -21,15 +21,18 @@
 #include <vector>
 
 #include "arrow/array/array_base.h"
+#include "arrow/type_fwd.h"
 #include "paimon/predicate/function.h"
 #include "paimon/predicate/literal.h"
 #include "paimon/status.h"
 namespace paimon {
 class LeafFunction : public Function {
  public:
     // input array is the exact single field array
+    // `pool` is where any arrow buffer the evaluation allocates comes from, it must not be null.
     virtual Result<std::vector<char>> Test(const arrow::Array& array,
-                                           const std::vector<Literal>& literals) const = 0;
+                                           const std::vector<Literal>& literals,
+                                           arrow::MemoryPool* pool) const = 0;
 
     virtual Result<bool> Test(const Literal& value, const std::vector<Literal>& literals) const = 0;
 

src/paimon/common/predicate/leaf_predicate_impl.h

@@ -41,15 +41,16 @@ class LeafPredicateImpl : public LeafPredicate, public PredicateFilter {
         return leaf_function_;
     }
 
-    Result<std::vector<char>> Test(const arrow::Array& array) const override {
+    Result<std::vector<char>> Test(const arrow::Array& array,
+                                   arrow::MemoryPool* pool) const override {
         const auto& struct_array = checked_cast<const arrow::StructArray&>(array);
         if (field_index_ >= static_cast<int32_t>(struct_array.fields().size())) {
             return Status::Invalid(
                 fmt::format("field index {} exceed field count {} in struct array", field_index_,
                             struct_array.fields().size()));
         }
         const auto& field_array = struct_array.field(field_index_);
-        return leaf_function_.Test(*field_array, literals_);
+        return leaf_function_.Test(*field_array, literals_, pool);
     }
 
     Result<bool> Test(const std::shared_ptr<arrow::Schema>& schema,

src/paimon/common/predicate/leaf_unary_function.h

@@ -30,8 +30,8 @@
 namespace paimon {
 class LeafUnaryFunction : public LeafFunction {
  public:
-    Result<std::vector<char>> Test(const arrow::Array& array,
-                                   const std::vector<Literal>& literals) const override {
+    Result<std::vector<char>> Test(const arrow::Array& array, const std::vector<Literal>& literals,
+                                   arrow::MemoryPool* pool) const override {
         std::vector<char> is_valid(array.length(), false);
         PAIMON_ASSIGN_OR_RAISE(
             std::vector<Literal> array_values,

src/paimon/common/predicate/literal_converter.cpp

@@ -25,12 +25,18 @@
 #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"
@@ -284,4 +290,165 @@ std::vector<Literal> LiteralConverter::GetLiteralFromTimestampArray(const arrow:
     }
     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

src/paimon/common/predicate/literal_converter.h

@@ -27,6 +27,7 @@
 #include <vector>
 
 #include "arrow/array/array_dict.h"
+#include "arrow/type_fwd.h"
 #include "arrow/type_traits.h"
 #include "paimon/common/utils/checked_cast.h"
 #include "paimon/predicate/literal.h"
@@ -50,6 +51,29 @@ class PAIMON_EXPORT LiteralConverter {
     static Result<std::vector<Literal>> ConvertLiteralsFromArray(const arrow::Array& array,
                                                                  bool own_data);
 
+    /// Collects the literals into an arrow array, the reverse of `ConvertLiteralsFromArray`.
+    ///
+    /// @param field_type The field type shared by every literal, it picks the arrow type.
+    /// @param literals The literals to convert, a null literal is written as a null, so the result
+    ///                 has one entry per literal.
+    /// @param pool The pool every buffer of the result is allocated from, it must not be null.
+    /// @return `Status::Invalid` for a field type this does not write, which is every one outside
+    ///         `BOOLEAN`, `TINYINT`, `SMALLINT`, `INT`, `BIGINT`, `FLOAT`, `DOUBLE`, `DATE`,
+    ///         `STRING`, `BINARY`, `DECIMAL` and `TIMESTAMP`. `DECIMAL` is written with the
+    ///         precision and the scale its literals carry, so at least one literal has to be non
+    ///         null and every non null one has to carry the same pair. `TIMESTAMP` is written with
+    ///         the finest time unit that keeps every value, so at least one literal has to be non
+    ///         null to settle the unit.
+    static Result<std::shared_ptr<arrow::Array>> ConvertLiteralsToArray(
+        const FieldType& field_type, const std::vector<Literal>& literals, arrow::MemoryPool* pool);
+
+    /// The finest time unit any of the non null timestamp literals needs to keep its value, which
+    /// is the unit `ConvertLiteralsToArray` writes them with. A null literal does not constrain the
+    /// unit, and literals that are all null leave it at the coarsest one.
+    ///
+    /// Every literal has to be of `TIMESTAMP` type.
+    static arrow::TimeUnit::type MinRequiredTimeUnit(const std::vector<Literal>& literals);
+
     static Result<Literal> ConvertLiteralsFromString(const FieldType& type,
                                                      const std::string& value_str);