diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index cd952055666..1721da096f5 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -30,6 +30,7 @@ #include "arrow/compute/exec_internal.h" #include "arrow/compute/expression_internal.h" #include "arrow/compute/function_internal.h" +#include "arrow/compute/kernels/scalar_set_lookup_internal.h" #include "arrow/compute/util.h" #include "arrow/io/memory.h" #ifdef ARROW_IPC @@ -1270,12 +1271,10 @@ struct Inequality { auto options = checked_pointer_cast(is_in_call->options); - // The maximum number of values in the is_in expression set of values - // in order to use the simplification. - // If the set is large there are performance implications, see: - // https://github.com/apache/arrow/issues/46777 - constexpr int16_t kIsInSimplificationMaxValueSet = 50; - if (options->value_set.length() > kIsInSimplificationMaxValueSet) { + // Filtering a large value set and rebinding the call for every guarantee is + // expensive (https://github.com/apache/arrow/issues/46777). Large value sets + // are handled by SimplifyIsInWithInequalities instead. + if (options->value_set.length() > internal::kIsInSimplificationMaxValueSet) { return std::nullopt; } @@ -1412,6 +1411,158 @@ struct Inequality { } }; +/// \brief Whether an `is_in` call can be simplified to `false` once we know that +/// none of its (non-null) value set members can occur in the data. +/// +/// \param nullable whether the data matched against the value set may be null +bool IsInWithoutMatchesIsFalse(const SetLookupOptions& options, bool value_set_has_null, + bool nullable) { + switch (options.null_matching_behavior) { + case SetLookupOptions::MATCH: + // A null in the value set matches nulls in the data. + return !(nullable && value_set_has_null); + case SetLookupOptions::SKIP: + // Nulls in the data produce false. + return true; + case SetLookupOptions::EMIT_NULL: + // Nulls in the data produce null. + return !nullable; + case SetLookupOptions::INCONCLUSIVE: + // Nulls in the data produce null and so does any non-match if the value + // set contains a null. + return !nullable && !value_set_has_null; + } + return false; +} + +/// \brief Simplify an `is_in` call with a large value set against all inequality +/// guarantees on its argument at once. +/// +/// Rather than filtering the value set (see Inequality::SimplifyIsIn), this binary +/// searches the sorted copy of the value set kept in the kernel state for the +/// range of values admitted by the inequalities. If that range is empty the call +/// is simplified to `false`; otherwise it is left unchanged. This is exact for any +/// number of values and costs O(log N) per guarantee, with no rebinding. +/// +/// Applying all inequalities on a field together matters: a row group's +/// statistics yield `x >= min and x <= max`, and neither bound on its own can +/// rule out a value set whose hull spans the row group. +Result SimplifyIsInWithInequalities( + const std::vector& inequalities, Expression expr) { + auto call = expr.call(); + if (!call || call->function_name != "is_in") return expr; + if (!call->kernel_state) return expr; + + auto options = checked_pointer_cast(call->options); + if (options->value_set.length() <= internal::kIsInSimplificationMaxValueSet) { + // Handled by Inequality::SimplifyIsIn + return expr; + } + + const auto& lhs = Comparison::StripOrderPreservingCasts(call->arguments[0]); + if (!lhs.field_ref()) return expr; + const FieldRef& target = *lhs.field_ref(); + // The casts stripped above are nondecreasing but, when the argument is floating + // point, not necessarily injective (e.g. float64 to float32 or int64 to float64), + // so `x > b` only implies `cast(x) >= cast(b)`: strict bounds have to be + // weakened to inclusive ones below. Integer widenings are injective. + const bool lhs_was_cast = &lhs != &call->arguments[0]; + const bool weaken_strict_bounds = + lhs_was_cast && is_floating(call->arguments[0].type()->id()); + + std::vector applicable; + for (const Inequality& inequality : inequalities) { + if (inequality.target != target) continue; + if (inequality.cmp == Comparison::NA || inequality.cmp == Comparison::NOT_EQUAL) { + continue; + } + applicable.push_back(&inequality); + } + // Only materialize the (lazily sorted) value set once we know it can be used. + if (applicable.empty()) return expr; + + const auto& state = + checked_cast(*call->kernel_state); + std::shared_ptr sorted = state.GetSortedValueSet(); + if (!sorted) return expr; + // If the value set could not be cast to the argument type when the call was + // bound, the kernel casts the argument instead and the argument's order (which + // the guarantee refers to) need not be the value set's order. + if (!sorted->type()->Equals(*call->arguments[0].type())) return expr; + + // Indices into `sorted` delimiting the values admitted by the inequalities. + int64_t begin = 0; + int64_t end = sorted->length(); + // The guarantee holds for every row, and an inequality which does not admit + // nulls can only hold for non-null data. + bool nullable = true; + + ExecContext exec_context; + for (const Inequality* inequality : applicable) { + nullable = nullable && inequality->nullable; + + std::shared_ptr bound = inequality->bound.scalar(); + if (!bound->type->Equals(*sorted->type())) { + // The sorted value set has the type of the (cast) argument. A bound which + // cannot be cast to it exactly is ignored, which only loses precision: + // "safe" casts between floating point types or to decimals may still + // round, so require the cast to round trip. + auto maybe_cast = + Cast(inequality->bound, sorted->type(), CastOptions::Safe(), &exec_context); + if (!maybe_cast.ok()) continue; + auto maybe_round_trip = + Cast(*maybe_cast, bound->type, CastOptions::Safe(), &exec_context); + if (!maybe_round_trip.ok() || !maybe_round_trip->scalar()->Equals(*bound)) { + continue; + } + bound = maybe_cast->scalar(); + } + // A comparison against null or NaN cannot be guaranteed true, but be defensive. + if (!bound->is_valid || internal::IsNaNScalar(*bound)) continue; + + Comparison::type cmp = inequality->cmp; + if (weaken_strict_bounds) { + if (cmp == Comparison::GREATER) cmp = Comparison::GREATER_EQUAL; + if (cmp == Comparison::LESS) cmp = Comparison::LESS_EQUAL; + } + + ARROW_ASSIGN_OR_RAISE( + int64_t lower, internal::SearchSortedValueSet(*sorted, *bound, /*upper=*/false)); + ARROW_ASSIGN_OR_RAISE( + int64_t upper, internal::SearchSortedValueSet(*sorted, *bound, /*upper=*/true)); + switch (cmp) { + case Comparison::GREATER_EQUAL: + begin = std::max(begin, lower); + break; + case Comparison::GREATER: + begin = std::max(begin, upper); + break; + case Comparison::LESS_EQUAL: + end = std::min(end, upper); + break; + case Comparison::LESS: + end = std::min(end, lower); + break; + case Comparison::EQUAL: + begin = std::max(begin, lower); + end = std::min(end, upper); + break; + case Comparison::NA: + case Comparison::NOT_EQUAL: + break; + } + } + + // Some value set member may occur in the data: leave the call unchanged rather + // than rebinding it against a narrowed value set. + if (begin < end) return expr; + + if (!IsInWithoutMatchesIsFalse(*options, state.value_set_has_null, nullable)) { + return expr; + } + return literal(false); +} + /// \brief Simplify an expression given a guarantee, if the guarantee /// is is_valid(). Result SimplifyIsValidGuarantee(Expression expr, @@ -1455,6 +1606,28 @@ Result SimplifyWithGuarantee(Expression expr, }; RETURN_NOT_OK(CanonicalizeAndFoldConstants()); + // `is_in` calls with large value sets are simplified against all inequalities + // on their argument at once, see SimplifyIsInWithInequalities. + std::vector inequalities; + for (const auto& guarantee : conjunction_members) { + if (auto inequality = Inequality::ExtractOne(guarantee)) { + inequalities.push_back(*inequality); + } + } + if (!inequalities.empty()) { + Expression before = expr; + ARROW_ASSIGN_OR_RAISE(expr, ModifyExpression( + std::move(expr), [](Expression expr) { return expr; }, + [&](Expression expr, ...) -> Result { + return SimplifyIsInWithInequalities( + inequalities, std::move(expr)); + })); + + if (!Expression::Identical(expr, before)) { + RETURN_NOT_OK(CanonicalizeAndFoldConstants()); + } + } + for (const auto& guarantee : conjunction_members) { if (!guarantee.call()) continue; diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 10c1b348e62..94ca1dd627e 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -90,6 +90,14 @@ std::string make_range_json(int start, int end) { return result; } +// Concatenate JSON arrays: "[1,2]" + "[3]" -> "[1,2,3]" +std::string concat_json(const std::string& left, const std::string& right) { + std::string result = left.substr(0, left.size() - 1); + if (result.size() > 1 && right.size() > 2) result += ","; + result += right.substr(1); + return result; +} + const auto no_change = std::nullopt; TEST(ExpressionUtils, Comparison) { @@ -2063,8 +2071,8 @@ TEST(Expression, SimplifyIsIn) { .WithGuarantee(greater(field_ref("u32"), literal(10))) .Expect(is_in(field_ref("u32"), int64(), make_range_json(11, 40), null_matching)); - // For large ranges we don't do any simplification, see - // `kIsInSimplificationMaxValueSet` in expression.cc. + // Large value sets (longer than `kIsInSimplificationMaxValueSet`) are never + // narrowed, only simplified to false. See SimplifyIsInLargeValueSet. Simplify{is_in(field_ref("u32"), int64(), make_range_json(1, 100), null_matching)} .WithGuarantee(greater(field_ref("u32"), literal(3))) .ExpectUnchanged(); @@ -2151,6 +2159,332 @@ TEST(Expression, SimplifyIsIn) { .ExpectUnchanged(); } +TEST(Expression, SimplifyIsInLargeValueSet) { + // Value sets longer than kIsInSimplificationMaxValueSet are binary searched for + // the range admitted by the inequalities of the guarantee. The call is simplified + // to false if the range is empty and is left unchanged otherwise. + auto is_in = [](Expression field, std::shared_ptr value_set_type, + std::string json_array, + SetLookupOptions::NullMatchingBehavior null_matching_behavior) { + SetLookupOptions options{ArrayFromJSON(value_set_type, json_array), + null_matching_behavior}; + return call("is_in", {field}, options); + }; + + auto u32 = field_ref("u32"); + auto i32 = field_ref("i32"); + auto i64 = field_ref("i64"); + auto f64 = field_ref("f64"); + auto str = field_ref("str"); + + // 100 contiguous values + auto range = make_range_json(1, 100); + // 101 values with a gap: [1, 50] and [100, 150] + auto gapped = concat_json(make_range_json(1, 50), make_range_json(100, 150)); + + for (SetLookupOptions::NullMatchingBehavior null_matching : { + SetLookupOptions::MATCH, + SetLookupOptions::SKIP, + SetLookupOptions::EMIT_NULL, + SetLookupOptions::INCONCLUSIVE, + }) { + // Guarantees overlapping the value set leave the call unchanged + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(greater(u32, literal(3))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(less_equal(u32, literal(1))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(100)), less_equal(u32, literal(200)))) + .ExpectUnchanged(); + + // Guarantees disjoint from the value set's hull + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(greater(u32, literal(100))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(greater_equal(u32, literal(101))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(less(u32, literal(1))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(less_equal(u32, literal(0))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(101)), less_equal(u32, literal(200)))) + .Expect(false); + + // Guarantees inside the value set's hull but disjoint from its values. This + // requires considering both bounds of the guarantee at once. + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(and_(greater(u32, literal(50)), less(u32, literal(100)))) + .Expect(false); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(51)), less_equal(u32, literal(99)))) + .Expect(false); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(less_equal(u32, literal(99)), greater_equal(u32, literal(51)))) + .Expect(false); + // ... but a single value of the set in the range is enough to keep the call + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(50)), less_equal(u32, literal(99)))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(51)), less_equal(u32, literal(100)))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(greater(u32, literal(50))) + .ExpectUnchanged(); + // Redundant bounds on the same field are all applied + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(and_(and_(greater(u32, literal(10)), greater(u32, literal(50))), + less(u32, literal(100)))) + .Expect(false); + + // Guarantees on other fields are irrelevant + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) + .ExpectUnchanged(); + + // Casts: the value set is cast to the field type when the call is bound and + // the bounds of the guarantee are cast to the value set type. + Simplify{is_in(i64, int32(), range, null_matching)} + .WithGuarantee(greater(i64, literal(100))) + .Expect(false); + Simplify{is_in(i64, int32(), range, null_matching)} + .WithGuarantee(greater(i64, literal(100))) + .Expect(false); + Simplify{is_in(cast(i32, int64()), int64(), range, null_matching)} + .WithGuarantee(greater(i32, literal(100))) + .Expect(false); + + // Unsorted value sets with duplicates are sorted when the call is bound + std::string unsorted = "["; + for (int i = 100; i >= 1; --i) { + unsorted += std::to_string(i) + "," + std::to_string(i) + ","; + } + unsorted.back() = ']'; + Simplify{is_in(u32, int64(), unsorted, null_matching)} + .WithGuarantee(greater(u32, literal(100))) + .Expect(false); + Simplify{is_in(u32, int64(), unsorted, null_matching)} + .WithGuarantee(and_(greater(u32, literal(50)), less(u32, literal(51)))) + .Expect(false); + Simplify{is_in(u32, int64(), unsorted, null_matching)} + .WithGuarantee(and_(greater_equal(u32, literal(50)), less(u32, literal(51)))) + .ExpectUnchanged(); + + // Strings + std::string strings = "["; + for (int i = 0; i < 60; ++i) { + if (i > 0) strings += ","; + strings += "\"k" + std::to_string(100 + i) + "\""; + } + strings += "]"; + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(greater(str, literal("k159"))) + .Expect(false); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(less(str, literal("k100"))) + .Expect(false); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(and_(greater(str, literal("k110")), less(str, literal("k111")))) + .Expect(false); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee( + and_(greater_equal(str, literal("k110")), less(str, literal("k111")))) + .ExpectUnchanged(); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(greater(str, literal("k1"))) + .ExpectUnchanged(); + // Unsorted strings with duplicates take the generic sorting path + std::string unsorted_strings = "["; + for (int i = 59; i >= 0; --i) { + unsorted_strings += "\"k" + std::to_string(100 + i) + "\","; + unsorted_strings += "\"k" + std::to_string(100 + i) + "\","; + } + unsorted_strings.back() = ']'; + Simplify{is_in(str, utf8(), unsorted_strings, null_matching)} + .WithGuarantee(and_(greater(str, literal("k110")), less(str, literal("k111")))) + .Expect(false); + Simplify{is_in(str, utf8(), unsorted_strings, null_matching)} + .WithGuarantee( + and_(greater(str, literal("k110")), less_equal(str, literal("k111")))) + .ExpectUnchanged(); + // ... as do value sets with nulls, which INCONCLUSIVE cannot simplify + auto with_null = Simplify{is_in(str, utf8(), concat_json(unsorted_strings, "[null]"), + null_matching)} + .WithGuarantee(and_(greater(str, literal("k110")), + less(str, literal("k111")))); + if (null_matching == SetLookupOptions::INCONCLUSIVE) { + with_null.ExpectUnchanged(); + } else { + with_null.Expect(false); + } + + // Floats + Simplify{is_in(f64, float64(), range, null_matching)} + .WithGuarantee(greater(f64, literal(100.5))) + .Expect(false); + Simplify{is_in(f64, float64(), range, null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.0)), less(f64, literal(51.0)))) + .Expect(false); + Simplify{is_in(f64, float64(), range, null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.0)), less_equal(f64, literal(51.0)))) + .ExpectUnchanged(); + // A double bound which does not round trip through the float32 value set type + // is ignored rather than rounded: 0.1 rounds up to 0.1f, which would wrongly + // exclude the member 0.1f (0.1f > 0.1 holds in double). + auto f32 = field_ref("f32"); + Simplify{is_in(f32, float32(), concat_json("[0.1]", range), null_matching)} + .WithGuarantee(and_(greater(f32, literal(0.1)), less(f32, literal(0.5)))) + .ExpectUnchanged(); + Simplify{is_in(f32, float32(), concat_json("[0.7]", range), null_matching)} + .WithGuarantee(less(f32, literal(0.7))) + .ExpectUnchanged(); + Simplify{is_in(f32, float32(), concat_json("[0.5]", range), null_matching)} + .WithGuarantee(less(f32, literal(0.5))) + .Expect(false); + Simplify{is_in(f32, float32(), concat_json("[0.5]", range), null_matching)} + .WithGuarantee(and_(greater(f32, literal(0.25)), less(f32, literal(0.5)))) + .Expect(false); + // Statistics cannot bound NaN, so a value set containing NaN is never simplified + Simplify{is_in(f64, float64(), concat_json("[NaN]", range), null_matching)} + .WithGuarantee(greater(f64, literal(100.5))) + .ExpectUnchanged(); + Simplify{is_in(f64, float64(), concat_json(range, "[NaN]"), null_matching)} + .WithGuarantee(greater(f64, literal(100.5))) + .ExpectUnchanged(); + + // Strict bounds through a stripped narrowing cast are weakened to inclusive + // ones: f64 = 51.00000001 satisfies f64 > 51 and casts to 51f, a set member. + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(greater(f64, literal(51.0))) + .ExpectUnchanged(); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(less(f64, literal(1.0))) + .ExpectUnchanged(); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(greater(f64, literal(51.5))) + .Expect(false); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.0)), less(f64, literal(51.0)))) + .ExpectUnchanged(); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.25)), less(f64, literal(50.75)))) + .Expect(false); + + // When the value set cannot be cast to the argument type the kernel casts the + // argument instead, in which case the guarantee's order (here lexicographic) + // says nothing about the value set's order: "10" <= x <= "2" admits "19". + auto binary_literal = [](const std::string& value) { + return literal(std::make_shared(Buffer::FromString(value))); + }; + Simplify{is_in(field_ref("binary"), int64(), make_range_json(10, 60), null_matching)} + .WithGuarantee(and_(greater_equal(field_ref("binary"), binary_literal("10")), + less_equal(field_ref("binary"), binary_literal("2")))) + .ExpectUnchanged(); + + // Dictionaries are not searchable + Simplify{is_in(field_ref("dict_i32"), int32(), range, null_matching)} + .WithGuarantee(greater(field_ref("dict_i32"), literal(100))) + .ExpectUnchanged(); + } + + // Null matching behavior. The field may be null if every inequality of the + // guarantee admits nulls. + auto nullable_guarantee = or_(greater(i32, literal(100)), is_null(i32)); + auto non_nullable_guarantee = greater(i32, literal(100)); + auto mixed_guarantee = + and_(or_(greater(i32, literal(100)), is_null(i32)), less_equal(i32, literal(200))); + auto range_with_null = concat_json(range, "[null]"); + + // MATCH: nulls match a null in the value set + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::MATCH)} + .WithGuarantee(nullable_guarantee) + .ExpectUnchanged(); + Simplify{is_in(i32, int32(), range, SetLookupOptions::MATCH)} + .WithGuarantee(nullable_guarantee) + .Expect(false); + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::MATCH)} + .WithGuarantee(non_nullable_guarantee) + .Expect(false); + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::MATCH)} + .WithGuarantee(mixed_guarantee) + .Expect(false); + + // SKIP: nulls produce false + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::SKIP)} + .WithGuarantee(nullable_guarantee) + .Expect(false); + Simplify{is_in(i32, int32(), range, SetLookupOptions::SKIP)} + .WithGuarantee(nullable_guarantee) + .Expect(false); + + // EMIT_NULL: nulls produce null + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::EMIT_NULL)} + .WithGuarantee(nullable_guarantee) + .ExpectUnchanged(); + Simplify{is_in(i32, int32(), range, SetLookupOptions::EMIT_NULL)} + .WithGuarantee(nullable_guarantee) + .ExpectUnchanged(); + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::EMIT_NULL)} + .WithGuarantee(non_nullable_guarantee) + .Expect(false); + Simplify{is_in(i32, int32(), range, SetLookupOptions::EMIT_NULL)} + .WithGuarantee(mixed_guarantee) + .Expect(false); + + // INCONCLUSIVE: nulls produce null, as does any non-match if the value set + // contains a null + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::INCONCLUSIVE)} + .WithGuarantee(nullable_guarantee) + .ExpectUnchanged(); + Simplify{is_in(i32, int32(), range, SetLookupOptions::INCONCLUSIVE)} + .WithGuarantee(nullable_guarantee) + .ExpectUnchanged(); + Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::INCONCLUSIVE)} + .WithGuarantee(non_nullable_guarantee) + .ExpectUnchanged(); + Simplify{is_in(i32, int32(), range, SetLookupOptions::INCONCLUSIVE)} + .WithGuarantee(non_nullable_guarantee) + .Expect(false); + Simplify{is_in(i32, int32(), range, SetLookupOptions::INCONCLUSIVE)} + .WithGuarantee(mixed_guarantee) + .Expect(false); + + // Equality guarantees which admit nulls are not known values but inequalities + Simplify{is_in(i32, int32(), gapped, SetLookupOptions::MATCH)} + .WithGuarantee(or_(equal(i32, literal(75)), is_null(i32))) + .Expect(false); + Simplify{is_in(i32, int32(), gapped, SetLookupOptions::MATCH)} + .WithGuarantee(or_(equal(i32, literal(50)), is_null(i32))) + .ExpectUnchanged(); + + // The large value set is simplified in combination with other simplifications + Simplify{and_(is_in(i32, int32(), gapped, SetLookupOptions::MATCH), + greater(field_ref("u32"), literal(5)))} + .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) + .Expect(false); + Simplify{or_(is_in(i32, int32(), gapped, SetLookupOptions::MATCH), + greater(field_ref("u32"), literal(5)))} + .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) + .Expect(greater(field_ref("u32"), literal(5))); +} + TEST(Expression, SimplifyThenExecute) { auto filter = or_({equal(field_ref("f32"), literal(0)), @@ -2178,6 +2512,46 @@ TEST(Expression, SimplifyThenExecute) { AssertDatumsEqual(evaluated, simplified_evaluated, /*verbose=*/true); } +TEST(Expression, SimplifyLargeIsInThenExecute) { + // Rows satisfying the guarantee must evaluate identically under the original + // and the simplified filter, in particular through a narrowing cast. + auto value_set = ArrayFromJSON(float32(), make_range_json(1, 51)); + auto filter = call("is_in", {cast(field_ref("f64"), float32())}, + compute::SetLookupOptions{value_set}); + ASSERT_OK_AND_ASSIGN(filter, filter.Bind(*kBoringSchema)); + + struct Case { + Expression guarantee; + std::string input; // rows satisfying the guarantee + }; + for ( + const Case& c : std::vector{ + {greater(field_ref("f64"), literal(51.0)), + R"([{"f64": 51.00000001}, {"f64": 51.0000001}, {"f64": 51.5}, {"f64": 1e9}])"}, + {greater_equal(field_ref("f64"), literal(51.0)), + R"([{"f64": 51.0}, {"f64": 51.00000001}, {"f64": 52.0}])"}, + {and_(greater(field_ref("f64"), literal(51.0)), + less(field_ref("f64"), literal(52.0))), + R"([{"f64": 51.00000001}, {"f64": 51.5}, {"f64": 51.9}])"}, + {less(field_ref("f64"), literal(1.0)), + R"([{"f64": 0.99999999}, {"f64": 0.5}, {"f64": -1e9}])"}, + {greater(field_ref("f64"), literal(51.5)), + R"([{"f64": 51.6}, {"f64": 52.0}, {"f64": 1e9}])"}, + }) { + auto input = RecordBatchFromJSON(kBoringSchema, c.input); + ASSERT_OK_AND_ASSIGN(auto simplified, SimplifyWithGuarantee(filter, c.guarantee)); + Datum evaluated, simplified_evaluated; + ExpectExecute(filter, input, &evaluated); + ExpectExecute(simplified, input, &simplified_evaluated); + if (simplified_evaluated.is_scalar()) { + ASSERT_OK_AND_ASSIGN( + simplified_evaluated, + MakeArrayFromScalar(*simplified_evaluated.scalar(), input->num_rows())); + } + AssertDatumsEqual(evaluated, simplified_evaluated, /*verbose=*/true); + } +} + TEST(Expression, SimplifyIsInThenExecute) { auto input = RecordBatchFromJSON(kBoringSchema, R"([ {"i64": 2, "i32": 5}, diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc index c6a3562233c..676661b0287 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc @@ -16,9 +16,11 @@ // under the License. #include "arrow/array/array_base.h" +#include "arrow/array/concatenate.h" #include "arrow/compute/api_scalar.h" #include "arrow/compute/cast.h" #include "arrow/compute/kernels/common_internal.h" +#include "arrow/compute/kernels/scalar_set_lookup_internal.h" #include "arrow/compute/kernels/util_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/type.h" @@ -36,11 +38,6 @@ using internal::HashTraits; namespace compute::internal { namespace { -// This base class enables non-templated access to the value set type -struct SetLookupStateBase : public KernelState { - std::shared_ptr value_set_type; -}; - template struct SetLookupState : public SetLookupStateBase { explicit SetLookupState(MemoryPool* pool) : memory_pool(pool) {} @@ -69,6 +66,7 @@ struct SetLookupState : public SetLookupStateBase { } else { return Status::Invalid("value_set should be an array or chunked array"); } + value_set_has_null = lookup_table->GetNull() >= 0; if (this->null_matching_behavior != SetLookupOptions::SKIP && lookup_table->GetNull() >= 0) { null_index = memo_index_to_value_index[lookup_table->GetNull()]; @@ -130,13 +128,12 @@ struct SetLookupState : public SetLookupStateBase { Status Init(SetLookupOptions& options) { null_matching_behavior = options.GetNullMatchingBehavior(); - value_set_has_null = (options.value_set.length() > 0) && - this->null_matching_behavior != SetLookupOptions::SKIP; + // The kernels consult null_matching_behavior before using this. + value_set_has_null = options.value_set.length() > 0; value_set_type = null(); return Status::OK(); } - bool value_set_has_null; SetLookupOptions::NullMatchingBehavior null_matching_behavior; }; @@ -171,10 +168,15 @@ struct InitStateVisitor { TypeHolder arg_type; std::unique_ptr result; - InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args) + // Whether to populate SetLookupStateBase::sorted_value_set (is_in only). + bool compute_sorted_value_set; + + InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args, + bool compute_sorted_value_set) : ctx(ctx), options(*checked_cast(args.options)), - arg_type(args.inputs[0]) {} + arg_type(args.inputs[0]), + compute_sorted_value_set(compute_sorted_value_set) {} template Status Init() { @@ -254,18 +256,49 @@ struct InitStateVisitor { } RETURN_NOT_OK(VisitTypeInline(*options.value_set.type(), this)); + + if (compute_sorted_value_set && + options.value_set.length() > kIsInSimplificationMaxValueSet && + IsSearchableValueSetType(*options.value_set.type())) { + // Best effort: a missing sorted copy only disables the simplification of + // `is_in` against guarantees, see SetLookupStateBase::GetSortedValueSet. + auto maybe_values = MaterializedValueSet(); + if (maybe_values.ok()) { + ARROW_UNUSED(checked_cast(result.get()) + ->SetSearchableValueSet(*std::move(maybe_values))); + } + } return std::move(result); } + + Result> MaterializedValueSet() { + if (options.value_set.is_array()) { + return options.value_set.make_array(); + } + return Concatenate(options.value_set.chunked_array()->chunks(), + ctx->exec_context()->memory_pool()); + } }; Result> InitSetLookup(KernelContext* ctx, - const KernelInitArgs& args) { + const KernelInitArgs& args, + bool compute_sorted_value_set) { if (args.options == nullptr) { return Status::Invalid( "Attempted to call a set lookup function without SetLookupOptions"); } - return InitStateVisitor{ctx, args}.GetResult(); + return InitStateVisitor{ctx, args, compute_sorted_value_set}.GetResult(); +} + +Result> InitIsIn(KernelContext* ctx, + const KernelInitArgs& args) { + return InitSetLookup(ctx, args, /*compute_sorted_value_set=*/true); +} + +Result> InitIndexIn(KernelContext* ctx, + const KernelInitArgs& args) { + return InitSetLookup(ctx, args, /*compute_sorted_value_set=*/false); } struct IndexInVisitor { @@ -648,7 +681,7 @@ void RegisterScalarSetLookup(FunctionRegistry* registry) { // IsIn writes its boolean output into preallocated memory { ScalarKernel isin_base; - isin_base.init = InitSetLookup; + isin_base.init = InitIsIn; isin_base.exec = ExecIsIn; isin_base.null_handling = NullHandling::COMPUTED_PREALLOCATE; auto is_in = std::make_shared("is_in", Arity::Unary(), is_in_doc); @@ -665,7 +698,7 @@ void RegisterScalarSetLookup(FunctionRegistry* registry) { // IndexIn writes its int32 output into preallocated memory { ScalarKernel index_in_base; - index_in_base.init = InitSetLookup; + index_in_base.init = InitIndexIn; index_in_base.exec = ExecIndexIn; index_in_base.null_handling = NullHandling::COMPUTED_PREALLOCATE; auto index_in = diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h new file mode 100644 index 00000000000..60011e4a676 --- /dev/null +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -0,0 +1,386 @@ +// 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 + +// Shared between the set lookup kernels (libarrow_compute) and expression +// simplification (libarrow). Everything here must stay header-only. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_binary.h" +#include "arrow/array/array_decimal.h" +#include "arrow/array/array_primitive.h" +#include "arrow/array/concatenate.h" +#include "arrow/array/data.h" +#include "arrow/array/util.h" +#include "arrow/buffer.h" +#include "arrow/compute/api_vector.h" +#include "arrow/compute/kernel.h" +#include "arrow/datum.h" +#include "arrow/result.h" +#include "arrow/scalar.h" +#include "arrow/status.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/decimal.h" +#include "arrow/util/logging.h" +#include "arrow/visit_type_inline.h" + +namespace arrow::compute::internal { + +/// Value sets up to this length are simplified against a guarantee by filtering +/// the value set and rebinding the `is_in` call (see SimplifyIsIn in expression.cc). +/// Longer value sets are instead binary searched through SetLookupStateBase:: +/// sorted_value_set, which only ever simplifies the call to `false`, see +/// https://github.com/apache/arrow/issues/46777 +constexpr int64_t kIsInSimplificationMaxValueSet = 50; + +/// Whether a (valid) scalar of a floating point type is NaN. +inline bool IsNaNScalar(const Scalar& value) { + if (value.type->id() == Type::FLOAT) { + return std::isnan(::arrow::internal::checked_cast(value).value); + } + if (value.type->id() == Type::DOUBLE) { + return std::isnan(::arrow::internal::checked_cast(value).value); + } + return false; +} + +inline Result IsSortedValueSet(const Array& values); +inline Result> SortPrimitiveValueSet(const Array& values); + +/// Common state for the is_in and index_in kernels. +struct SetLookupStateBase : public KernelState { + std::shared_ptr value_set_type; + + /// Whether the value set contains a null (regardless of null matching behavior). + bool value_set_has_null = false; + + /// Make the value set available to GetSortedValueSet(). + /// + /// Called once by the `is_in` kernel init, before the state is shared. If the + /// value set is sorted already it is used as is; otherwise sorting is deferred + /// to the first call of GetSortedValueSet() so that only expressions which are + /// simplified against guarantees pay for it. + Status SetSearchableValueSet(std::shared_ptr values) { + if (values->null_count() == 0) { + ARROW_ASSIGN_OR_RAISE(bool sorted, IsSortedValueSet(*values)); + if (sorted) { + sorted_value_set_ = std::move(values); + sorted_value_set_ready_.store(true, std::memory_order_release); + return Status::OK(); + } + } + unsorted_value_set_ = std::move(values); + return Status::OK(); + } + + /// A copy of the value set sorted ascending in the order of the comparison + /// kernels, with nulls removed, or nullptr if there is none. + /// + /// This is only available for `is_in` when the value set is longer than + /// kIsInSimplificationMaxValueSet, has a type for which IsSearchableValueSetType() + /// holds and contains no NaN. It is computed at most once; concurrent calls are + /// safe. + std::shared_ptr GetSortedValueSet() const { + if (!sorted_value_set_ready_.load(std::memory_order_acquire)) { + std::lock_guard lock(sorted_value_set_mutex_); + if (!sorted_value_set_ready_.load(std::memory_order_relaxed)) { + if (unsorted_value_set_) { + // Best effort: a failure only disables the simplification of `is_in` + // against guarantees. + sorted_value_set_ = SortValueSet(*unsorted_value_set_).ValueOr(nullptr); + unsorted_value_set_.reset(); + } + sorted_value_set_ready_.store(true, std::memory_order_release); + } + } + return sorted_value_set_; + } + + private: + /// \return `values` sorted with nulls removed, or nullptr if it contains NaN + static Result> SortValueSet(const Array& values) { + // Sorting a copy of primitive values directly is several times faster than + // the generic sort_indices + take path. + auto maybe_sorted = SortPrimitiveValueSet(values); + if (maybe_sorted.ok()) return maybe_sorted; + if (!maybe_sorted.status().IsNotImplemented()) return maybe_sorted.status(); + + ARROW_ASSIGN_OR_RAISE(auto indices, + SortIndices(values, ArraySortOptions(SortOrder::Ascending, + NullPlacement::AtEnd))); + ARROW_ASSIGN_OR_RAISE(Datum taken, + Take(values, indices, TakeOptions::NoBoundsCheck())); + auto sorted = taken.make_array(); + sorted = sorted->Slice(0, sorted->length() - sorted->null_count()); + // NaN sorts after every other value and before nulls. Statistics based + // guarantees cannot bound NaN, so such a value set is not searchable. + if (sorted->length() > 0 && is_floating(sorted->type_id())) { + ARROW_ASSIGN_OR_RAISE(auto last, sorted->GetScalar(sorted->length() - 1)); + if (IsNaNScalar(*last)) return std::shared_ptr(nullptr); + } + return sorted; + } + + mutable std::mutex sorted_value_set_mutex_; + mutable std::atomic sorted_value_set_ready_{false}; + mutable std::shared_ptr unsorted_value_set_; + mutable std::shared_ptr sorted_value_set_; +}; + +/// Whether value sets of this type can be binary searched with a comparator that +/// agrees with both the sorting kernels and the comparison kernels. +inline bool IsSearchableValueSetType(const DataType& type) { + switch (type.id()) { + case Type::INT8: + case Type::INT16: + case Type::INT32: + case Type::INT64: + case Type::UINT8: + case Type::UINT16: + case Type::UINT32: + case Type::UINT64: + case Type::FLOAT: + case Type::DOUBLE: + case Type::DATE32: + case Type::DATE64: + case Type::TIME32: + case Type::TIME64: + case Type::TIMESTAMP: + case Type::DURATION: + case Type::DECIMAL32: + case Type::DECIMAL64: + case Type::DECIMAL128: + case Type::DECIMAL256: + case Type::STRING: + case Type::LARGE_STRING: + case Type::BINARY: + case Type::LARGE_BINARY: + case Type::FIXED_SIZE_BINARY: + return true; + default: + return false; + } +} + +namespace set_lookup_detail { + +template +struct ValueSetAccessor; + +template +struct ValueSetAccessor< + T, enable_if_t::value && !is_boolean_type::value && + !is_half_float_type::value && !is_interval_type::value>> { + using ArrayType = typename TypeTraits::ArrayType; + using ScalarType = typename TypeTraits::ScalarType; + using ValueType = typename T::c_type; + + static ValueType Get(const Array& values, int64_t i) { + return ::arrow::internal::checked_cast(values).Value(i); + } + static ValueType FromScalar(const Scalar& scalar) { + return ::arrow::internal::checked_cast(scalar).value; + } +}; + +template +struct ValueSetAccessor> { + using ArrayType = typename TypeTraits::ArrayType; + using ScalarType = typename TypeTraits::ScalarType; + using ValueType = typename ScalarType::ValueType; + + static ValueType Get(const Array& values, int64_t i) { + return ValueType( + ::arrow::internal::checked_cast(values).GetValue(i)); + } + static ValueType FromScalar(const Scalar& scalar) { + return ::arrow::internal::checked_cast(scalar).value; + } +}; + +template +struct ValueSetAccessor::value || + std::is_same::value>> { + using ArrayType = typename TypeTraits::ArrayType; + using ValueType = std::string_view; + + static ValueType Get(const Array& values, int64_t i) { + return ::arrow::internal::checked_cast(values).GetView(i); + } + static ValueType FromScalar(const Scalar& scalar) { + return ::arrow::internal::checked_cast(scalar).view(); + } +}; + +template +struct is_searchable_value_set_type : std::false_type {}; + +template +struct is_searchable_value_set_type< + T, std::void_t::Get(std::declval(), 0))>> + : std::true_type {}; + +template +bool IsNaNValue(const ValueType&) { + return false; +} +inline bool IsNaNValue(float value) { return std::isnan(value); } +inline bool IsNaNValue(double value) { return std::isnan(value); } + +struct IsSortedVisitor { + const Array& values; + bool result = false; + + template + enable_if_t::value, Status> Visit(const T&) { + using Accessor = ValueSetAccessor; + if (values.length() == 0) { + result = true; + return Status::OK(); + } + auto previous = Accessor::Get(values, 0); + if (IsNaNValue(previous)) return Status::OK(); + for (int64_t i = 1; i < values.length(); ++i) { + auto current = Accessor::Get(values, i); + if (IsNaNValue(current) || current < previous) return Status::OK(); + previous = current; + } + result = true; + return Status::OK(); + } + + Status Visit(const DataType& type) { + return Status::NotImplemented("Searching value sets of type ", type); + } +}; + +struct SortPrimitiveVisitor { + const Array& values; + std::shared_ptr result; + + template + enable_if_t::value && is_searchable_value_set_type::value, Status> + Visit(const T&) { + using c_type = typename T::c_type; + const auto* data = values.data()->template GetValues(1); + std::vector sorted; + sorted.reserve(values.length() - values.null_count()); + if (values.null_count() == 0) { + sorted.assign(data, data + values.length()); + } else { + for (int64_t i = 0; i < values.length(); ++i) { + if (values.IsValid(i)) sorted.push_back(data[i]); + } + } + // NaN would break the strict weak ordering required by std::sort. Statistics + // based guarantees cannot bound NaN anyway, so such a value set is not + // searchable (nullptr). + for (const auto& value : sorted) { + if (IsNaNValue(value)) return Status::OK(); + } + std::sort(sorted.begin(), sorted.end()); + const auto length = static_cast(sorted.size()); + result = MakeArray(ArrayData::Make(values.type(), length, + {nullptr, Buffer::FromVector(std::move(sorted))}, + /*null_count=*/0)); + return Status::OK(); + } + + Status Visit(const DataType& type) { + return Status::NotImplemented("Sorting value sets of type ", type); + } +}; + +struct SearchSortedVisitor { + const Array& sorted; + const Scalar& bound; + bool upper; + int64_t result = 0; + + template + enable_if_t::value, Status> Visit(const T&) { + using Accessor = ValueSetAccessor; + const auto target = Accessor::FromScalar(bound); + int64_t lo = 0, hi = sorted.length(); + while (lo < hi) { + const int64_t mid = lo + (hi - lo) / 2; + const auto value = Accessor::Get(sorted, mid); + const bool go_right = upper ? !(target < value) : (value < target); + if (go_right) { + lo = mid + 1; + } else { + hi = mid; + } + } + result = lo; + return Status::OK(); + } + + Status Visit(const DataType& type) { + return Status::NotImplemented("Searching value sets of type ", type); + } +}; + +} // namespace set_lookup_detail + +/// Whether `values` (which must contain no nulls) is sorted ascending in the order +/// of the comparison kernels and contains no NaN. +inline Result IsSortedValueSet(const Array& values) { + ARROW_DCHECK_EQ(values.null_count(), 0); + set_lookup_detail::IsSortedVisitor visitor{values, /*result=*/false}; + RETURN_NOT_OK(VisitTypeInline(*values.type(), &visitor)); + return visitor.result; +} + +/// Sort a value set of a primitive type, dropping nulls. +/// +/// \return the sorted value set, or nullptr if it contains NaN +/// \return NotImplemented for non-primitive types +inline Result> SortPrimitiveValueSet(const Array& values) { + set_lookup_detail::SortPrimitiveVisitor visitor{values, /*result=*/nullptr}; + RETURN_NOT_OK(VisitTypeInline(*values.type(), &visitor)); + return visitor.result; +} + +/// Binary search a sorted value set for `bound`. +/// +/// \pre `sorted` satisfies IsSortedValueSet() +/// \pre `bound` is valid, not NaN, and has the same type as `sorted` +/// \return the index of the first element not less than `bound` (`upper == false`), +/// or of the first element greater than `bound` (`upper == true`) +inline Result SearchSortedValueSet(const Array& sorted, const Scalar& bound, + bool upper) { + ARROW_DCHECK(bound.is_valid); + ARROW_DCHECK(sorted.type()->Equals(*bound.type)); + set_lookup_detail::SearchSortedVisitor visitor{sorted, bound, upper, /*result=*/0}; + RETURN_NOT_OK(VisitTypeInline(*sorted.type(), &visitor)); + return visitor.result; +} + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 0d86f5eaa0d..26bf8561f98 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -34,6 +35,7 @@ #include "arrow/io/util_internal.h" #include "arrow/record_batch.h" #include "arrow/table.h" +#include "arrow/testing/builder.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/util.h" #include "arrow/type.h" @@ -747,6 +749,82 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragments) { equal(field_ref(FieldRef("struct", "str")), literal("2"))); } +TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsIn) { + // Row group i holds i64 == 10 * i, ..., 10 * i + 9, so that its statistics + // yield a range guarantee rather than a single known value. + constexpr int kNumRowGroups = 8; + auto iota64 = [](int64_t start, int64_t end) { + std::vector values(end - start); + std::iota(values.begin(), values.end(), start); + return values; + }; + auto dataset_schema = schema({field("i64", int64())}); + std::vector> batches; + for (int i = 0; i < kNumRowGroups; ++i) { + std::shared_ptr values; + ArrayFromVector(iota64(10 * i, 10 * i + 10), &values); + batches.push_back(RecordBatch::Make(dataset_schema, values->length(), {values})); + } + ASSERT_OK_AND_ASSIGN(auto reader, + RecordBatchReader::Make(std::move(batches), dataset_schema)); + auto source = GetFileSource(reader.get()); + + SetSchema(dataset_schema->fields()); + ASSERT_OK_AND_ASSIGN(auto fragment, format_->MakeFragment(*source)); + auto parquet_fragment = checked_pointer_cast(fragment); + + auto is_in = [](std::shared_ptr set) { + return call("is_in", {field_ref("i64")}, compute::SetLookupOptions{std::move(set)}); + }; + auto as_int64_array = [](const std::vector& values) { + std::shared_ptr array; + ArrayFromVector(values, &array); + return array; + }; + auto expect_row_groups = [&](compute::Expression filter, + const std::vector& expected_row_groups) { + SetFilter(std::move(filter)); + ASSERT_OK_AND_ASSIGN(auto fragments, + parquet_fragment->SplitByRowGroup(opts_->filter)); + std::vector row_groups; + for (const auto& split : fragments) { + for (int row_group : + checked_pointer_cast(split)->row_groups()) { + row_groups.push_back(row_group); + } + } + EXPECT_EQ(row_groups, expected_row_groups); + }; + + // A value set longer than the exact simplification threshold whose hull + // [0, 79] spans every row group but which contains no value in [20, 39]. + std::vector values = iota64(0, 20); + for (int64_t v : iota64(40, 80)) values.push_back(v); + ASSERT_GT(values.size(), 50); + + // Only the row groups holding a value of the set are read. + expect_row_groups(is_in(as_int64_array(values)), {0, 1, 4, 5, 6, 7}); + + // Unsorted value sets with duplicates behave the same. + std::vector shuffled(values.rbegin(), values.rend()); + shuffled.insert(shuffled.end(), values.begin(), values.end()); + expect_row_groups(is_in(as_int64_array(shuffled)), {0, 1, 4, 5, 6, 7}); + + // Values beyond the data select nothing. + expect_row_groups(is_in(as_int64_array(iota64(1000, 1100))), {}); + + // Values cast from a narrower type are handled too. + std::shared_ptr int32_set; + ArrayFromVector(std::vector(values.begin(), values.end()), + &int32_set); + expect_row_groups(is_in(int32_set), {0, 1, 4, 5, 6, 7}); + + // Combined with other predicates. + expect_row_groups( + and_(is_in(as_int64_array(values)), less(field_ref("i64"), literal(50))), + {0, 1, 4}); +} + TEST_P(TestParquetFileFormatScan, ExplicitRowGroupSelection) { constexpr int64_t kNumRowGroups = 16; constexpr int64_t kTotalNumRows = kNumRowGroups * (kNumRowGroups + 1) / 2;