diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 1749634f4..857c19b71 100644 --- a/src/paimon/CMakeLists.txt +++ b/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 diff --git a/src/paimon/common/predicate/arrow_set_lookup.cpp b/src/paimon/common/predicate/arrow_set_lookup.cpp new file mode 100644 index 000000000..eb0bccefb --- /dev/null +++ b/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 + +#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 +std::shared_ptr BuildValueSet(const std::vector& literals, Extract extract) { + BuilderType builder; + if (!builder.Reserve(static_cast(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 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 BuildTypedValueSet(FieldType field_type, + const std::vector& literals) { + switch (field_type) { + case FieldType::BOOLEAN: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::TINYINT: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::SMALLINT: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::INT: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::BIGINT: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::DATE: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::STRING: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + case FieldType::BINARY: + return BuildValueSet( + literals, [](const Literal& literal) { return literal.GetValue(); }); + 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 MakeInValueSet(const std::vector& 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* out) { + if (out == nullptr || static_cast(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(*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(matched.Value(i) != negate); + } + return Status::OK(); +} + +} // namespace paimon diff --git a/src/paimon/common/predicate/arrow_set_lookup.h b/src/paimon/common/predicate/arrow_set_lookup.h new file mode 100644 index 000000000..7afa65c8d --- /dev/null +++ b/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 +#include + +#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 MakeInValueSet(const std::vector& 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* out); + +} // namespace paimon diff --git a/src/paimon/common/predicate/arrow_set_lookup_test.cpp b/src/paimon/common/predicate/arrow_set_lookup_test.cpp new file mode 100644 index 000000000..937554694 --- /dev/null +++ b/src/paimon/common/predicate/arrow_set_lookup_test.cpp @@ -0,0 +1,373 @@ +/* + * 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 +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/array_dict.h" +#include "arrow/ipc/json_simple.h" +#include "gtest/gtest.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" +#include "paimon/status.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +class ArrowSetLookupTest : public ::testing::Test { + public: + void SetUp() override {} + void TearDown() override {} + + static Literal StringLiteral(const std::string& value) { + return Literal(FieldType::STRING, value.data(), value.size()); + } + + static Literal BinaryLiteral(const std::string& value) { + return Literal(FieldType::BINARY, value.data(), value.size()); + } + + // Builds the value set of an `IN` predicate. `negate` only rejects a null literal, so as long + // as the literals hold none the same value set probes `NOT IN` too. + static std::shared_ptr InValueSet(const std::vector& literals) { + return MakeInValueSet(literals, /*negate=*/false); + } + + // Probes `array` and returns the per row result, asserting the whole call succeeded. + static std::vector Probe(const arrow::Array& value_set, + const std::shared_ptr& array, bool negate) { + std::vector is_valid(array->length(), 0); + Status status = ProbeInValueSet(*array, value_set, negate, &is_valid); + EXPECT_TRUE(status.ok()) << status.ToString(); + return is_valid; + } +}; + +TEST_F(ArrowSetLookupTest, TestInt) { + auto value_set = InValueSet({Literal(1), Literal(2), Literal(3), Literal(5)}); + ASSERT_TRUE(value_set); + + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([0, 1, 2, 3, 4, 5, 6, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), + std::vector({0, 1, 1, 1, 0, 1, 0, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), + std::vector({1, 0, 0, 0, 1, 0, 1, 0})); +} + +TEST_F(ArrowSetLookupTest, TestBigInt) { + auto value_set = + InValueSet({Literal(int64_t{-1000000}), Literal(int64_t{0}), Literal(int64_t{1000000})}); + ASSERT_TRUE(value_set); + + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-1000000, -999999, 0, 1, 1000000, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 0, 1, 0, 1, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({0, 1, 0, 1, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestInt64Boundaries) { + auto value_set = InValueSet({Literal(std::numeric_limits::min()), + Literal(std::numeric_limits::max())}); + ASSERT_TRUE(value_set); + + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-9223372036854775808, 0, 9223372036854775807, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 0, 1, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({0, 1, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestTinyIntAndSmallInt) { + auto tinyint_set = + InValueSet({Literal(int8_t{-128}), Literal(int8_t{0}), Literal(int8_t{127})}); + ASSERT_TRUE(tinyint_set); + auto tinyint_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int8(), R"([-128, -127, 0, 126, 127])") + .ValueOrDie(); + ASSERT_EQ(Probe(*tinyint_set, tinyint_array, /*negate=*/false), + std::vector({1, 0, 1, 0, 1})); + + auto smallint_set = InValueSet({Literal(int16_t{-30000}), Literal(int16_t{30000})}); + ASSERT_TRUE(smallint_set); + auto smallint_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int16(), R"([-30000, 0, 30000, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*smallint_set, smallint_array, /*negate=*/false), + std::vector({1, 0, 1, 0})); +} + +TEST_F(ArrowSetLookupTest, TestDate) { + auto value_set = InValueSet({Literal(FieldType::DATE, 100), Literal(FieldType::DATE, 20000)}); + ASSERT_TRUE(value_set); + + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::date32(), R"([100, 101, 20000, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 0, 1, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({0, 1, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestBoolean) { + auto true_only = InValueSet({Literal(true)}); + ASSERT_TRUE(true_only); + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::boolean(), R"([true, false, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*true_only, array, /*negate=*/false), std::vector({1, 0, 0})); + ASSERT_EQ(Probe(*true_only, array, /*negate=*/true), std::vector({0, 1, 0})); + + auto both = InValueSet({Literal(true), Literal(false)}); + ASSERT_TRUE(both); + ASSERT_EQ(Probe(*both, array, /*negate=*/false), std::vector({1, 1, 0})); + ASSERT_EQ(Probe(*both, array, /*negate=*/true), std::vector({0, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestString) { + auto value_set = + InValueSet({StringLiteral("apple"), StringLiteral(""), StringLiteral("banana")}); + ASSERT_TRUE(value_set); + + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::utf8(), R"(["apple", "", "banana", "cherry", "app", null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 1, 1, 0, 0, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({0, 0, 0, 1, 1, 0})); +} + +TEST_F(ArrowSetLookupTest, TestStringWithoutEmptyLiteral) { + // An empty column value must not match when no empty literal was given. + auto value_set = InValueSet({StringLiteral("a")}); + ASSERT_TRUE(value_set); + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "", "b"])").ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestBinary) { + // Binary literals keep their embedded zero bytes, the value set must not truncate them. + auto value_set = InValueSet({BinaryLiteral(std::string("\x00\x01", 2)), BinaryLiteral("xyz")}); + ASSERT_TRUE(value_set); + + arrow::BinaryBuilder builder; + ASSERT_TRUE(builder.Append(std::string("\x00\x01", 2)).ok()); + ASSERT_TRUE(builder.Append("xyz").ok()); + ASSERT_TRUE(builder.Append("xyw").ok()); + ASSERT_TRUE(builder.Append(std::string("\x00", 1)).ok()); + ASSERT_TRUE(builder.AppendNull().ok()); + std::shared_ptr array; + ASSERT_TRUE(builder.Finish(&array).ok()); + + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 1, 0, 0, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({0, 0, 1, 1, 0})); +} + +TEST_F(ArrowSetLookupTest, TestDictionaryString) { + auto value_set = InValueSet({StringLiteral("a"), StringLiteral("c")}); + ASSERT_TRUE(value_set); + + auto dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "c"])").ValueOrDie(); + auto indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([0, 1, 2, 2, null, 0])") + .ValueOrDie(); + auto dict_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + auto array = arrow::DictionaryArray::FromArrays(dict_type, indices, dictionary).ValueOrDie(); + + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 0, 1, 1, 0, 1})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({0, 1, 0, 0, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestLargeStringDictionary) { + auto value_set = InValueSet({StringLiteral("c")}); + ASSERT_TRUE(value_set); + + auto dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::large_utf8(), R"(["a", "b", "c"])") + .ValueOrDie(); + auto indices = arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), R"([0, 2, null, 1])") + .ValueOrDie(); + auto dict_type = arrow::dictionary(arrow::int64(), arrow::large_utf8()); + auto array = arrow::DictionaryArray::FromArrays(dict_type, indices, dictionary).ValueOrDie(); + + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({0, 1, 0, 0})); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/true), std::vector({1, 0, 0, 1})); +} + +TEST_F(ArrowSetLookupTest, TestDictionaryWithNullValue) { + // `is_in` decodes the dictionary, so a row pointing at a null dictionary value becomes a null + // row and is false for both `IN` and `NOT IN`. An empty literal must not match it. + auto dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"([null, "a"])").ValueOrDie(); + auto indices = arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([0, 1, null, 0])") + .ValueOrDie(); + auto dict_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + auto array = arrow::DictionaryArray::FromArrays(dict_type, indices, dictionary).ValueOrDie(); + + auto set_with_empty = InValueSet({StringLiteral("a"), StringLiteral("")}); + ASSERT_TRUE(set_with_empty); + ASSERT_EQ(Probe(*set_with_empty, array, /*negate=*/false), std::vector({0, 1, 0, 0})); + ASSERT_EQ(Probe(*set_with_empty, array, /*negate=*/true), std::vector({0, 0, 0, 0})); + + auto set_without_empty = InValueSet({StringLiteral("a")}); + ASSERT_TRUE(set_without_empty); + ASSERT_EQ(Probe(*set_without_empty, array, /*negate=*/false), std::vector({0, 1, 0, 0})); + ASSERT_EQ(Probe(*set_without_empty, array, /*negate=*/true), std::vector({0, 0, 0, 0})); +} + +TEST_F(ArrowSetLookupTest, TestNullLiteralIgnoredForIn) { + const std::vector literals = {Literal(int64_t{1}), Literal(FieldType::BIGINT), + Literal(int64_t{3})}; + auto value_set = MakeInValueSet(literals, /*negate=*/false); + ASSERT_TRUE(value_set); + + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), R"([1, 2, 3, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*value_set, array, /*negate=*/false), std::vector({1, 0, 1, 0})); + + // A null literal makes `NOT IN` false for every row, which `NotIn` already handles. + ASSERT_FALSE(MakeInValueSet(literals, /*negate=*/true)); +} + +TEST_F(ArrowSetLookupTest, TestOnlyNullLiterals) { + auto int_set = MakeInValueSet({Literal(FieldType::BIGINT)}, + /*negate=*/false); + ASSERT_TRUE(int_set); + ASSERT_EQ(int_set->length(), 0); + auto int_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-9007199254740993, 0, 9007199254740993, null])") + .ValueOrDie(); + ASSERT_EQ(Probe(*int_set, int_array, /*negate=*/false), std::vector({0, 0, 0, 0})); + + auto string_set = MakeInValueSet({Literal(FieldType::STRING)}, + /*negate=*/false); + ASSERT_TRUE(string_set); + auto string_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["", "a", null])").ValueOrDie(); + ASSERT_EQ(Probe(*string_set, string_array, /*negate=*/false), std::vector({0, 0, 0})); + + ASSERT_FALSE(MakeInValueSet({Literal(FieldType::BIGINT)}, /*negate=*/true)); +} + +TEST_F(ArrowSetLookupTest, TestMakeInValueSetUnsupported) { + // Empty literals. + ASSERT_FALSE(InValueSet({})); + // Literals of mixed types make `Literal::CompareTo` fail, so they must keep reporting the error + // through the generic path instead of being built into one typed value set. + ASSERT_FALSE(InValueSet({Literal(1), Literal(int64_t{2})})); + ASSERT_FALSE(InValueSet({StringLiteral("a"), Literal(int64_t{1})})); + // Types whose equality is not the hash based equality of `is_in`. + ASSERT_FALSE(InValueSet({Literal(1.0)})); + ASSERT_FALSE(InValueSet({Literal(1.0f)})); + ASSERT_FALSE(InValueSet( + {Literal(Decimal::FromUnscaledLong(/*unscaled_long=*/10, /*precision=*/10, /*scale=*/1))})); + ASSERT_FALSE(InValueSet({Literal(Timestamp::FromEpochMillis(1))})); +} + +TEST_F(ArrowSetLookupTest, TestProbeResolvesArrowTypeOnItsOwn) { + // A column read as a narrower arrow type than the field type is promoted to the common type, so + // the probe stays correct instead of having to fall back. + auto bigint_set = InValueSet({Literal(int64_t{1}), Literal(int64_t{2})}); + ASSERT_TRUE(bigint_set); + auto int32_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([1, 3, null])").ValueOrDie(); + ASSERT_EQ(Probe(*bigint_set, int32_array, /*negate=*/false), std::vector({1, 0, 0})); + + // Promotion widens both sides, so a value outside the value set type range cannot alias. + auto int_set = InValueSet({Literal(1)}); + ASSERT_TRUE(int_set); + auto int64_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), R"([1, 4294967297])") + .ValueOrDie(); + ASSERT_EQ(Probe(*int_set, int64_array, /*negate=*/false), std::vector({1, 0})); + + // Dictionary layouts that `LiteralConverter::ConvertLiteralsFromArray` rejects are decoded by + // `is_in`, so they no longer fail the whole evaluation. + auto int64_dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), R"([10, 20])").ValueOrDie(); + auto int32_indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([0, 1])").ValueOrDie(); + auto int64_dict_array = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::int64()), + int32_indices, int64_dictionary) + .ValueOrDie(); + auto ten_set = InValueSet({Literal(int64_t{10})}); + ASSERT_TRUE(ten_set); + ASSERT_EQ(Probe(*ten_set, int64_dict_array, /*negate=*/false), std::vector({1, 0})); + + auto string_dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b"])").ValueOrDie(); + auto int8_indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int8(), R"([0, 1])").ValueOrDie(); + auto int8_dict_array = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int8(), arrow::utf8()), + int8_indices, string_dictionary) + .ValueOrDie(); + auto a_set = InValueSet({StringLiteral("a")}); + ASSERT_TRUE(a_set); + ASSERT_EQ(Probe(*a_set, int8_dict_array, /*negate=*/false), std::vector({1, 0})); +} + +TEST_F(ArrowSetLookupTest, TestProbeFailsOnUnrelatedArrowType) { + // Nothing casts an int32 value set to a string column, so the probe reports the mismatch. This + // only happens when the field type disagrees with the column being evaluated. + auto value_set = InValueSet({Literal(1)}); + ASSERT_TRUE(value_set); + auto string_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["1"])").ValueOrDie(); + std::vector is_valid(string_array->length(), 0); + ASSERT_FALSE(ProbeInValueSet(*string_array, *value_set, /*negate=*/false, &is_valid).ok()); +} + +TEST_F(ArrowSetLookupTest, TestSlicedArray) { + auto value_set = InValueSet({Literal(2), Literal(4)}); + ASSERT_TRUE(value_set); + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([1, 2, 3, 4, 5, null])") + .ValueOrDie(); + auto sliced = array->Slice(2, 4); + ASSERT_EQ(Probe(*value_set, sliced, /*negate=*/false), std::vector({0, 1, 0, 0})); + ASSERT_EQ(Probe(*value_set, sliced, /*negate=*/true), std::vector({1, 0, 1, 0})); + + auto string_set = InValueSet({StringLiteral("c")}); + ASSERT_TRUE(string_set); + auto string_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "c", "d"])") + .ValueOrDie(); + auto sliced_string = string_array->Slice(1, 3); + ASSERT_EQ(Probe(*string_set, sliced_string, /*negate=*/false), std::vector({0, 1, 0})); +} + +TEST_F(ArrowSetLookupTest, TestOutputBufferValidation) { + auto value_set = InValueSet({Literal(1)}); + ASSERT_TRUE(value_set); + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([1, 2])").ValueOrDie(); + ASSERT_FALSE(ProbeInValueSet(*array, *value_set, /*negate=*/false, nullptr).ok()); + std::vector too_small(1, 0); + ASSERT_FALSE(ProbeInValueSet(*array, *value_set, /*negate=*/false, &too_small).ok()); +} +} // namespace paimon::test diff --git a/src/paimon/common/predicate/in.h b/src/paimon/common/predicate/in.h index c2301ad04..8490d3672 100644 --- a/src/paimon/common/predicate/in.h +++ b/src/paimon/common/predicate/in.h @@ -19,11 +19,14 @@ #pragma once #include +#include #include #include #include #include +#include "arrow/array/array_base.h" +#include "paimon/common/predicate/arrow_set_lookup.h" #include "paimon/common/predicate/multi_literals_leaf_function.h" #include "paimon/predicate/literal.h" #include "paimon/result.h" @@ -38,6 +41,20 @@ class In : public MultiLiteralsLeafFunction { static const In instance = In(); return instance; } + /// Probes the whole batch with `arrow::compute::is_in` rather than comparing every row against + /// every literal. 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> Test(const arrow::Array& array, + const std::vector& literals) const override { + std::shared_ptr value_set = MakeInValueSet(literals, /*negate=*/false); + if (value_set == nullptr) { + return MultiLiteralsLeafFunction::Test(array, literals); + } + std::vector is_valid(array.length(), 0); + PAIMON_RETURN_NOT_OK(ProbeInValueSet(array, *value_set, /*negate=*/false, &is_valid)); + return is_valid; + } + Result InnerTest(const Literal& field, const std::vector& literals) const override { for (const auto& literal : literals) { diff --git a/src/paimon/common/predicate/leaf_predicate_impl.h b/src/paimon/common/predicate/leaf_predicate_impl.h index c9a8758f5..338234fe6 100644 --- a/src/paimon/common/predicate/leaf_predicate_impl.h +++ b/src/paimon/common/predicate/leaf_predicate_impl.h @@ -87,13 +87,13 @@ class LeafPredicateImpl : public LeafPredicate, public PredicateFilter { } std::shared_ptr NewLeafPredicate(int32_t new_field_index) const { - return std::make_shared(leaf_function_, new_field_index, field_name_, - field_type_, literals_); + return std::shared_ptr(new LeafPredicateImpl( + leaf_function_, new_field_index, field_name_, field_type_, literals_)); } std::shared_ptr NewLeafPredicate(const std::string& new_field_name) const { - return std::make_shared(leaf_function_, field_index_, new_field_name, - field_type_, literals_); + return std::shared_ptr(new LeafPredicateImpl( + leaf_function_, field_index_, new_field_name, field_type_, literals_)); } }; } // namespace paimon diff --git a/src/paimon/common/predicate/not_in.h b/src/paimon/common/predicate/not_in.h index 5de896242..17d02c8fa 100644 --- a/src/paimon/common/predicate/not_in.h +++ b/src/paimon/common/predicate/not_in.h @@ -18,11 +18,14 @@ #pragma once #include +#include #include #include #include #include +#include "arrow/array/array_base.h" +#include "paimon/common/predicate/arrow_set_lookup.h" #include "paimon/common/predicate/multi_literals_leaf_function.h" #include "paimon/predicate/literal.h" #include "paimon/result.h" @@ -38,6 +41,20 @@ class NotIn : public MultiLiteralsLeafFunction { return instance; } + /// Probes the whole batch with `arrow::compute::is_in` rather than comparing every row against + /// every literal. 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> Test(const arrow::Array& array, + const std::vector& literals) const override { + std::shared_ptr value_set = MakeInValueSet(literals, /*negate=*/true); + if (value_set == nullptr) { + return MultiLiteralsLeafFunction::Test(array, literals); + } + std::vector is_valid(array.length(), 0); + PAIMON_RETURN_NOT_OK(ProbeInValueSet(array, *value_set, /*negate=*/true, &is_valid)); + return is_valid; + } + Result InnerTest(const Literal& field, const std::vector& literals) const override { for (const auto& literal : literals) { diff --git a/src/paimon/common/predicate/predicate_test.cpp b/src/paimon/common/predicate/predicate_test.cpp index fa4733cba..d8b9ccaac 100644 --- a/src/paimon/common/predicate/predicate_test.cpp +++ b/src/paimon/common/predicate/predicate_test.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -29,10 +30,12 @@ #include "arrow/api.h" #include "arrow/array/array_nested.h" #include "arrow/ipc/json_simple.h" +#include "fmt/format.h" #include "gtest/gtest.h" #include "paimon/common/data/binary_array.h" #include "paimon/common/data/binary_row.h" #include "paimon/common/data/binary_row_writer.h" +#include "paimon/common/predicate/leaf_predicate_impl.h" #include "paimon/common/predicate/predicate_filter.h" #include "paimon/defs.h" #include "paimon/memory/memory_pool.h" @@ -938,6 +941,95 @@ TEST_F(PredicateTest, TestLargeNotInNull) { ASSERT_FALSE(StatsCheck(*predicate, 3ll, {FieldStats(29ll, 32ll, 0ll)})); } +TEST_F(PredicateTest, TestLargeStringIn) { + auto string_type = arrow::utf8(); + std::vector literals; + literals.reserve(1000); + for (int32_t i = 0; i < 1000; i++) { + std::string value = fmt::format("key-{}", i); + literals.emplace_back(FieldType::STRING, value.data(), value.size()); + } + auto in_base = + PredicateBuilder::In(/*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, literals); + auto in_predicate = std::dynamic_pointer_cast(in_base); + ASSERT_TRUE(in_predicate); + auto not_in_base = PredicateBuilder::NotIn(/*field_index=*/0, /*field_name=*/"f0", + FieldType::STRING, literals); + auto not_in_predicate = std::dynamic_pointer_cast(not_in_base); + ASSERT_TRUE(not_in_predicate); + + auto f0 = arrow::ipc::internal::json::ArrayFromJSON( + string_type, R"(["key-0", "key-999", "key-1000", "other", "", null])") + .ValueOrDie(); + std::shared_ptr src_type = arrow::struct_({arrow::field("f0", string_type)}); + std::shared_ptr struct_array = + arrow::StructArray::Make({f0}, src_type->fields()).ValueOrDie(); + + ASSERT_OK_AND_ASSIGN(auto in_valid, in_predicate->Test(*struct_array)); + ASSERT_EQ(in_valid, std::vector({1, 1, 0, 0, 0, 0})); + ASSERT_OK_AND_ASSIGN(auto not_in_valid, not_in_predicate->Test(*struct_array)); + ASSERT_EQ(not_in_valid, std::vector({0, 0, 1, 1, 1, 0})); +} + +TEST_F(PredicateTest, TestInAfterRebind) { + auto bigint_type = arrow::int64(); + auto predicate_base = PredicateBuilder::In(/*field_index=*/0, /*field_name=*/"f0", + FieldType::BIGINT, {Literal(1l), Literal(3l)}); + auto leaf_predicate = std::dynamic_pointer_cast(predicate_base); + ASSERT_TRUE(leaf_predicate); + + // Rebinding shares the prebuilt lookup structure, results must stay identical. + auto renamed = leaf_predicate->NewLeafPredicate(/*new_field_name=*/"f1"); + ASSERT_EQ(renamed->FieldName(), "f1"); + auto rebound = renamed->NewLeafPredicate(/*new_field_index=*/1); + ASSERT_EQ(rebound->FieldIndex(), 1); + + auto f0 = + arrow::ipc::internal::json::ArrayFromJSON(bigint_type, R"([3, 2, 1, 0])").ValueOrDie(); + auto f1 = + arrow::ipc::internal::json::ArrayFromJSON(bigint_type, R"([1, 2, 3, null])").ValueOrDie(); + std::shared_ptr src_type = + arrow::struct_({arrow::field("f0", bigint_type), arrow::field("f1", bigint_type)}); + std::shared_ptr struct_array = + arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); + + ASSERT_OK_AND_ASSIGN(auto is_valid, rebound->Test(*struct_array)); + ASSERT_EQ(is_valid, std::vector({1, 0, 1, 0})); + + auto arrow_schema = arrow::schema( + arrow::FieldVector({arrow::field("f0", bigint_type), arrow::field("f1", bigint_type)})); + ASSERT_TRUE(rebound->Test(arrow_schema, CreateBigIntRow({0, 1})).value()); + ASSERT_FALSE(rebound->Test(arrow_schema, CreateBigIntRow({1, 2})).value()); +} + +TEST_F(PredicateTest, TestInt64BoundaryIn) { + // Building the lookup for the full int64 range used to crash with an out of bounds dense + // bitmap index; construction itself is part of what this test guards. + std::vector literals = {Literal(std::numeric_limits::min()), + Literal(std::numeric_limits::max())}; + auto in_base = + PredicateBuilder::In(/*field_index=*/0, /*field_name=*/"f0", FieldType::BIGINT, literals); + auto in_predicate = std::dynamic_pointer_cast(in_base); + ASSERT_TRUE(in_predicate); + auto not_in_base = PredicateBuilder::NotIn(/*field_index=*/0, /*field_name=*/"f0", + FieldType::BIGINT, literals); + auto not_in_predicate = std::dynamic_pointer_cast(not_in_base); + ASSERT_TRUE(not_in_predicate); + + auto f0 = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-9223372036854775808, 0, 9223372036854775807, null])") + .ValueOrDie(); + std::shared_ptr src_type = + arrow::struct_({arrow::field("f0", arrow::int64())}); + std::shared_ptr struct_array = + arrow::StructArray::Make({f0}, src_type->fields()).ValueOrDie(); + + ASSERT_OK_AND_ASSIGN(auto in_valid, in_predicate->Test(*struct_array)); + ASSERT_EQ(in_valid, std::vector({1, 0, 1, 0})); + ASSERT_OK_AND_ASSIGN(auto not_in_valid, not_in_predicate->Test(*struct_array)); + ASSERT_EQ(not_in_valid, std::vector({0, 1, 0, 0})); +} + TEST_F(PredicateTest, TestAnd) { auto bigint_type = arrow::int64(); ASSERT_OK_AND_ASSIGN(