Dev10

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
@@ -605,6 +606,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/leaf_predicate_impl.h

@@ -87,13 +87,13 @@ class LeafPredicateImpl : public LeafPredicate, public PredicateFilter {
     }
 
     std::shared_ptr<LeafPredicateImpl> NewLeafPredicate(int32_t new_field_index) const {
-        return std::make_shared<LeafPredicateImpl>(leaf_function_, new_field_index, field_name_,
-                                                   field_type_, literals_);
+        return std::shared_ptr<LeafPredicateImpl>(new LeafPredicateImpl(
+            leaf_function_, new_field_index, field_name_, field_type_, literals_));
     }
 
     std::shared_ptr<LeafPredicateImpl> NewLeafPredicate(const std::string& new_field_name) const {
-        return std::make_shared<LeafPredicateImpl>(leaf_function_, field_index_, new_field_name,
-                                                   field_type_, literals_);
+        return std::shared_ptr<LeafPredicateImpl>(new LeafPredicateImpl(
+            leaf_function_, field_index_, new_field_name, field_type_, literals_));
     }
 };
 }  // namespace paimon

src/paimon/common/predicate/multi_literals_leaf_function.cpp

@@ -0,0 +1,196 @@
+/*
+ * 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/multi_literals_leaf_function.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "arrow/array/array_base.h"
+#include "arrow/array/array_primitive.h"
+#include "arrow/array/builder_binary.h"
+#include "arrow/array/builder_primitive.h"
+#include "arrow/compute/api_scalar.h"
+#include "arrow/datum.h"
+#include "arrow/type.h"
+#include "paimon/common/predicate/literal_converter.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+#include "paimon/common/utils/checked_cast.h"
+
+namespace paimon {
+namespace {
+// Collects the non-null literals into an arrow array. Returns `nullptr` on any arrow failure, which
+// keeps the predicate on the row by row path instead of turning a build error into an evaluation
+// error.
+template <typename BuilderType, typename Extract>
+std::shared_ptr<arrow::Array> BuildValueSet(const std::vector<Literal>& literals, Extract extract) {
+    BuilderType builder;
+    if (!builder.Reserve(static_cast<int64_t>(literals.size())).ok()) {
+        return nullptr;
+    }
+    for (const auto& literal : literals) {
+        if (literal.IsNull()) {
+            continue;
+        }
+        if (!builder.Append(extract(literal)).ok()) {
+            return nullptr;
+        }
+    }
+    std::shared_ptr<arrow::Array> value_set;
+    if (!builder.Finish(&value_set).ok()) {
+        return nullptr;
+    }
+    return value_set;
+}
+
+// Builds the value set with the arrow type a column of `field_type` is read as, so that `is_in`
+// only has to promote the two sides when the column is read as another arrow type.
+std::shared_ptr<arrow::Array> BuildTypedValueSet(FieldType field_type,
+                                                 const std::vector<Literal>& literals) {
+    switch (field_type) {
+        case FieldType::BOOLEAN:
+            return BuildValueSet<arrow::BooleanBuilder>(
+                literals, [](const Literal& literal) { return literal.GetValue<bool>(); });
+        case FieldType::TINYINT:
+            return BuildValueSet<arrow::Int8Builder>(
+                literals, [](const Literal& literal) { return literal.GetValue<int8_t>(); });
+        case FieldType::SMALLINT:
+            return BuildValueSet<arrow::Int16Builder>(
+                literals, [](const Literal& literal) { return literal.GetValue<int16_t>(); });
+        case FieldType::INT:
+            return BuildValueSet<arrow::Int32Builder>(
+                literals, [](const Literal& literal) { return literal.GetValue<int32_t>(); });
+        case FieldType::BIGINT:
+            return BuildValueSet<arrow::Int64Builder>(
+                literals, [](const Literal& literal) { return literal.GetValue<int64_t>(); });
+        case FieldType::DATE:
+            return BuildValueSet<arrow::Date32Builder>(
+                literals, [](const Literal& literal) { return literal.GetValue<int32_t>(); });
+        case FieldType::STRING:
+            return BuildValueSet<arrow::StringBuilder>(
+                literals, [](const Literal& literal) { return literal.GetValue<std::string>(); });
+        case FieldType::BINARY:
+            return BuildValueSet<arrow::BinaryBuilder>(
+                literals, [](const Literal& literal) { return literal.GetValue<std::string>(); });
+        default:
+            // FLOAT / DOUBLE hash the raw bits, so canonicalized NaNs would stop matching,
+            // TIMESTAMP needs unit conversion and DECIMAL compares across scales. All of them keep
+            // the row by row comparison path.
+            return nullptr;
+    }
+}
+
+/// Builds the value set that `arrow::compute::is_in` takes from the literals of an `IN` / `NOT IN`
+/// predicate. The arrow type comes from the literals themselves, which all share one `FieldType`.
+///
+/// @param negate `false` for `IN`, `true` for `NOT IN`.
+/// @return `nullptr` when the literals cannot be probed by `is_in`, which covers the types whose
+///         equality is not its hash based equality, and `NOT IN` holding a null literal, which
+///         `NotIn::InnerTest` makes false for every row. This never fails.
+std::shared_ptr<arrow::Array> MakeInValueSet(const std::vector<Literal>& literals, bool negate) {
+    if (literals.empty()) {
+        return nullptr;
+    }
+    // The literals of one predicate share a type, so take it from the first non-null one. When
+    // every literal is null the value set ends up empty, but it still needs a type for `is_in` to
+    // compare it against the column, and a null `Literal` carries its type too.
+    FieldType field_type = literals.front().GetType();
+    for (const auto& literal : literals) {
+        if (!literal.IsNull()) {
+            field_type = literal.GetType();
+            break;
+        }
+    }
+    for (const auto& literal : literals) {
+        // A literal typed differently makes `Literal::CompareTo` fail, keep that on the row by row
+        // path.
+        if (!literal.IsNull() && literal.GetType() != field_type) {
+            return nullptr;
+        }
+        // `NotIn::InnerTest` returns false as soon as it meets a null literal, so no row can match
+        // and there is nothing worth building.
+        if (negate && literal.IsNull()) {
+            return nullptr;
+        }
+    }
+    return BuildTypedValueSet(field_type, literals);
+}
+
+/// Probes every non-null row of `array` against `value_set`.
+///
+/// @param negate `false` for `IN` semantics, `true` for `NOT IN` semantics.
+/// @return One entry per row, with the null rows left at 0 because `IN` and `NOT IN` are both false
+///         on null.
+///
+/// `is_in` resolves the comparison itself: it decodes a dictionary column, and promotes both sides
+/// to their common type when the column is read as a wider or narrower arrow type than the one
+/// `value_set` was built with. It fails when the two types have no common type at all, which only
+/// happens when the field type disagrees with the column the predicate is evaluated against.
+Result<std::vector<char>> ProbeInValueSet(const arrow::Array& array, const arrow::Array& value_set,
+                                          bool negate) {
+    // `EMIT_NULL` ignores the nulls of the value set and turns a null input into a null output, so
+    // the validity of `matches` marks exactly the rows that `In` / `NotIn` consider null. That also
+    // covers a dictionary column, whose null rows come either from the indices or from a null
+    // dictionary value once `is_in` decodes it.
+    arrow::compute::SetLookupOptions options(value_set.data(),
+                                             arrow::compute::SetLookupOptions::EMIT_NULL);
+    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(arrow::Datum matches,
+                                      arrow::compute::IsIn(arrow::Datum(array), options));
+    const auto& matched = checked_cast<const arrow::BooleanArray&>(*matches.make_array());
+    std::vector<char> is_valid(matched.length(), 0);
+    for (int64_t i = 0; i < matched.length(); i++) {
+        if (matched.IsNull(i)) {
+            // `IN` and `NOT IN` are both false on a null value, leave the row at 0.
+            continue;
+        }
+        is_valid[i] = static_cast<char>(matched.Value(i) != negate);
+    }
+    return is_valid;
+}
+}  // namespace
+
+Result<std::vector<char>> MultiLiteralsLeafFunction::Test(
+    const arrow::Array& array, const std::vector<Literal>& literals) const {
+    const Function::Type type = GetType();
+    // `In` and `NotIn` are the only subclasses today. The check keeps a future one off this path,
+    // because `MakeInValueSet` only looks at the literals and would silently give it `IN`
+    // semantics.
+    if (type == Function::Type::IN || type == Function::Type::NOT_IN) {
+        const bool negate = type == Function::Type::NOT_IN;
+        std::shared_ptr<arrow::Array> value_set = MakeInValueSet(literals, negate);
+        if (value_set != nullptr) {
+            return ProbeInValueSet(array, *value_set, negate);
+        }
+    }
+
+    // Materializing the column into `Literal` objects costs one heap allocation per row and then
+    // every row scans the literals linearly, so this only runs when `is_in` cannot probe them.
+    PAIMON_ASSIGN_OR_RAISE(std::vector<Literal> array_values,
+                           LiteralConverter::ConvertLiteralsFromArray(array, /*own_data=*/false));
+    std::vector<char> is_valid(array.length(), false);
+    for (int64_t i = 0; i < array.length(); i++) {
+        if (!array.IsNull(i)) {
+            PAIMON_ASSIGN_OR_RAISE(is_valid[i], Test(array_values[i], literals));
+        }
+    }
+    return is_valid;
+}
+
+}  // namespace paimon

src/paimon/common/predicate/multi_literals_leaf_function.h

@@ -23,27 +23,18 @@
 #include "arrow/array/array_nested.h"
 #include "arrow/c/bridge.h"
 #include "paimon/common/predicate/leaf_function.h"
-#include "paimon/common/predicate/literal_converter.h"
-#include "paimon/common/utils/arrow/status_utils.h"
 #include "paimon/status.h"
 
 namespace paimon {
 
 class MultiLiteralsLeafFunction : public LeafFunction {
  public:
+    /// Probes the whole batch with `arrow::compute::is_in` when the literals allow it, and falls
+    /// back to comparing every row against every literal otherwise. Every `LeafFunction` is a
+    /// shared stateless singleton, so the value set is built per batch; that costs `O(literals)`
+    /// and buys an `O(rows)` probe.
     Result<std::vector<char>> Test(const arrow::Array& array,
-                                   const std::vector<Literal>& literals) const override {
-        PAIMON_ASSIGN_OR_RAISE(
-            std::vector<Literal> array_values,
-            LiteralConverter::ConvertLiteralsFromArray(array, /*own_data=*/false));
-        std::vector<char> is_valid(array.length(), false);
-        for (int64_t i = 0; i < array.length(); i++) {
-            if (!array.IsNull(i)) {
-                PAIMON_ASSIGN_OR_RAISE(is_valid[i], Test(array_values[i], literals));
-            }
-        }
-        return is_valid;
-    }
+                                   const std::vector<Literal>& literals) const override;
 
     Result<bool> Test(int64_t row_count, const Literal& min_value, const Literal& max_value,
                       const std::optional<int64_t>& null_count,