Dev10

src/paimon/CMakeLists.txt

@@ -124,6 +124,7 @@ set(PAIMON_COMMON_SRCS
     common/predicate/less_than.cpp
     common/predicate/like.cpp
     common/predicate/literal_converter.cpp
+    common/predicate/arrow_set_lookup.cpp
     common/predicate/literal.cpp
     common/predicate/not_equal.cpp
     common/predicate/not_in.cpp
@@ -604,6 +605,7 @@ if(PAIMON_BUILD_TESTS)
                     common/options/memory_size_test.cpp
                     common/options/time_duration_test.cpp
                     common/predicate/literal_converter_test.cpp
+                    common/predicate/arrow_set_lookup_test.cpp
                     common/predicate/literal_test.cpp
                     common/predicate/predicate_test.cpp
                     common/predicate/predicate_utils_test.cpp

src/paimon/common/predicate/arrow_set_lookup.cpp

@@ -0,0 +1,150 @@
+/*
+ * 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/arrow_set_lookup.h"
+
+#include <string>
+
+#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/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 fallback path instead of surfacing an error from a constructor.
+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 generic comparison path.
+            return nullptr;
+    }
+}
+}  // namespace
+
+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 fallback
+        // 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);
+}
+
+Status ProbeInValueSet(const arrow::Array& array, const arrow::Array& value_set, bool negate,
+                       std::vector<char>* out) {
+    if (out == nullptr || static_cast<int64_t>(out->size()) != array.length()) {
+        return Status::Invalid("output buffer size must match the array length");
+    }
+
+    // `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());
+    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;
+        }
+        (*out)[i] = static_cast<char>(matched.Value(i) != negate);
+    }
+    return Status::OK();
+}
+
+}  // namespace paimon

src/paimon/common/predicate/arrow_set_lookup.h

@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "paimon/defs.h"
+#include "paimon/predicate/literal.h"
+#include "paimon/status.h"
+
+namespace arrow {
+class Array;
+}  // namespace arrow
+
+namespace paimon {
+
+/// 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`.
+///
+/// `MultiLiteralsLeafFunction` materializes the whole column into `Literal` objects (one heap
+/// allocation per row) and then linearly scans the literals for every row, which is
+/// `O(rows * literals)` with a heap allocation per row. Handing the literals to `is_in` as a value
+/// set lets it hash them instead, so a batch is probed in `O(rows)` without allocating per row.
+///
+/// @param negate `false` for `IN`, `true` for `NOT IN`.
+/// @return `nullptr` when the predicate must keep using the generic `LeafFunction` implementation.
+///         That covers the types whose equality is not the hash based equality of `is_in`, 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);
+
+/// Probes every non-null row of `array` against `value_set`.
+///
+/// @param negate `false` for `IN` semantics, `true` for `NOT IN` semantics.
+/// @param out Must be sized `array.length()` with all elements pre-set to 0. Only non-null rows are
+///            written, so null rows stay 0 (`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.
+Status ProbeInValueSet(const arrow::Array& array, const arrow::Array& value_set, bool negate,
+                       std::vector<char>* out);
+
+}  // namespace paimon