diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 051eba324..e5dd9c0f8 100644 --- a/src/paimon/CMakeLists.txt +++ b/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 @@ -608,6 +609,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 diff --git a/src/paimon/common/predicate/and.h b/src/paimon/common/predicate/and.h index b3663ea89..f0da67d74 100644 --- a/src/paimon/common/predicate/and.h +++ b/src/paimon/common/predicate/and.h @@ -49,9 +49,9 @@ class And : public CompoundFunction { return instance; } - Result> Test( - const arrow::Array& array, - const std::vector>& children) const override { + Result> Test(const arrow::Array& array, + const std::vector>& children, + arrow::MemoryPool* pool) const override { std::vector is_valid(array.length(), true); for (const auto& child : children) { auto child_filter = std::dynamic_pointer_cast(child); @@ -59,7 +59,7 @@ class And : public CompoundFunction { return Status::Invalid( fmt::format("child filter {} does not support Test", child->ToString())); } - PAIMON_ASSIGN_OR_RAISE(std::vector child_valid, child_filter->Test(array)); + PAIMON_ASSIGN_OR_RAISE(std::vector 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]); } diff --git a/src/paimon/common/predicate/compound_function.h b/src/paimon/common/predicate/compound_function.h index f093fc404..e32a3188b 100644 --- a/src/paimon/common/predicate/compound_function.h +++ b/src/paimon/common/predicate/compound_function.h @@ -22,6 +22,7 @@ #include #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> Test( - const arrow::Array& array, - const std::vector>& children) const = 0; + // `pool` is where any arrow buffer the evaluation allocates comes from, it must not be null. + virtual Result> Test(const arrow::Array& array, + const std::vector>& children, + arrow::MemoryPool* pool) const = 0; virtual Result Test(const std::shared_ptr& schema, const InternalRow& row, const std::vector>& children) const = 0; diff --git a/src/paimon/common/predicate/compound_predicate_impl.h b/src/paimon/common/predicate/compound_predicate_impl.h index 3cf6d9c7a..7ab5b613b 100644 --- a/src/paimon/common/predicate/compound_predicate_impl.h +++ b/src/paimon/common/predicate/compound_predicate_impl.h @@ -32,8 +32,9 @@ class CompoundPredicateImpl : public CompoundPredicate, public PredicateFilter { const std::vector>& children) : CompoundPredicate(compound_function, children) {} - Result> Test(const arrow::Array& array) const override { - return compound_function_.Test(array, children_); + Result> Test(const arrow::Array& array, + arrow::MemoryPool* pool) const override { + return compound_function_.Test(array, children_, pool); } Result Test(const std::shared_ptr& schema, diff --git a/src/paimon/common/predicate/leaf_function.h b/src/paimon/common/predicate/leaf_function.h index 9fdb4d8fe..bdc83699c 100644 --- a/src/paimon/common/predicate/leaf_function.h +++ b/src/paimon/common/predicate/leaf_function.h @@ -21,6 +21,7 @@ #include #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" @@ -28,8 +29,10 @@ 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> Test(const arrow::Array& array, - const std::vector& literals) const = 0; + const std::vector& literals, + arrow::MemoryPool* pool) const = 0; virtual Result Test(const Literal& value, const std::vector& literals) const = 0; diff --git a/src/paimon/common/predicate/leaf_predicate_impl.h b/src/paimon/common/predicate/leaf_predicate_impl.h index c9a8758f5..0d2984095 100644 --- a/src/paimon/common/predicate/leaf_predicate_impl.h +++ b/src/paimon/common/predicate/leaf_predicate_impl.h @@ -41,7 +41,8 @@ class LeafPredicateImpl : public LeafPredicate, public PredicateFilter { return leaf_function_; } - Result> Test(const arrow::Array& array) const override { + Result> Test(const arrow::Array& array, + arrow::MemoryPool* pool) const override { const auto& struct_array = checked_cast(array); if (field_index_ >= static_cast(struct_array.fields().size())) { return Status::Invalid( @@ -49,7 +50,7 @@ class LeafPredicateImpl : public LeafPredicate, public PredicateFilter { 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 Test(const std::shared_ptr& schema, diff --git a/src/paimon/common/predicate/leaf_unary_function.h b/src/paimon/common/predicate/leaf_unary_function.h index 544ce6738..04a939ae8 100644 --- a/src/paimon/common/predicate/leaf_unary_function.h +++ b/src/paimon/common/predicate/leaf_unary_function.h @@ -30,8 +30,8 @@ namespace paimon { class LeafUnaryFunction : public LeafFunction { public: - Result> Test(const arrow::Array& array, - const std::vector& literals) const override { + Result> Test(const arrow::Array& array, const std::vector& literals, + arrow::MemoryPool* pool) const override { std::vector is_valid(array.length(), false); PAIMON_ASSIGN_OR_RAISE( std::vector array_values, diff --git a/src/paimon/common/predicate/literal_converter.cpp b/src/paimon/common/predicate/literal_converter.cpp index 102348d70..4129fe706 100644 --- a/src/paimon/common/predicate/literal_converter.cpp +++ b/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 LiteralConverter::GetLiteralFromTimestampArray(const arrow: } return literals; } + +arrow::TimeUnit::type LiteralConverter::MinRequiredTimeUnit(const std::vector& literals) { + bool needs_micro = false; + for (const auto& literal : literals) { + if (literal.IsNull()) { + continue; + } + const int32_t nano = literal.GetValue().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 +Result> BuildArray(const std::vector& literals, + Extract extract, arrow::MemoryPool* pool) { + BuilderType builder(pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast(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 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> BuildDecimalArray(const std::vector& literals, + arrow::MemoryPool* pool) { + std::optional typed_value; + for (const auto& literal : literals) { + if (!literal.IsNull()) { + typed_value = literal.GetValue(); + 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 type, + arrow::Decimal128Type::Make(precision, scale)); + arrow::Decimal128Builder builder(type, pool); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast(literals.size()))); + for (const auto& literal : literals) { + if (literal.IsNull()) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.AppendNull()); + continue; + } + auto value = literal.GetValue(); + 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 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> BuildTimestampArray(const std::vector& 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(builder.type())); + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(static_cast(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(), time_type))); + } + std::shared_ptr array; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&array)); + return array; +} +} // namespace + +Result> LiteralConverter::ConvertLiteralsToArray( + const FieldType& field_type, const std::vector& literals, arrow::MemoryPool* pool) { + switch (field_type) { + case FieldType::BOOLEAN: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::TINYINT: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::SMALLINT: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::INT: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::BIGINT: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::FLOAT: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::DOUBLE: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::DATE: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, pool); + case FieldType::STRING: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, + pool); + case FieldType::BINARY: + return BuildArray( + literals, [](const Literal& literal) { return literal.GetValue(); }, + 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 diff --git a/src/paimon/common/predicate/literal_converter.h b/src/paimon/common/predicate/literal_converter.h index 09739533f..a3fde5881 100644 --- a/src/paimon/common/predicate/literal_converter.h +++ b/src/paimon/common/predicate/literal_converter.h @@ -27,6 +27,7 @@ #include #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> 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> ConvertLiteralsToArray( + const FieldType& field_type, const std::vector& 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& literals); + static Result ConvertLiteralsFromString(const FieldType& type, const std::string& value_str); diff --git a/src/paimon/common/predicate/literal_converter_test.cpp b/src/paimon/common/predicate/literal_converter_test.cpp index 7d03f9cb7..0d36cfe71 100644 --- a/src/paimon/common/predicate/literal_converter_test.cpp +++ b/src/paimon/common/predicate/literal_converter_test.cpp @@ -77,6 +77,17 @@ class LiteralConverterTest : public ::testing::Test { ASSERT_EQ(result, expected[i]); } } + + void CheckLiteralsToArray(const FieldType& type, const std::vector& literals, + const std::shared_ptr& expected) const { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr array, + LiteralConverter::ConvertLiteralsToArray(type, literals, arrow::default_memory_pool())); + ASSERT_TRUE(array->Equals(*expected)) + << "actual: " << array->ToString() << ", expected: " << expected->ToString(); + // the array reads back as the literals it was built from, the null ones included + CheckResult(array, literals); + } }; TEST_F(LiteralConverterTest, TestBooleanLiteral) { @@ -438,4 +449,150 @@ TEST_F(LiteralConverterTest, TestDictType) { Literal(FieldType::STRING, "foo", 3), Literal(FieldType::STRING)})); } +TEST_F(LiteralConverterTest, TestLiteralsToArray) { + // Every writable field type converts to an array and back, the null literal included, so each + // case asserts `ConvertLiteralsToArray` and `ConvertLiteralsFromArray` agree on every type. + CheckLiteralsToArray( + FieldType::BOOLEAN, {Literal(true), Literal(FieldType::BOOLEAN), Literal(false)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::boolean(), R"([true, null, false])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::TINYINT, + {Literal(static_cast(4)), Literal(FieldType::TINYINT), + Literal(static_cast(-5))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int8(), R"([4, null, -5])").ValueOrDie()); + CheckLiteralsToArray( + FieldType::SMALLINT, + {Literal(static_cast(45)), Literal(FieldType::SMALLINT), + Literal(static_cast(-55))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int16(), R"([45, null, -55])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::INT, {Literal(456), Literal(FieldType::INT), Literal(-567)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([456, null, -567])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::BIGINT, {Literal(4l), Literal(FieldType::BIGINT), Literal(-5l)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), R"([4, null, -5])").ValueOrDie()); + CheckLiteralsToArray( + FieldType::FLOAT, {Literal(4.5f), Literal(FieldType::FLOAT), Literal(-5.25f)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::float32(), R"([4.5, null, -5.25])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::DOUBLE, {Literal(4.5), Literal(FieldType::DOUBLE), Literal(-5.25)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::float64(), R"([4.5, null, -5.25])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::DATE, + {Literal(FieldType::DATE, 0), Literal(FieldType::DATE), Literal(FieldType::DATE, -5)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::date32(), R"([0, null, -5])") + .ValueOrDie()); + std::string str = "苹果"; + CheckLiteralsToArray( + FieldType::STRING, + {Literal(FieldType::STRING, "apple", 5), Literal(FieldType::STRING), + Literal(FieldType::STRING, str.data(), str.size())}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["apple", null, "苹果"])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::BINARY, + {Literal(FieldType::BINARY, "apple", 5), Literal(FieldType::BINARY), + Literal(FieldType::BINARY, str.data(), str.size())}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::binary(), R"(["apple", null, "苹果"])") + .ValueOrDie()); + // a decimal is written with the precision and the scale its literals carry + CheckLiteralsToArray( + FieldType::DECIMAL, + {Literal(Decimal(21, 3, DecimalUtils::StrToInt128("-123456789987654321234").value())), + Literal(FieldType::DECIMAL), Literal(Decimal(21, 3, 123456))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::decimal128(21, 3), + R"(["-123456789987654321.234", null, "123.456"])") + .ValueOrDie()); + // a timestamp is written with the finest time unit that keeps the values, one per case, and + // the JSON reader takes a raw integer in the unit of the type it is given + CheckLiteralsToArray( + FieldType::TIMESTAMP, + {Literal(Timestamp::FromEpochMillis(123000l)), Literal(FieldType::TIMESTAMP), + Literal(Timestamp::FromEpochMillis(-456000l))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::timestamp(arrow::TimeUnit::MILLI), + R"([123000, null, -456000])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::TIMESTAMP, + {Literal(Timestamp(123000l, 1000)), Literal(FieldType::TIMESTAMP), + Literal(Timestamp(-456000l, 2000))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::timestamp(arrow::TimeUnit::MICRO), + R"([123000001, null, -455999998])") + .ValueOrDie()); + CheckLiteralsToArray( + FieldType::TIMESTAMP, + {Literal(Timestamp(123000l, 456789)), Literal(FieldType::TIMESTAMP), + Literal(Timestamp(-456000l, 1))}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::timestamp(arrow::TimeUnit::NANO), + R"([123000456789, null, -455999999999])") + .ValueOrDie()); +} + +TEST_F(LiteralConverterTest, TestLiteralsToArrayWithoutValue) { + // no literal at all still gives an array typed after the field type + CheckLiteralsToArray( + FieldType::INT, {}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([])").ValueOrDie()); + // and a field type is all a null literal needs to be written + CheckLiteralsToArray( + FieldType::STRING, {Literal(FieldType::STRING), Literal(FieldType::STRING)}, + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"([null, null])").ValueOrDie()); +} + +TEST_F(LiteralConverterTest, TestLiteralsToArrayUnsupportedType) { + // a timestamp takes its arrow type from a literal, so it needs one that is not null + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray(FieldType::TIMESTAMP, {}, + arrow::default_memory_pool()), + "without a non null literal"); + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray( + FieldType::TIMESTAMP, {Literal(FieldType::TIMESTAMP)}, arrow::default_memory_pool()), + "without a non null literal"); + // a decimal takes its arrow type from a literal, so it needs one that is not null + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray(FieldType::DECIMAL, {}, + arrow::default_memory_pool()), + "without a non null literal"); + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray(FieldType::DECIMAL, {Literal(FieldType::DECIMAL)}, + arrow::default_memory_pool()), + "without a non null literal"); + // and every non null literal has to carry the same precision and scale + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray( + FieldType::DECIMAL, {Literal(Decimal(21, 3, 123456)), Literal(Decimal(21, 5, 123456))}, + arrow::default_memory_pool()), + "do not share one precision and scale"); + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray( + FieldType::DECIMAL, {Literal(Decimal(21, 3, 123456)), Literal(Decimal(20, 3, 123456))}, + arrow::default_memory_pool()), + "do not share one precision and scale"); + // a precision arrow rejects is reported instead of checked fatally + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray(FieldType::DECIMAL, {Literal(Decimal(0, 0, 1))}, + arrow::default_memory_pool()), + "Decimal precision out of range"); + // the field types no caller asks for yet + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray(FieldType::BLOB, {}, arrow::default_memory_pool()), + "Not support converting literals of BLOB type to an arrow array"); + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray(FieldType::ARRAY, {}, + arrow::default_memory_pool()), + "Not support converting literals of ARRAY type to an arrow array"); + ASSERT_NOK_WITH_MSG( + LiteralConverter::ConvertLiteralsToArray(FieldType::MAP, {}, arrow::default_memory_pool()), + "Not support converting literals of MAP type to an arrow array"); + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray(FieldType::STRUCT, {}, + arrow::default_memory_pool()), + "Not support converting literals of STRUCT type to an arrow array"); + ASSERT_NOK_WITH_MSG(LiteralConverter::ConvertLiteralsToArray(FieldType::UNKNOWN, {}, + arrow::default_memory_pool()), + "Not support converting literals of UNKNOWN"); +} + } // namespace paimon::test diff --git a/src/paimon/common/predicate/multi_literals_leaf_function.cpp b/src/paimon/common/predicate/multi_literals_leaf_function.cpp new file mode 100644 index 000000000..c8a41812d --- /dev/null +++ b/src/paimon/common/predicate/multi_literals_leaf_function.cpp @@ -0,0 +1,245 @@ +/* + * 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 +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_primitive.h" +#include "arrow/compute/api_scalar.h" +#include "arrow/compute/exec.h" +#include "arrow/datum.h" +#include "arrow/memory_pool.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" +#include "paimon/data/decimal.h" + +namespace paimon { +namespace { +/// The field types `arrow::compute::is_in` can stand in for `Literal::CompareTo`. These are also +/// the field types `LiteralConverter::ConvertLiteralsToArray` writes today, and the check keeps one +/// it starts writing from reaching `is_in` on its own. +/// +/// `TIMESTAMP` is parameterized by a time unit, but a timestamp literal settles one itself: the +/// finest unit that keeps the values of the literals, which `MakeInValueSet` only probes against +/// a column of that very unit. `DECIMAL` is parameterized too, by precision and scale, but a +/// decimal literal carries both, and `MakeInValueSet` only probes a column that carries the same +/// scale. `FLOAT` and `DOUBLE` agree on every value but NaN, which `MakeInValueSet` keeps off this +/// path on its own. +bool CanProbeWithIsIn(FieldType field_type) { + switch (field_type) { + case FieldType::BOOLEAN: + case FieldType::TINYINT: + case FieldType::SMALLINT: + case FieldType::INT: + case FieldType::BIGINT: + case FieldType::FLOAT: + case FieldType::DOUBLE: + case FieldType::DATE: + case FieldType::STRING: + case FieldType::BINARY: + case FieldType::DECIMAL: + case FieldType::TIMESTAMP: + return true; + default: + return false; + } +} + +/// Whether `data_type` is the arrow type of a decimal column of `scale`. +bool IsDecimalOfScale(const arrow::DataType& data_type, int32_t scale) { + return data_type.id() == arrow::Type::DECIMAL128 && + checked_cast(data_type).scale() == scale; +} + +/// Whether `data_type` is the arrow type of a timestamp column of `unit` without a time zone. +bool IsTimestampOfUnit(const arrow::DataType& data_type, arrow::TimeUnit::type unit) { + return data_type.id() == arrow::Type::TIMESTAMP && + checked_cast(data_type).unit() == unit && + checked_cast(data_type).timezone().empty(); +} + +/// Whether `literal` holds a floating point NaN. +bool IsNanLiteral(const Literal& literal) { + switch (literal.GetType()) { + case FieldType::FLOAT: + return std::isnan(literal.GetValue()); + case FieldType::DOUBLE: + return std::isnan(literal.GetValue()); + default: + return false; + } +} + +/// 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`. +/// @param data_type The arrow type of the column the predicate is evaluated against, which settles +/// whether a decimal or a timestamp value set can be probed against it. +/// @param pool The pool the value set is allocated from. +/// @return `nullptr` when the literals cannot be probed by `is_in`, which covers the field types +/// `CanProbeWithIsIn` rejects, a NaN literal, a decimal column of another scale, a +/// timestamp column of another unit or with a time zone, 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, + const arrow::DataType& data_type, + arrow::MemoryPool* pool) { + 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 is all nulls, which `is_in` ignores, but it still needs a + // type to compare against the column, and a null `Literal` carries its type too. + const Literal* typed_literal = &literals.front(); + for (const auto& literal : literals) { + if (!literal.IsNull()) { + typed_literal = &literal; + break; + } + } + const FieldType field_type = typed_literal->GetType(); + if (!CanProbeWithIsIn(field_type)) { + return nullptr; + } + // A decimal value set carries the precision and the scale of its literals, and `is_in` compares + // it against a column of another scale by casting one side to the other. Such a cast fails on a + // value that does not fit the other scale, where `Literal::CompareTo` rescales one value at a + // time and merely finds no match, so only a column carrying the scale of the literals is + // probed. Nothing but null literals leaves no scale to compare and no arrow type to write. + if (field_type == FieldType::DECIMAL && + (typed_literal->IsNull() || + !IsDecimalOfScale(data_type, typed_literal->GetValue().Scale()))) { + return nullptr; + } + // A timestamp value set carries the time unit the values of the literals need, and `is_in` + // compares it against a column of another unit by casting one side to the other. Casting to a + // coarser unit fails on a value the unit does not keep, casting to a finer one can overflow + // int64, where `Literal::CompareTo` merely finds no match, so only a column of that very unit + // is probed. A column with a time zone is one as well: `is_in` refuses to compare a zoned + // timestamp against an unzoned one, and a literal has no zone to speak of. Nothing but null + // literals leaves no unit to compare and no arrow type to write. + if (field_type == FieldType::TIMESTAMP && + (typed_literal->IsNull() || + !IsTimestampOfUnit(data_type, LiteralConverter::MinRequiredTimeUnit(literals)))) { + return nullptr; + } + 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; + } + // `is_in` hashes the raw bits of a float, so a NaN literal would only match the column NaNs + // carrying the very same bit pattern, while `FieldsComparator::CompareFloatingPoint` makes + // every NaN equal. Keep a NaN literal on the row by row path. With none in the value set + // the two agree, because a column NaN then matches no literal either way. + if (!literal.IsNull() && IsNanLiteral(literal)) { + 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; + } + } + Result> value_set = + LiteralConverter::ConvertLiteralsToArray(field_type, literals, pool); + // A failure only says the value set is not there, and the row by row path still is. Decimal + // literals of mixed precision and scale end up here, anything else is an arrow failure. + if (!value_set.ok()) { + return nullptr; + } + return std::move(value_set).value(); +} + +/// Probes every non-null row of `array` against `value_set`. +/// +/// @param negate `false` for `IN` semantics, `true` for `NOT IN` semantics. +/// @param pool The pool the match bitmap `is_in` writes is allocated from. +/// @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> ProbeInValueSet(const arrow::Array& array, const arrow::Array& value_set, + bool negate, arrow::MemoryPool* pool) { + // `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); + arrow::compute::ExecContext exec_context(pool); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum matches, arrow::compute::IsIn(arrow::Datum(array), options, &exec_context)); + // `make_array` hands out a new `shared_ptr` that owns the array, so it has to be held for as + // long as the array is read. Binding a reference straight to what it points at would drop the + // last owner at the end of the statement and leave that reference dangling. + std::shared_ptr matches_array = matches.make_array(); + const auto& matched = checked_cast(*matches_array); + std::vector 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(matched.Value(i) != negate); + } + return is_valid; +} +} // namespace + +Result> MultiLiteralsLeafFunction::Test(const arrow::Array& array, + const std::vector& literals, + arrow::MemoryPool* pool) 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 value_set = + MakeInValueSet(literals, negate, *array.type(), pool); + if (value_set != nullptr) { + return ProbeInValueSet(array, *value_set, negate, pool); + } + } + + // 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 array_values, + LiteralConverter::ConvertLiteralsFromArray(array, /*own_data=*/false)); + std::vector 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 diff --git a/src/paimon/common/predicate/multi_literals_leaf_function.h b/src/paimon/common/predicate/multi_literals_leaf_function.h index 804466471..c94b23843 100644 --- a/src/paimon/common/predicate/multi_literals_leaf_function.h +++ b/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: - Result> Test(const arrow::Array& array, - const std::vector& literals) const override { - PAIMON_ASSIGN_OR_RAISE( - std::vector array_values, - LiteralConverter::ConvertLiteralsFromArray(array, /*own_data=*/false)); - std::vector 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; - } + /// 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> Test(const arrow::Array& array, const std::vector& literals, + arrow::MemoryPool* pool) const override; Result Test(int64_t row_count, const Literal& min_value, const Literal& max_value, const std::optional& null_count, diff --git a/src/paimon/common/predicate/multi_literals_leaf_function_test.cpp b/src/paimon/common/predicate/multi_literals_leaf_function_test.cpp new file mode 100644 index 000000000..981334eed --- /dev/null +++ b/src/paimon/common/predicate/multi_literals_leaf_function_test.cpp @@ -0,0 +1,524 @@ +/* + * 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 +#include +#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/common/predicate/in.h" +#include "paimon/common/predicate/literal_converter.h" +#include "paimon/common/predicate/not_in.h" +#include "paimon/common/utils/decimal_utils.h" +#include "paimon/data/decimal.h" +#include "paimon/data/timestamp.h" +#include "paimon/defs.h" +#include "paimon/testing/utils/testharness.h" + +namespace paimon::test { +class MultiLiteralsLeafFunctionTest : 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()); + } + + // A decimal literal of `unscaled` as written by the digits of `precision` and `scale`. + static Literal DecimalLiteral(int32_t precision, int32_t scale, const std::string& unscaled) { + EXPECT_OK_AND_ASSIGN(Decimal::int128_t value, DecimalUtils::StrToInt128(unscaled)); + return Literal(Decimal(precision, scale, value)); + } + + // A timestamp literal of `millis` since the epoch and `nanos` within the millisecond. + static Literal TimestampLiteral(int64_t millis, int32_t nanos = 0) { + return Literal(Timestamp::FromEpochMillis(millis, nanos)); + } + + // Evaluates the whole batch and returns the per row result, asserting the call succeeded. + static std::vector Eval(const LeafFunction& function, + const std::vector& literals, + const std::shared_ptr& array) { + EXPECT_OK_AND_ASSIGN(std::vector is_valid, + function.Test(*array, literals, arrow::default_memory_pool())); + return is_valid; + } + + static std::vector EvalIn(const std::vector& literals, + const std::shared_ptr& array) { + return Eval(In::Instance(), literals, array); + } + + static std::vector EvalNotIn(const std::vector& literals, + const std::shared_ptr& array) { + return Eval(NotIn::Instance(), literals, array); + } +}; + +TEST_F(MultiLiteralsLeafFunctionTest, TestInt) { + const std::vector literals = {Literal(1), Literal(2), Literal(3), Literal(5)}; + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([0, 1, 2, 3, 4, 5, 6, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({0, 1, 1, 1, 0, 1, 0, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({1, 0, 0, 0, 1, 0, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestBigInt) { + const std::vector literals = {Literal(int64_t{-1000000}), Literal(int64_t{0}), + Literal(int64_t{1000000})}; + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-1000000, -999999, 0, 1, 1000000, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 1, 0, 1, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 1, 0, 1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestInt64Boundaries) { + const std::vector literals = {Literal(std::numeric_limits::min()), + Literal(std::numeric_limits::max())}; + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-9223372036854775808, 0, 9223372036854775807, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 1, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestTinyIntAndSmallInt) { + auto tinyint_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int8(), R"([-128, -127, 0, 126, 127])") + .ValueOrDie(); + ASSERT_EQ( + EvalIn({Literal(int8_t{-128}), Literal(int8_t{0}), Literal(int8_t{127})}, tinyint_array), + std::vector({1, 0, 1, 0, 1})); + + auto smallint_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int16(), R"([-30000, 0, 30000, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({Literal(int16_t{-30000}), Literal(int16_t{30000})}, smallint_array), + std::vector({1, 0, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestDate) { + const std::vector literals = {Literal(FieldType::DATE, 100), + Literal(FieldType::DATE, 20000)}; + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::date32(), R"([100, 101, 20000, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 1, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestBoolean) { + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::boolean(), R"([true, false, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({Literal(true)}, array), std::vector({1, 0, 0})); + ASSERT_EQ(EvalNotIn({Literal(true)}, array), std::vector({0, 1, 0})); + + const std::vector both = {Literal(true), Literal(false)}; + ASSERT_EQ(EvalIn(both, array), std::vector({1, 1, 0})); + ASSERT_EQ(EvalNotIn(both, array), std::vector({0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestString) { + const std::vector literals = {StringLiteral("apple"), StringLiteral(""), + StringLiteral("banana")}; + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::utf8(), R"(["apple", "", "banana", "cherry", "app", null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 1, 1, 0, 0, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 0, 0, 1, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestStringWithoutEmptyLiteral) { + // An empty column value must not match when no empty literal was given. + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "", "b"])").ValueOrDie(); + ASSERT_EQ(EvalIn({StringLiteral("a")}, array), std::vector({1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestBinary) { + // Binary literals keep their embedded zero bytes, the value set must not truncate them. The + // JSON reader cannot spell a value like that, so the column is built with a builder. + const std::vector literals = {BinaryLiteral(std::string("\x00\x01", 2)), + BinaryLiteral("xyz")}; + + 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(EvalIn(literals, array), std::vector({1, 1, 0, 0, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 0, 1, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestDictionaryString) { + const std::vector literals = {StringLiteral("a"), StringLiteral("c")}; + 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()); + std::shared_ptr array = + arrow::DictionaryArray::FromArrays(dict_type, indices, dictionary).ValueOrDie(); + + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 1, 1, 0, 1})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 1, 0, 0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestLargeStringDictionary) { + 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()); + std::shared_ptr array = + arrow::DictionaryArray::FromArrays(dict_type, indices, dictionary).ValueOrDie(); + + ASSERT_EQ(EvalIn({StringLiteral("c")}, array), std::vector({0, 1, 0, 0})); + ASSERT_EQ(EvalNotIn({StringLiteral("c")}, array), std::vector({1, 0, 0, 1})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, 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()); + std::shared_ptr array = + arrow::DictionaryArray::FromArrays(dict_type, indices, dictionary).ValueOrDie(); + + const std::vector with_empty = {StringLiteral("a"), StringLiteral("")}; + ASSERT_EQ(EvalIn(with_empty, array), std::vector({0, 1, 0, 0})); + ASSERT_EQ(EvalNotIn(with_empty, array), std::vector({0, 0, 0, 0})); + + ASSERT_EQ(EvalIn({StringLiteral("a")}, array), std::vector({0, 1, 0, 0})); + ASSERT_EQ(EvalNotIn({StringLiteral("a")}, array), std::vector({0, 0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestNullLiteralIgnoredForIn) { + const std::vector literals = {Literal(int64_t{1}), Literal(FieldType::BIGINT), + Literal(int64_t{3})}; + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), R"([1, 2, 3, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 1, 0})); + // A null literal makes `NOT IN` false for every row. + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestOnlyNullLiterals) { + const std::vector int_literals = {Literal(FieldType::BIGINT)}; + auto int_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::int64(), R"([-9007199254740993, 0, 9007199254740993, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(int_literals, int_array), std::vector({0, 0, 0, 0})); + ASSERT_EQ(EvalNotIn(int_literals, int_array), std::vector({0, 0, 0, 0})); + + auto string_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["", "a", null])").ValueOrDie(); + ASSERT_EQ(EvalIn({Literal(FieldType::STRING)}, string_array), std::vector({0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestEmptyLiterals) { + auto array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([1, 2, null])").ValueOrDie(); + ASSERT_EQ(EvalIn({}, array), std::vector({0, 0, 0})); + ASSERT_EQ(EvalNotIn({}, array), std::vector({1, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestTimestamp) { + // A timestamp literal settles the time unit of its value set itself, the finest unit that + // keeps every literal, so a column of that very unit is probed by the value set. + auto milli_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::timestamp(arrow::TimeUnit::MILLI), R"([1000, 2000, 3000, null])") + .ValueOrDie(); + const std::vector milli_literals = {TimestampLiteral(1000), TimestampLiteral(3000)}; + ASSERT_EQ(EvalIn(milli_literals, milli_array), std::vector({1, 0, 1, 0})); + ASSERT_EQ(EvalNotIn(milli_literals, milli_array), std::vector({0, 1, 0, 0})); + + // A nanos-of-millisecond outside a microsecond forces the nanosecond unit, which the values + // keep through the value set. + auto nano_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::timestamp(arrow::TimeUnit::NANO), + R"([1000000000000, 1000000456789, null])") + .ValueOrDie(); + const std::vector nano_literals = {TimestampLiteral(1000000, 456789)}; + ASSERT_EQ(EvalIn(nano_literals, nano_array), std::vector({0, 1, 0})); + ASSERT_EQ(EvalNotIn(nano_literals, nano_array), std::vector({1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestDecimal) { + // A decimal literal carries the precision and the scale of its own value, so a column of that + // scale is probed by the value set. + const std::vector literals = {DecimalLiteral(10, 2, "100"), + DecimalLiteral(10, 2, "-250")}; + auto array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::decimal128(10, 2), R"(["1.00", "1.01", "-2.50", "0.00", null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 1, 0, 0})); + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 1, 0, 1, 0})); + + // A column of the same scale but another precision is probed too, the value set is widened to + // it without ever changing a value. + auto wider_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::decimal128(38, 2), R"(["1.00", "1.01", "-2.50", null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, wider_array), std::vector({1, 0, 1, 0})); + ASSERT_EQ(EvalNotIn(literals, wider_array), std::vector({0, 1, 0, 0})); + + // A value beyond 64 bits keeps its high bits through the value set. + const std::vector wide_literals = {DecimalLiteral(38, 2, "12345678998765432134567"), + DecimalLiteral(38, 2, "-12345678998765432134567")}; + ASSERT_EQ(EvalIn(wide_literals, + arrow::ipc::internal::json::ArrayFromJSON( + arrow::decimal128(38, 2), + R"(["123456789987654321345.67", "-123456789987654321345.67", "1.00"])") + .ValueOrDie()), + std::vector({1, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestDecimalNullLiterals) { + const std::vector literals = {DecimalLiteral(10, 2, "100"), + Literal(FieldType::DECIMAL)}; + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::decimal128(10, 2), + R"(["1.00", "2.00", null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 0})); + // A null literal makes `NOT IN` false for every row. + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 0, 0})); + + // Nothing but null literals leaves no precision and scale to write the value set with, so the + // row by row path takes over and matches nothing either way. + ASSERT_EQ(EvalIn({Literal(FieldType::DECIMAL)}, array), std::vector({0, 0, 0})); + ASSERT_EQ(EvalNotIn({Literal(FieldType::DECIMAL)}, array), std::vector({0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestTimestampNullLiterals) { + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::timestamp(arrow::TimeUnit::MILLI), + R"([1000, 2000, null])") + .ValueOrDie(); + const std::vector literals = {TimestampLiteral(1000), Literal(FieldType::TIMESTAMP)}; + ASSERT_EQ(EvalIn(literals, array), std::vector({1, 0, 0})); + // A null literal makes `NOT IN` false for every row. + ASSERT_EQ(EvalNotIn(literals, array), std::vector({0, 0, 0})); + + // Nothing but null literals leaves no time unit to write the value set with, so the row by row + // path takes over and matches nothing either way. + ASSERT_EQ(EvalIn({Literal(FieldType::TIMESTAMP)}, array), std::vector({0, 0, 0})); + ASSERT_EQ(EvalNotIn({Literal(FieldType::TIMESTAMP)}, array), std::vector({0, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestDecimalOffTheValueSetPath) { + // `Literal::CompareTo` compares two decimals by value, so a literal of another scale still + // matches the same number. The value set would have `is_in` cast one side to the other, so a + // column of another scale keeps the row by row path and the result stays the same either way. + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::decimal128(10, 2), + R"(["1.00", "2.00", null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({DecimalLiteral(20, 4, "10000")}, array), std::vector({1, 0, 0})); + ASSERT_EQ(EvalNotIn({DecimalLiteral(20, 4, "10000")}, array), std::vector({0, 1, 0})); + + // Literals that do not share one precision and scale cannot be one arrow array, so they keep + // the row by row path as well. + ASSERT_EQ(EvalIn({DecimalLiteral(10, 2, "100"), DecimalLiteral(20, 4, "20000")}, array), + std::vector({1, 1, 0})); + + // Casting between those scales is what the row by row path spares: neither side of this pair + // fits the scale of the other, which would make `is_in` report an error where comparing by + // value merely finds no match. + auto lossy_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::decimal128(38, 2), R"(["1.01", "2.00"])") + .ValueOrDie(); + const std::vector huge_literals = { + DecimalLiteral(38, 0, "99999999999999999999999999999999999999")}; + ASSERT_EQ(EvalIn(huge_literals, lossy_array), std::vector({0, 0})); + ASSERT_EQ(EvalNotIn(huge_literals, lossy_array), std::vector({1, 1})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestTimestampOffTheValueSetPath) { + // A column of a coarser unit than the literals need keeps the row by row path, because `is_in` + // would cast the value set to the column unit and fail on a value the unit does not keep. A + // value that needs the finer unit cannot exist in the coarser column, so comparing by value + // merely finds no match. + auto milli_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::timestamp(arrow::TimeUnit::MILLI), R"([1000, 2000, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({TimestampLiteral(1, 500000)}, milli_array), std::vector({0, 0, 0})); + ASSERT_EQ(EvalNotIn({TimestampLiteral(1, 500000)}, milli_array), std::vector({1, 1, 0})); + + // A column of a finer unit than the literals need keeps the row by row path as well, casting + // to it can overflow int64 where comparing by value still matches the same instant. + auto nano_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::timestamp(arrow::TimeUnit::NANO), R"([1000000000, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({TimestampLiteral(1000)}, nano_array), std::vector({1, 0})); + ASSERT_EQ(EvalNotIn({TimestampLiteral(1000)}, nano_array), std::vector({0, 0})); + + // A column with a time zone as well: `is_in` refuses to compare a zoned timestamp against an + // unzoned value set, so the row by row path compares by value instead. + auto zoned_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::timestamp(arrow::TimeUnit::MILLI, "UTC"), R"([1000, 2000, null])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({TimestampLiteral(1000)}, zoned_array), std::vector({1, 0, 0})); + ASSERT_EQ(EvalNotIn({TimestampLiteral(1000)}, zoned_array), std::vector({0, 1, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestFloatAndDouble) { + auto double_array = arrow::ipc::internal::json::ArrayFromJSON( + arrow::float64(), R"([1.0, 2.5, -3.25, Inf, -Inf, null])") + .ValueOrDie(); + const std::vector double_literals = {Literal(1.0), Literal(-3.25), + Literal(std::numeric_limits::infinity())}; + ASSERT_EQ(EvalIn(double_literals, double_array), std::vector({1, 0, 1, 1, 0, 0})); + ASSERT_EQ(EvalNotIn(double_literals, double_array), std::vector({0, 1, 0, 0, 1, 0})); + + auto float_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::float32(), R"([1.5, 2.5, -Inf, null])") + .ValueOrDie(); + const std::vector float_literals = {Literal(1.5f), + Literal(-std::numeric_limits::infinity())}; + ASSERT_EQ(EvalIn(float_literals, float_array), std::vector({1, 0, 1, 0})); + ASSERT_EQ(EvalNotIn(float_literals, float_array), std::vector({0, 1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestNegativeZeroDoesNotMatchPositiveZero) { + // `FieldsComparator::CompareFloatingPoint` orders `-0.0 < +0.0`, so they are two distinct + // values. `is_in` agrees only because it hashes the raw bits of a float, which arrow's own + // comment marks as something it would rather change. Should it ever hash equal floats alike, + // `-0.0` would start matching `IN (0.0)` and this assertion is what catches the divergence. + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::float64(), R"([0.0, -0.0, 1.0])") + .ValueOrDie(); + ASSERT_EQ(EvalIn({Literal(0.0)}, array), std::vector({1, 0, 0})); + ASSERT_EQ(EvalIn({Literal(-0.0)}, array), std::vector({0, 1, 0})); + ASSERT_EQ(EvalNotIn({Literal(0.0)}, array), std::vector({0, 1, 1})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestNanLiteralStaysOffTheValueSetPath) { + // `is_in` hashes the raw bits of a float, so a NaN literal would only match the column NaNs + // carrying the very same bit pattern, while `FieldsComparator::CompareFloatingPoint` makes + // every NaN equal. A NaN literal therefore keeps the row by row path, where the sign flipped + // NaN below still matches. The JSON reader cannot spell a sign flipped NaN, so the column is + // built with a builder. + const double canonical_nan = std::numeric_limits::quiet_NaN(); + arrow::DoubleBuilder builder; + ASSERT_TRUE(builder.Append(canonical_nan).ok()); + ASSERT_TRUE(builder.Append(-canonical_nan).ok()); + ASSERT_TRUE(builder.Append(1.0).ok()); + ASSERT_TRUE(builder.AppendNull().ok()); + std::shared_ptr array; + ASSERT_TRUE(builder.Finish(&array).ok()); + + ASSERT_EQ(EvalIn({Literal(canonical_nan)}, array), std::vector({1, 1, 0, 0})); + ASSERT_EQ(EvalNotIn({Literal(canonical_nan)}, array), std::vector({0, 0, 1, 0})); + ASSERT_EQ(EvalIn({Literal(1.0), Literal(canonical_nan)}, array), + std::vector({1, 1, 1, 0})); + + // With no NaN in the value set the two paths agree, because a column NaN then matches no + // literal either way. + ASSERT_EQ(EvalIn({Literal(1.0)}, array), std::vector({0, 0, 1, 0})); + ASSERT_EQ(EvalNotIn({Literal(1.0)}, array), std::vector({1, 1, 0, 0})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestMixedLiteralTypesReportTheError) { + // Literals of mixed types make `Literal::CompareTo` fail, so they must keep reporting the error + // through the row by row path instead of being built into one typed value set. The column value + // matches none of them, otherwise the comparison would stop before reaching the odd literal. + auto array = arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), R"([0])").ValueOrDie(); + ASSERT_NOK(In::Instance().Test(*array, {Literal(1), Literal(int64_t{2})}, + arrow::default_memory_pool())); + ASSERT_NOK(In::Instance().Test(*array, {Literal(1), StringLiteral("a")}, + arrow::default_memory_pool())); + ASSERT_NOK(NotIn::Instance().Test(*array, {Literal(1), Literal(int64_t{2})}, + arrow::default_memory_pool())); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestDictionaryLayoutsLiteralConversionRejects) { + // 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(); + std::shared_ptr int64_dict_array = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::int64()), + int32_indices, int64_dictionary) + .ValueOrDie(); + ASSERT_NOK(LiteralConverter::ConvertLiteralsFromArray(*int64_dict_array, /*own_data=*/false)); + ASSERT_EQ(EvalIn({Literal(int64_t{10})}, int64_dict_array), std::vector({1, 0})); + ASSERT_EQ(EvalNotIn({Literal(int64_t{10})}, int64_dict_array), std::vector({0, 1})); + + 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(); + std::shared_ptr int8_dict_array = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int8(), arrow::utf8()), + int8_indices, string_dictionary) + .ValueOrDie(); + ASSERT_NOK(LiteralConverter::ConvertLiteralsFromArray(*int8_dict_array, /*own_data=*/false)); + ASSERT_EQ(EvalIn({StringLiteral("a")}, int8_dict_array), std::vector({1, 0})); + ASSERT_EQ(EvalNotIn({StringLiteral("a")}, int8_dict_array), std::vector({0, 1})); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestProbeFailsOnUnrelatedArrowType) { + // Nothing casts an int32 value set to a string column, so the probe reports the mismatch + // instead of matching by chance. `PredicateValidator` already rejects such a predicate against + // the schema, this only pins what the evaluation does if one reaches it anyway. + auto string_array = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["1"])").ValueOrDie(); + ASSERT_NOK(In::Instance().Test(*string_array, {Literal(1)}, arrow::default_memory_pool())); +} + +TEST_F(MultiLiteralsLeafFunctionTest, TestSlicedArray) { + 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(EvalIn({Literal(2), Literal(4)}, sliced), std::vector({0, 1, 0, 0})); + ASSERT_EQ(EvalNotIn({Literal(2), Literal(4)}, sliced), std::vector({1, 0, 1, 0})); + + 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(EvalIn({StringLiteral("c")}, sliced_string), std::vector({0, 1, 0})); +} +} // namespace paimon::test diff --git a/src/paimon/common/predicate/null_false_leaf_binary_function.h b/src/paimon/common/predicate/null_false_leaf_binary_function.h index 87de32490..b9b6850b3 100644 --- a/src/paimon/common/predicate/null_false_leaf_binary_function.h +++ b/src/paimon/common/predicate/null_false_leaf_binary_function.h @@ -31,8 +31,8 @@ namespace paimon { class NullFalseLeafBinaryFunction : public LeafFunction { public: - Result> Test(const arrow::Array& array, - const std::vector& literals) const override { + Result> Test(const arrow::Array& array, const std::vector& literals, + arrow::MemoryPool* pool) const override { if (literals.size() < LITERAL_LIMIT) { return Status::Invalid("NullFalseLeafBinaryFunction needs single literal for field"); } diff --git a/src/paimon/common/predicate/or.h b/src/paimon/common/predicate/or.h index fb1809131..f74046407 100644 --- a/src/paimon/common/predicate/or.h +++ b/src/paimon/common/predicate/or.h @@ -45,9 +45,9 @@ class Or : public CompoundFunction { public: static const Or& Instance(); - Result> Test( - const arrow::Array& array, - const std::vector>& children) const override { + Result> Test(const arrow::Array& array, + const std::vector>& children, + arrow::MemoryPool* pool) const override { std::vector is_valid(array.length(), false); for (const auto& child : children) { auto child_filter = std::dynamic_pointer_cast(child); @@ -55,7 +55,7 @@ class Or : public CompoundFunction { return Status::Invalid( fmt::format("child filter {} does not support Test", child->ToString())); } - PAIMON_ASSIGN_OR_RAISE(std::vector child_valid, child_filter->Test(array)); + PAIMON_ASSIGN_OR_RAISE(std::vector 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]); } diff --git a/src/paimon/common/predicate/predicate_filter.h b/src/paimon/common/predicate/predicate_filter.h index 799ed1560..d14b47476 100644 --- a/src/paimon/common/predicate/predicate_filter.h +++ b/src/paimon/common/predicate/predicate_filter.h @@ -30,7 +30,10 @@ namespace paimon { class PredicateFilter : virtual public Predicate { public: /// @param array is the struct array of all fields - virtual Result> Test(const arrow::Array& array) const = 0; + /// @param pool is where any arrow buffer the evaluation allocates comes from, it must not be + /// null + virtual Result> Test(const arrow::Array& array, + arrow::MemoryPool* pool) const = 0; virtual Result Test(const std::shared_ptr& schema, const InternalRow& row) const = 0; virtual Result Test(const std::shared_ptr& schema, int64_t row_count, diff --git a/src/paimon/common/predicate/predicate_test.cpp b/src/paimon/common/predicate/predicate_test.cpp index fa4733cba..332fa276d 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" @@ -185,7 +188,7 @@ TEST_F(PredicateTest, TestInvalidFieldIndex) { arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); // with array - ASSERT_NOK_WITH_MSG(predicate->Test(*struct_array), + ASSERT_NOK_WITH_MSG(predicate->Test(*struct_array, arrow::default_memory_pool()), "field index 2 exceed field count 2 in struct array"); // with internal row @@ -208,7 +211,8 @@ TEST_F(PredicateTest, TestEqual) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 1, 0})); ASSERT_EQ(*predicate->Negate(), @@ -248,7 +252,8 @@ TEST_F(PredicateTest, TestEqualNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -276,7 +281,8 @@ TEST_F(PredicateTest, TestNotEqual) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0, 0})); auto predicate_negate = std::dynamic_pointer_cast(predicate->Negate()); @@ -312,7 +318,8 @@ TEST_F(PredicateTest, TestNotEqualNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row auto arrow_schema = arrow::schema(arrow::FieldVector({arrow::field("f0", bigint_type)})); @@ -340,7 +347,8 @@ TEST_F(PredicateTest, TestGreater) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 1, 0})); ASSERT_EQ(*predicate->Negate(), @@ -375,7 +383,8 @@ TEST_F(PredicateTest, TestGreaterNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -404,7 +413,8 @@ TEST_F(PredicateTest, TestGreaterOrEqual) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 1, 1, 0})); ASSERT_EQ(*predicate->Negate(), @@ -439,7 +449,8 @@ TEST_F(PredicateTest, TestGreaterOrEqualNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -467,7 +478,8 @@ TEST_F(PredicateTest, TestLess) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0, 0, 0})); ASSERT_EQ(*predicate->Negate(), @@ -501,7 +513,8 @@ TEST_F(PredicateTest, TestLessNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -529,7 +542,8 @@ TEST_F(PredicateTest, TestLessOrEqual) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 1, 0, 0})); ASSERT_EQ(*predicate->Negate(), @@ -563,7 +577,8 @@ TEST_F(PredicateTest, TestLessOrEqualNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -589,7 +604,8 @@ TEST_F(PredicateTest, TestIsNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 1})); ASSERT_EQ(*predicate->Negate(), *PredicateBuilder::IsNotNull( @@ -618,7 +634,8 @@ TEST_F(PredicateTest, TestIsNotNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0})); ASSERT_EQ(*predicate->Negate(), @@ -650,7 +667,8 @@ TEST_F(PredicateTest, TestIn) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0, 1, 0})); ASSERT_EQ(*predicate->Negate(), @@ -686,7 +704,8 @@ TEST_F(PredicateTest, TestInNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0, 1, 0})); // with internal row @@ -717,7 +736,8 @@ TEST_F(PredicateTest, TestNotIn) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 1, 0, 0})); ASSERT_EQ(*predicate->Negate(), @@ -756,7 +776,8 @@ TEST_F(PredicateTest, TestNotInNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 0, 0})); // with internal row @@ -797,7 +818,8 @@ TEST_F(PredicateTest, TestLargeIn) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0, 1, 0})); // with internal row @@ -837,7 +859,8 @@ TEST_F(PredicateTest, TestLargeInNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 0, 1, 0})); // with internal row @@ -876,7 +899,8 @@ TEST_F(PredicateTest, TestLargeNotIn) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 1, 0, 0})); // with internal row @@ -919,7 +943,8 @@ TEST_F(PredicateTest, TestLargeNotInNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 0, 0})); // with internal row @@ -938,6 +963,98 @@ 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, arrow::default_memory_pool())); + 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, arrow::default_memory_pool())); + 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 builds a new predicate around the same literals, 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, arrow::default_memory_pool())); + 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) { + // The value set carries the whole int64 range, the boundary values included. + 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, arrow::default_memory_pool())); + ASSERT_EQ(in_valid, std::vector({1, 0, 1, 0})); + ASSERT_OK_AND_ASSIGN(auto not_in_valid, + not_in_predicate->Test(*struct_array, arrow::default_memory_pool())); + ASSERT_EQ(not_in_valid, std::vector({0, 1, 0, 0})); +} + TEST_F(PredicateTest, TestAnd) { auto bigint_type = arrow::int64(); ASSERT_OK_AND_ASSIGN( @@ -957,7 +1074,8 @@ TEST_F(PredicateTest, TestAnd) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 1, 0})); ASSERT_OK_AND_ASSIGN( @@ -1003,7 +1121,8 @@ TEST_F(PredicateTest, TestOr) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 1, 1, 1})); ASSERT_OK_AND_ASSIGN( @@ -1045,7 +1164,8 @@ TEST_F(PredicateTest, TestBetween) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({1, 1, 1, 0, 0, 0})); auto less_than = PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"f0", @@ -1087,7 +1207,8 @@ TEST_F(PredicateTest, TestBetweenNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -1120,7 +1241,8 @@ TEST_F(PredicateTest, TestStartsWith) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 1, 0})); ASSERT_EQ(predicate->Negate(), nullptr); @@ -1165,7 +1287,8 @@ TEST_F(PredicateTest, TestStartsWithNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -1199,7 +1322,8 @@ TEST_F(PredicateTest, TestEndsWith) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 1, 0})); ASSERT_EQ(predicate->Negate(), nullptr); @@ -1236,7 +1360,8 @@ TEST_F(PredicateTest, TestEndsWithNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row @@ -1270,7 +1395,8 @@ TEST_F(PredicateTest, TestContains) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0, 1, 0})); ASSERT_EQ(predicate->Negate(), nullptr); @@ -1307,7 +1433,8 @@ TEST_F(PredicateTest, TestContainsNull) { std::shared_ptr struct_array = arrow::StructArray::Make({f0, f1}, src_type->fields()).ValueOrDie(); - ASSERT_OK_AND_ASSIGN(auto is_valid, predicate->Test(*struct_array)); + ASSERT_OK_AND_ASSIGN(auto is_valid, + predicate->Test(*struct_array, arrow::default_memory_pool())); ASSERT_EQ(is_valid, std::vector({0, 0})); // with internal row diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 23b6960fd..a5984c703 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -90,7 +90,8 @@ Result LateMaterializingFileBatchReader::FilterProbeBatch( const std::shared_ptr& array, const std::shared_ptr& bound_filter) { // TODO(zhouhonfeng.zhf): use arrow::compute::Filter instead of PredicateFilter - PAIMON_ASSIGN_OR_RAISE(std::vector results, bound_filter->Test(*array)); + PAIMON_ASSIGN_OR_RAISE(std::vector results, + bound_filter->Test(*array, arrow_pool_.get())); if (results.size() != static_cast(array->length())) { return Status::Invalid( fmt::format("predicate result size {} does not match probe batch length {}", diff --git a/src/paimon/common/reader/predicate_batch_reader.cpp b/src/paimon/common/reader/predicate_batch_reader.cpp index 621d7568c..fb9a5f974 100644 --- a/src/paimon/common/reader/predicate_batch_reader.cpp +++ b/src/paimon/common/reader/predicate_batch_reader.cpp @@ -61,6 +61,9 @@ Result> PredicateBatchReader::Create( return Status::Invalid( fmt::format("predicate {} does not support Test", predicate->ToString())); } + if (arrow_pool == nullptr) { + return Status::Invalid("create predicate batch reader failed. arrow pool is nullptr"); + } return std::unique_ptr( new PredicateBatchReader(std::move(reader), predicate, arrow_pool)); } @@ -124,7 +127,8 @@ Result PredicateBatchReader::Filter(const std::shared_ptr result, predicate_filter_->Test(*array)); + PAIMON_ASSIGN_OR_RAISE(std::vector result, + predicate_filter_->Test(*array, arrow_pool_.get())); assert(result.size() == static_cast(array->length())); RoaringBitmap32 is_valid; for (int32_t i = 0; i < static_cast(result.size()); i++) {