From 3b86b13ad5085c2e2ad1f413572191eff7a34599 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 27 Oct 2026 20:26:21 +0100 Subject: [PATCH 1/4] GH-55292: [C++][Dataset] Prune row groups for large is_in value sets using a sorted value set Since GH-46777 is_in predicates with more than 50 values are not simplified against row group statistics at all, because filtering the value set and rebinding the call for every guarantee was prohibitively expensive. Large value sets are now simplified with a binary search instead: - The is_in kernel init keeps a copy of the value set sorted in the order of the comparison kernels, with nulls removed. An already sorted value set is reused as is after a single O(N) check; an unsorted one is sorted lazily and thread-safely on the first simplification, so eager kernel calls never pay for sorting. Value sets containing NaN are not searchable, since statistics cannot bound NaN. - SimplifyWithGuarantee applies all inequality guarantees on a field to a large is_in call at once, narrowing an index range into the sorted value set with lower/upper bound searches. If the range is empty the call simplifies to false (subject to the null matching behavior), otherwise it is left unchanged. This is exact for any number of values, costs O(log N) per guarantee and never rebinds. Applying both bounds together is what allows pruning a row group inside the hull of the value set but holding none of its values, which no per-bound check can do. Value sets up to 50 values keep using the existing filter-and-rebind path. --- cpp/src/arrow/compute/expression.cc | 211 ++++++++- cpp/src/arrow/compute/expression_internal.h | 8 + cpp/src/arrow/compute/expression_test.cc | 429 +++++++++++++++++- .../compute/kernels/scalar_set_lookup.cc | 61 ++- .../kernels/scalar_set_lookup_internal.h | 327 +++++++++++++ cpp/src/arrow/dataset/file_parquet_test.cc | 78 ++++ 6 files changed, 1092 insertions(+), 22 deletions(-) create mode 100644 cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index cd952055666c..6965a6da55a9 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,184 @@ 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; + } + + // Strip the order preserving casts wrapping the argument one at a time, keeping + // them (outermost first) so that bounds can be put through the same casts. + std::vector casts; + const Expression* lhs = &call->arguments[0]; + while (lhs->call() && lhs->call()->function_name == "cast" && + &Comparison::StripOrderPreservingCasts(*lhs) != lhs) { + casts.push_back(lhs->call()); + lhs = &lhs->call()->arguments[0]; + } + if (!lhs->field_ref()) return expr; + const FieldRef& target = *lhs->field_ref(); + // The stripped casts 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 weaken_strict_bounds = + !casts.empty() && 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); + } + if (applicable.empty()) return expr; + + const auto& state = + checked_cast(*call->kernel_state); + // 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 (!state.value_set_type->Equals(*call->arguments[0].type())) return expr; + + // Only materialize the (lazily sorted) value set once we know it can be used. + std::shared_ptr sorted = state.GetSortedValueSet(); + if (!sorted) return expr; + DCHECK(sorted->type()->Equals(*call->arguments[0].type())); + + // 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; + + // Ignoring a bound only loses precision, so every conversion below bails out + // by skipping the inequality. + std::shared_ptr bound = inequality->bound.scalar(); + // A comparison against null or NaN cannot be guaranteed true, but be defensive. + if (!bound->is_valid || internal::IsNaNScalar(*bound)) continue; + if (!bound->type->Equals(*lhs->type())) { + // The guarantee compares the field with a bound of another type (in the + // promoted type). Bring the bound into the field's type exactly: "safe" + // casts between floating point types or to decimals may still round, so + // require the cast to round trip. + auto maybe_cast = + Cast(Datum(bound), lhs->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(); + } + // Put the bound through the casts the field goes through, innermost first and + // with the same options. Since each cast is nondecreasing, `x <= b` implies + // `cast(x) <= cast(b)` even if `cast(b)` rounds, so no exactness is needed. + bool cast_failed = false; + for (auto it = casts.rbegin(); it != casts.rend(); ++it) { + const auto& cast_options = checked_cast(*(*it)->options); + auto maybe_cast = Cast(Datum(bound), cast_options, &exec_context); + if (!maybe_cast.ok()) { + cast_failed = true; + break; + } + bound = maybe_cast->scalar(); + } + if (cast_failed || !bound->type->Equals(*sorted->type())) continue; + 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 +1632,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_internal.h b/cpp/src/arrow/compute/expression_internal.h index 50b08a3d7a0f..05b41c8d7f54 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -27,6 +27,7 @@ #include "arrow/compute/registry.h" #include "arrow/record_batch.h" #include "arrow/table.h" +#include "arrow/util/checked_cast.h" #include "arrow/util/logging.h" namespace arrow { @@ -139,6 +140,13 @@ struct Comparison { if (is_signed_integer(to_id)) { if (is_integer(from_id) && bit_width(to_id) >= bit_width(from_id)) { + // An unsigned to signed cast of the same width wraps negative when overflow + // is allowed, which does not preserve ordering. + if (is_unsigned_integer(from_id) && bit_width(to_id) == bit_width(from_id) && + ::arrow::internal::checked_cast(*call->options) + .allow_int_overflow) { + return expr; + } return StripOrderPreservingCasts(from); } return expr; diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 10c1b348e62d..092901ad93ae 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) { @@ -1702,6 +1710,16 @@ struct Simplify { TEST(Expression, SingleComparisonGuarantees) { auto i32 = field_ref("i32"); + // An unsafe unsigned to signed cast of the same width may wrap, so it is not + // stripped: u32 = 4294967221 casts to -75. + Simplify{ + equal(call("cast", {field_ref("u32")}, CastOptions::Unsafe(int32())), literal(-75))} + .WithGuarantee(greater_equal(field_ref("u32"), literal(1))) + .ExpectUnchanged(); + Simplify{equal(cast(field_ref("u32"), int32()), literal(-75))} + .WithGuarantee(greater_equal(field_ref("u32"), literal(1))) + .Expect(false); + // i32 is guaranteed equal to 3, so the projection can just materialize that constant // and need not incur IO Simplify{project({add(i32, literal(1))}, {"i32 + 1"})} @@ -2063,8 +2081,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 +2169,337 @@ 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(); + + // 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); + + // Bounds go through the whole cast chain: 0.1 rounds up to 0.1f and back to + // double(0.1f), which is a set member, so f64 <= 0.1 cannot be ruled out. + auto f32_roundtrip = cast(cast(f64, float32()), float64()); + auto roundtrip_set = concat_json("[0.10000000149011612]", make_range_json(1, 50)); + Simplify{is_in(f32_roundtrip, float64(), roundtrip_set, null_matching)} + .WithGuarantee(less_equal(f64, literal(0.1))) + .ExpectUnchanged(); + Simplify{is_in(f32_roundtrip, float64(), roundtrip_set, null_matching)} + .WithGuarantee(less_equal(f64, literal(0.05))) + .Expect(false); + // ... and strict bounds are weakened: f64 = 50.00000001 round trips to 50.0 + Simplify{is_in(f32_roundtrip, float64(), roundtrip_set, null_matching)} + .WithGuarantee(greater(f64, literal(50.0))) + .ExpectUnchanged(); + Simplify{is_in(f32_roundtrip, float64(), roundtrip_set, null_matching)} + .WithGuarantee(greater(f64, literal(50.5))) + .Expect(false); + + // An unsafe unsigned to signed cast of the same width wraps and is not order + // preserving: u32 = 4294967221 casts to -75, a set member. + auto negative_set = make_range_json(-100, -50); + Simplify{is_in(call("cast", {u32}, CastOptions::Unsafe(int32())), int32(), + negative_set, null_matching)} + .WithGuarantee(greater_equal(u32, literal(1))) + .ExpectUnchanged(); + Simplify{is_in(cast(u32, int32()), int32(), negative_set, null_matching)} + .WithGuarantee(greater_equal(u32, literal(1))) + .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 +2527,82 @@ 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); + } + + // A chain of casts which rounds an intermediate value and an unsafe cast which + // wraps must not make the simplified filter drop matching rows either. + struct FilterCase { + Expression filter; + Expression guarantee; + std::string input; + }; + for (const FilterCase& c : std::vector{ + {call("is_in", {cast(cast(field_ref("f64"), float32()), float64())}, + compute::SetLookupOptions{ArrayFromJSON( + float64(), + concat_json("[0.10000000149011612]", make_range_json(1, 50)))}), + less_equal(field_ref("f64"), literal(0.1)), + R"([{"f64": 0.09}, {"f64": 0.1}])"}, + {call("is_in", + {call("cast", {field_ref("u32")}, CastOptions::Unsafe(int32()))}, + compute::SetLookupOptions{ + ArrayFromJSON(int32(), make_range_json(-100, -50))}), + greater_equal(field_ref("u32"), literal(1)), + R"([{"u32": 1}, {"u32": 4294967221}])"}, + }) { + ASSERT_OK_AND_ASSIGN(auto bound_filter, c.filter.Bind(*kBoringSchema)); + auto input = RecordBatchFromJSON(kBoringSchema, c.input); + ASSERT_OK_AND_ASSIGN(auto simplified, + SimplifyWithGuarantee(bound_filter, c.guarantee)); + Datum evaluated, simplified_evaluated; + ExpectExecute(bound_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 c6a3562233c0..676661b02877 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 000000000000..37150853a2c3 --- /dev/null +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -0,0 +1,327 @@ +// 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 "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/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); + +/// 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) { + 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 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; +} + +/// 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 0d86f5eaa0d3..26bf8561f986 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; From ebe8e731b3e7ba1392572348785ba6d1056f2f98 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 27 Oct 2026 20:29:36 +0100 Subject: [PATCH 2/4] GH-55292: [C++] Sort primitive is_in value sets directly instead of via sort_indices + take Sorting a copy of the values with std::sort is about 14x faster than the generic sort_indices + take path for a shuffled set of 10M int64 values (0.19s vs 2.6s). Nulls are dropped while copying and a set containing NaN is reported as not searchable before sorting, since NaN would break the strict weak ordering std::sort requires. Non-primitive types (binary, decimal) keep using the generic path. --- cpp/src/arrow/compute/expression_test.cc | 24 ++++++++ .../kernels/scalar_set_lookup_internal.h | 59 +++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 092901ad93ae..500c3d34dcbb 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2315,6 +2315,30 @@ TEST(Expression, SimplifyIsInLargeValueSet) { 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)} diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h index 37150853a2c3..60011e4a6760 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -20,18 +20,23 @@ // 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" @@ -66,6 +71,7 @@ inline bool IsNaNScalar(const Scalar& value) { } 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 { @@ -119,6 +125,12 @@ struct SetLookupStateBase : public KernelState { 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))); @@ -268,6 +280,43 @@ struct IsSortedVisitor { } }; +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; @@ -309,6 +358,16 @@ inline Result IsSortedValueSet(const Array& values) { 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() From 691fa5be68f0aeb9a8d517e6346993328af58775 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Thu, 29 Oct 2026 19:16:11 +0100 Subject: [PATCH 3/4] GH-55292: [C++] Handle nullable ranges and legacy null options in is_in pruning Extract nullable min/max bounds from Parquet statistics for large value sets. Respect GetNullMatchingBehavior() in both simplification paths and when rebinding narrowed sets. Add expression and Parquet regressions for nullable ranges and conflicting deprecated skip_nulls options. --- cpp/src/arrow/compute/expression.cc | 33 +++++- cpp/src/arrow/compute/expression_test.cc | 118 +++++++++++++++++++++ cpp/src/arrow/dataset/file_parquet_test.cc | 66 ++++++++++++ 3 files changed, 212 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 6965a6da55a9..148bae8c45ce 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1205,6 +1205,29 @@ struct Inequality { return ExtractOneFromComparison(guarantee); } + // Also extract the bounds of a nullable range, as produced by Parquet: + // (cmp(a, min) and cmp(a, max)) or is_null(a). + // Each bound is only guaranteed for non-null values of the same field. + static std::vector Extract(const Expression& guarantee) { + if (auto inequality = ExtractOne(guarantee)) return {*inequality}; + + auto call = guarantee.call(); + if (!call || call->function_name != "or_kleene") return {}; + auto null_call = call->arguments[1].call(); + if (!null_call || null_call->function_name != "is_null") return {}; + auto target = null_call->arguments[0].field_ref(); + if (!target) return {}; + + std::vector inequalities; + for (const auto& member : GuaranteeConjunctionMembers(call->arguments[0])) { + auto inequality = ExtractOneFromComparison(member); + if (!inequality || inequality->target != *target) continue; + inequality->nullable = true; + inequalities.push_back(*inequality); + } + return inequalities; + } + static std::optional ExtractOneFromComparison(const Expression& guarantee) { auto call = guarantee.call(); if (!call) return std::nullopt; @@ -1283,7 +1306,7 @@ struct Inequality { if (*lhs.field_ref() != guarantee.target) return std::nullopt; FilterOptions::NullSelectionBehavior null_selection{}; - switch (options->null_matching_behavior) { + switch (options->GetNullMatchingBehavior()) { case SetLookupOptions::MATCH: null_selection = guarantee.nullable ? FilterOptions::EMIT_NULL : FilterOptions::DROP; @@ -1320,7 +1343,7 @@ struct Inequality { simplified_call.function_name = "is_in"; simplified_call.arguments = is_in_call->arguments; simplified_call.options = std::make_shared( - simplified_value_set, options->null_matching_behavior); + simplified_value_set, options->GetNullMatchingBehavior()); ARROW_ASSIGN_OR_RAISE( Expression simplified_expr, BindNonRecursive(std::move(simplified_call), @@ -1417,7 +1440,7 @@ struct Inequality { /// \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) { + switch (options.GetNullMatchingBehavior()) { case SetLookupOptions::MATCH: // A null in the value set matches nulls in the data. return !(nullable && value_set_has_null); @@ -1636,8 +1659,8 @@ Result SimplifyWithGuarantee(Expression expr, // 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); + for (const auto& inequality : Inequality::Extract(guarantee)) { + inequalities.push_back(inequality); } } if (!inequalities.empty()) { diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 500c3d34dcbb..0762eeb5d093 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2524,6 +2524,124 @@ TEST(Expression, SimplifyIsInLargeValueSet) { .Expect(greater(field_ref("u32"), literal(5))); } +TEST(Expression, SimplifyIsInLargeValueSetNullableRange) { + auto i32 = field_ref("i32"); + auto u32 = field_ref("u32"); + auto gapped = concat_json(make_range_json(1, 50), make_range_json(100, 150)); + // This is the shape of Parquet statistics when a row group may contain nulls. + auto range = and_(greater_equal(i32, literal(51)), less_equal(i32, literal(99))); + auto nullable_range = or_(range, is_null(i32)); + + for (auto null_matching : + {SetLookupOptions::MATCH, SetLookupOptions::SKIP, SetLookupOptions::EMIT_NULL, + SetLookupOptions::INCONCLUSIVE}) { + for (bool value_set_has_null : {false, true}) { + SCOPED_TRACE(null_matching); + SCOPED_TRACE(value_set_has_null); + auto values = ArrayFromJSON( + int32(), value_set_has_null ? concat_json(gapped, "[null]") : gapped); + auto filter = call("is_in", {i32}, SetLookupOptions{values, null_matching}); + auto simplified = Simplify{filter}.WithGuarantee(nullable_range); + if (null_matching == SetLookupOptions::SKIP || + (null_matching == SetLookupOptions::MATCH && !value_set_has_null)) { + simplified.Expect(false); + } else { + simplified.ExpectUnchanged(); + } + + // A matching non-null member prevents pruning regardless of null behavior. + Simplify{filter} + .WithGuarantee( + or_(and_(greater_equal(i32, literal(50)), less_equal(i32, literal(99))), + is_null(i32))) + .ExpectUnchanged(); + + // Do not infer bounds on a different field from the null branch. + Simplify{filter}.WithGuarantee(or_(range, is_null(u32))).ExpectUnchanged(); + auto u32_range = + and_(greater_equal(u32, literal(51)), less_equal(u32, literal(99))); + Simplify{filter} + .WithGuarantee(or_(and_(range, u32_range), is_null(u32))) + .ExpectUnchanged(); + + // Another inequality can exclude nulls while the nullable range supplies + // both bounds needed to find the gap in the value set. + auto non_nullable = + Simplify{filter}.WithGuarantee(and_(nullable_range, greater(i32, literal(0)))); + if (null_matching == SetLookupOptions::INCONCLUSIVE && value_set_has_null) { + non_nullable.ExpectUnchanged(); + } else { + non_nullable.Expect(false); + } + + // Verify that simplification also preserves the values on null input. + ASSERT_OK_AND_ASSIGN(auto bound_filter, filter.Bind(*kBoringSchema)); + ASSERT_OK_AND_ASSIGN(auto simplified_filter, + SimplifyWithGuarantee(bound_filter, nullable_range)); + auto input = RecordBatchFromJSON( + kBoringSchema, R"([{"i32": null}, {"i32": 51}, {"i32": 75}, {"i32": 99}])"); + Datum evaluated, simplified_evaluated; + ExpectExecute(bound_filter, input, &evaluated); + ExpectExecute(simplified_filter, 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, SimplifyIsInDeprecatedSkipNulls) { + auto i32 = field_ref("i32"); + auto input = RecordBatchFromJSON(kBoringSchema, + R"([{"i32": null}, {"i32": 1001}, {"i32": 1002}])"); + // Cover both the filter-and-rebind and the large-value-set paths, including + // options whose deprecated field overrides a conflicting modern enum. + for (int num_values : {10, 100}) { + for (auto behavior : {SetLookupOptions::MATCH, SetLookupOptions::SKIP, + SetLookupOptions::EMIT_NULL, SetLookupOptions::INCONCLUSIVE}) { + for (bool skip_nulls : {false, true}) { + SCOPED_TRACE(num_values); + SCOPED_TRACE(behavior); + SCOPED_TRACE(skip_nulls); + SetLookupOptions options{ + ArrayFromJSON(int32(), concat_json(make_range_json(1, num_values), "[null]")), + behavior}; + options.skip_nulls = skip_nulls; + ASSERT_OK_AND_ASSIGN(auto filter, + call("is_in", {i32}, options).Bind(*kBoringSchema)); + for (int bound : {3, 1000}) { + SCOPED_TRACE(bound); + auto guarantee = or_(greater(i32, literal(bound)), is_null(i32)); + ASSERT_OK_AND_ASSIGN(auto simplified, SimplifyWithGuarantee(filter, guarantee)); + if (bound == 1000) { + if (skip_nulls) { + EXPECT_EQ(simplified, literal(false)); + } else { + EXPECT_TRUE(simplified.IsSatisfiable()); + } + } + + 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())); + } + auto expected = ArrayFromJSON( + boolean(), skip_nulls ? "[false, false, false]" : "[true, false, false]"); + AssertDatumsEqual(expected, evaluated, /*verbose=*/true); + AssertDatumsEqual(evaluated, simplified_evaluated, /*verbose=*/true); + } + } + } + } +} + TEST(Expression, SimplifyThenExecute) { auto filter = or_({equal(field_ref("f32"), literal(0)), diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 26bf8561f986..71f30b6b5683 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -825,6 +826,71 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsIn) { {0, 1, 4}); } +TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsInWithNulls) { + // Each row group contains ten distinct values and a null, so its statistics + // yield (i64 >= min and i64 <= max) or is_null(i64). + auto dataset_schema = schema({field("i64", int64())}); + std::vector> batches; + for (int i = 0; i < 3; ++i) { + std::vector values(11); + std::iota(values.begin(), values.end(), 10 * i); + std::vector valid(values.size(), true); + valid.back() = false; + std::shared_ptr array; + ArrayFromVector(valid, values, &array); + batches.push_back(RecordBatch::Make(dataset_schema, array->length(), {array})); + } + 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); + + // The hull spans all three row groups but no member occurs in [10, 19]. + std::vector values(70); + std::iota(values.begin(), values.begin() + 10, 0); + std::iota(values.begin() + 10, values.end(), 20); + ASSERT_GT(values.size(), 50); + auto expect_row_groups = [&](bool value_set_has_null, + compute::SetLookupOptions::NullMatchingBehavior behavior, + const std::vector& expected, + std::optional skip_nulls = std::nullopt) { + SCOPED_TRACE(value_set_has_null); + SCOPED_TRACE(behavior); + auto set_values = values; + std::vector valid(values.size(), true); + if (value_set_has_null) { + set_values.push_back(0); + valid.push_back(false); + } + std::shared_ptr value_set; + ArrayFromVector(valid, set_values, &value_set); + compute::SetLookupOptions options{std::move(value_set), behavior}; + options.skip_nulls = skip_nulls; + SetFilter(call("is_in", {field_ref("i64")}, options)); + 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); + }; + + expect_row_groups(false, compute::SetLookupOptions::MATCH, {0, 2}); + // A null in a MATCH set matches the nulls in every row group. + expect_row_groups(true, compute::SetLookupOptions::MATCH, {0, 1, 2}); + expect_row_groups(false, compute::SetLookupOptions::SKIP, {0, 2}); + expect_row_groups(true, compute::SetLookupOptions::SKIP, {0, 2}); + // The deprecated flag takes precedence over a conflicting enum value. + expect_row_groups(true, compute::SetLookupOptions::SKIP, {0, 1, 2}, false); + expect_row_groups(true, compute::SetLookupOptions::MATCH, {0, 2}, true); +} + TEST_P(TestParquetFileFormatScan, ExplicitRowGroupSelection) { constexpr int64_t kNumRowGroups = 16; constexpr int64_t kTotalNumRows = kNumRowGroups * (kNumRowGroups + 1) / 2; From a3f317c7aeb85e94b7b4a538a2c47cc1b388d176 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Thu, 29 Oct 2026 20:43:41 +0100 Subject: [PATCH 4/4] GH-55292: [C++] Restate guarantees through non-injective casts and prune nullable ranges in every simplification path Inequality guarantees are now restated through the order preserving casts wrapping a filter's argument before they are applied: strict bounds are weakened to inclusive ones and the bound is put through the same casts, since e.g. f64 = 51.00000001 satisfies f64 > 51 but casts to 51f. This applies to plain comparisons and small is_in value sets as well as the large value set path, which previously was the only one handling it. The nullable range guarantee produced by Parquet statistics for row groups with nulls, (x >= min and x <= max) or is_null(x), is now extracted for all simplifications rather than only for large is_in value sets, and an is_valid(x) member of the guarantee rules out nulls for the inequalities on x. --- cpp/src/arrow/compute/expression.cc | 258 ++++++++++++--------- cpp/src/arrow/compute/expression_test.cc | 133 +++++++++++ cpp/src/arrow/dataset/file_parquet_test.cc | 61 +++-- 3 files changed, 326 insertions(+), 126 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 148bae8c45ce..731ba0d42ee3 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1161,6 +1161,64 @@ Result Canonicalize(Expression expr, compute::ExecContext* exec_cont namespace { +// The order preserving casts wrapping an Expression, outermost first, and the +// wrapped Expression. See Comparison::StripOrderPreservingCasts. +struct OrderPreservingCasts { + std::vector casts; + const Expression* inner; + // Whether the composed casts map distinct values to distinct values. Each cast + // is nondecreasing, so `x > b` implies `cast(x) >= cast(b)`, but only through + // injective casts does it imply `cast(x) > cast(b)`. Integer widenings are + // injective, as are casts to a floating point type whose mantissa holds every + // value of the source type; float64 to float32 or int64 to float64 are not. + bool injective = true; + + static OrderPreservingCasts Strip(const Expression& expr) { + OrderPreservingCasts out; + const Expression* lhs = &expr; + while (lhs->call() && lhs->call()->function_name == "cast" && + &Comparison::StripOrderPreservingCasts(*lhs) != lhs) { + const Expression& from = lhs->call()->arguments[0]; + out.injective = out.injective && IsInjective(from.type()->id(), lhs->type()->id()); + out.casts.push_back(lhs->call()); + lhs = &from; + } + out.inner = lhs; + return out; + } + + static bool IsInjective(Type::type from, Type::type to) { + // Only integer widenings are stripped between integer types. + if (!is_floating(to)) return true; + if (is_floating(from)) return bit_width(from) <= bit_width(to); + const int mantissa_bits = to == Type::DOUBLE ? 53 : to == Type::FLOAT ? 24 : 11; + return bit_width(from) <= mantissa_bits; + } + + // Put a bound on the wrapped Expression through the casts, innermost first and + // with the same options. + Result> ApplyToBound(std::shared_ptr bound, + ExecContext* exec_context) const { + for (auto it = casts.rbegin(); it != casts.rend(); ++it) { + const auto& options = checked_cast(*(*it)->options); + ARROW_ASSIGN_OR_RAISE(Datum cast, + Cast(Datum(std::move(bound)), options, exec_context)); + bound = cast.scalar(); + } + return bound; + } + + // The comparison between the cast Expression and a cast bound implied by `cmp` + // between the wrapped Expression and the bound: strict bounds are weakened to + // inclusive ones unless the casts are injective. + Comparison::type ApplyToComparison(Comparison::type cmp) const { + if (injective) return cmp; + if (cmp == Comparison::GREATER) return Comparison::GREATER_EQUAL; + if (cmp == Comparison::LESS) return Comparison::LESS_EQUAL; + return cmp; + } +}; + // An inequality comparison which a target Expression is known to satisfy. If nullable, // the target may evaluate to null in addition to values satisfying the comparison. struct Inequality { @@ -1173,51 +1231,27 @@ struct Inequality { // Whether target can be null bool nullable; - // Extract an Inequality if possible, derived from "less", - // "greater", "less_equal", and "greater_equal" expressions, - // possibly disjuncted with an "is_null" Expression. - // cmp(a, 2) - // cmp(a, 2) or is_null(a) - static std::optional ExtractOne(const Expression& guarantee) { - auto call = guarantee.call(); - if (!call) return std::nullopt; - - if (call->function_name == "or_kleene") { - // expect the LHS to be a usable field inequality - auto out = ExtractOneFromComparison(call->arguments[0]); - if (!out) return std::nullopt; - - // expect the RHS to be an is_null expression - auto call_rhs = call->arguments[1].call(); - if (!call_rhs) return std::nullopt; - if (call_rhs->function_name != "is_null") return std::nullopt; - - // ... and that it references the same target - auto target = call_rhs->arguments[0].field_ref(); - if (!target) return std::nullopt; - if (*target != out->target) return std::nullopt; - - out->nullable = true; - return out; - } - - // fall back to a simple comparison with no "is_null" - return ExtractOneFromComparison(guarantee); - } - - // Also extract the bounds of a nullable range, as produced by Parquet: - // (cmp(a, min) and cmp(a, max)) or is_null(a). - // Each bound is only guaranteed for non-null values of the same field. + // Extract the inequalities implied by a guarantee, derived from "less", + // "greater", "less_equal", "greater_equal" and "equal" expressions, possibly + // disjuncted with an "is_null" Expression: + // cmp(a, 2) + // cmp(a, 2) or is_null(a) + // (cmp(a, 2) and cmp(a, 5)) or is_null(a) + // The last is the shape of a Parquet row group's statistics when it may contain + // nulls. In a disjunction each bound only holds for non-null values of the field. static std::vector Extract(const Expression& guarantee) { - if (auto inequality = ExtractOne(guarantee)) return {*inequality}; + if (auto inequality = ExtractOneFromComparison(guarantee)) return {*inequality}; auto call = guarantee.call(); if (!call || call->function_name != "or_kleene") return {}; + + // expect the RHS to be an is_null expression ... auto null_call = call->arguments[1].call(); if (!null_call || null_call->function_name != "is_null") return {}; auto target = null_call->arguments[0].field_ref(); if (!target) return {}; + // ... and the LHS to be a conjunction of usable inequalities on the same field std::vector inequalities; for (const auto& member : GuaranteeConjunctionMembers(call->arguments[0])) { auto inequality = ExtractOneFromComparison(member); @@ -1249,6 +1283,20 @@ struct Inequality { return std::nullopt; } + /// This inequality restated for the target wrapped in `casts`: through casts + /// which are not injective `x > b` only implies `cast(x) >= cast(b)`, so the + /// comparison is weakened and the bound is put through the casts (`b` may round + /// to a value the expression compares with). Nullopt if the bound cannot be cast. + std::optional> ThroughCasts( + const OrderPreservingCasts& casts) const { + Comparison::type cast_cmp = casts.ApplyToComparison(cmp); + if (casts.injective) return std::make_pair(cast_cmp, bound); + ExecContext exec_context; + auto cast_bound = casts.ApplyToBound(bound.scalar(), &exec_context); + if (!cast_bound.ok()) return std::nullopt; + return std::make_pair(cast_cmp, Datum(*cast_bound)); + } + /// The given expression simplifies to `value` if the inequality /// target is not nullable. Otherwise, it simplifies to either a /// call to true_unless_null or !true_unless_null. @@ -1301,10 +1349,17 @@ struct Inequality { return std::nullopt; } - const auto& lhs = Comparison::StripOrderPreservingCasts(is_in_call->arguments[0]); + auto casts = OrderPreservingCasts::Strip(is_in_call->arguments[0]); + const Expression& lhs = *casts.inner; if (!lhs.field_ref()) return std::nullopt; if (*lhs.field_ref() != guarantee.target) return std::nullopt; + // The guarantee bounds the field but the value set is matched against the + // cast field. + auto through_casts = guarantee.ThroughCasts(casts); + if (!through_casts) return std::nullopt; + const auto& [cmp, bound] = *through_casts; + FilterOptions::NullSelectionBehavior null_selection{}; switch (options->GetNullMatchingBehavior()) { case SetLookupOptions::MATCH: @@ -1327,9 +1382,9 @@ struct Inequality { break; } - std::string func_name = Comparison::GetName(guarantee.cmp); + std::string func_name = Comparison::GetName(cmp); DCHECK_NE(func_name, "na"); - std::vector args{options->value_set, guarantee.bound}; + std::vector args{options->value_set, bound}; ARROW_ASSIGN_OR_RAISE(Datum filter_mask, CallFunction(func_name, args)); FilterOptions filter_options(null_selection); ARROW_ASSIGN_OR_RAISE(Datum simplified_value_set, @@ -1352,7 +1407,7 @@ struct Inequality { } /// \brief Simplify the given expression given this inequality as a guarantee. - Result Simplify(Expression expr) { + Result Simplify(Expression expr) const { const auto& guarantee = *this; auto call = expr.call(); @@ -1380,26 +1435,32 @@ struct Inequality { if (!rhs) return expr; if (!rhs->is_scalar()) return expr; - const auto& lhs = Comparison::StripOrderPreservingCasts(call->arguments[0]); + auto casts = OrderPreservingCasts::Strip(call->arguments[0]); + const Expression& lhs = *casts.inner; if (!lhs.field_ref()) return expr; if (*lhs.field_ref() != guarantee.target) return expr; + // The guarantee bounds the field but the expression compares the cast field. + auto through_casts = guarantee.ThroughCasts(casts); + if (!through_casts) return expr; + const auto& [guarantee_cmp, bound] = *through_casts; + // Whether the RHS of the expression is EQUAL, LESS, or GREATER than the // RHS of the guarantee. N.B. Comparison::type is a bitmask ARROW_ASSIGN_OR_RAISE(const Comparison::type cmp_rhs_bound, - Comparison::Execute(*rhs, guarantee.bound)); + Comparison::Execute(*rhs, bound)); DCHECK_NE(cmp_rhs_bound, Comparison::NA); if (cmp_rhs_bound == Comparison::EQUAL) { // RHS of filter is equal to RHS of guarantee - if ((*cmp & guarantee.cmp) == guarantee.cmp) { + if ((*cmp & guarantee_cmp) == guarantee_cmp) { // guarantee is a subset of filter, so all data will be included // x > 1, x >= 1, x != 1 guaranteed by x > 1 return simplified_to(lhs, true); } - if ((*cmp & guarantee.cmp) == 0) { + if ((*cmp & guarantee_cmp) == 0) { // guarantee disjoint with filter, so all data will be excluded // x > 1, x >= 1 unsatisfiable if x == 1 return simplified_to(lhs, false); @@ -1408,12 +1469,12 @@ struct Inequality { return expr; } - if (guarantee.cmp & cmp_rhs_bound) { + if (guarantee_cmp & cmp_rhs_bound) { // We guarantee (x (?) N) and are trying to simplify (x (?) M). We know // either M < N or M > N (i.e. cmp_rhs_bound is either LESS or GREATER). // If M > N, then if the guarantee is (x > N), (x >= N), or (x != N) - // (i.e. guarantee.cmp & cmp_rhs_bound), we cannot do anything with the + // (i.e. guarantee_cmp & cmp_rhs_bound), we cannot do anything with the // guarantee, and bail out here. // For example, take M = 5, N = 3. Then cmp_rhs_bound = GREATER. @@ -1425,7 +1486,7 @@ struct Inequality { if (*cmp & Comparison::GetFlipped(cmp_rhs_bound)) { // x > 1, x >= 1, x != 1 guaranteed by x >= 3 - // (where `guarantee.cmp` is GREATER_EQUAL, `cmp_rhs_bound` is LESS) + // (where `guarantee_cmp` is GREATER_EQUAL, `cmp_rhs_bound` is LESS) return simplified_to(lhs, true); } else { // x < 1, x <= 1, x == 1 unsatisfiable if x >= 3 @@ -1482,23 +1543,12 @@ Result SimplifyIsInWithInequalities( return expr; } - // Strip the order preserving casts wrapping the argument one at a time, keeping - // them (outermost first) so that bounds can be put through the same casts. - std::vector casts; - const Expression* lhs = &call->arguments[0]; - while (lhs->call() && lhs->call()->function_name == "cast" && - &Comparison::StripOrderPreservingCasts(*lhs) != lhs) { - casts.push_back(lhs->call()); - lhs = &lhs->call()->arguments[0]; - } + // Strip the order preserving casts wrapping the argument, keeping them so that + // bounds can be put through the same casts. + auto casts = OrderPreservingCasts::Strip(call->arguments[0]); + const Expression* lhs = casts.inner; if (!lhs->field_ref()) return expr; const FieldRef& target = *lhs->field_ref(); - // The stripped casts 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 weaken_strict_bounds = - !casts.empty() && is_floating(call->arguments[0].type()->id()); std::vector applicable; for (const Inequality& inequality : inequalities) { @@ -1553,27 +1603,18 @@ Result SimplifyIsInWithInequalities( } bound = maybe_cast->scalar(); } - // Put the bound through the casts the field goes through, innermost first and - // with the same options. Since each cast is nondecreasing, `x <= b` implies - // `cast(x) <= cast(b)` even if `cast(b)` rounds, so no exactness is needed. - bool cast_failed = false; - for (auto it = casts.rbegin(); it != casts.rend(); ++it) { - const auto& cast_options = checked_cast(*(*it)->options); - auto maybe_cast = Cast(Datum(bound), cast_options, &exec_context); - if (!maybe_cast.ok()) { - cast_failed = true; - break; - } - bound = maybe_cast->scalar(); - } - if (cast_failed || !bound->type->Equals(*sorted->type())) continue; + // Put the bound through the casts the field goes through. Since each cast is + // nondecreasing, `x <= b` implies `cast(x) <= cast(b)` even if `cast(b)` + // rounds, so no exactness is needed. + auto cast_bound = casts.ApplyToBound(std::move(bound), &exec_context); + if (!cast_bound.ok()) continue; + bound = *cast_bound; + if (!bound->type->Equals(*sorted->type())) continue; 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; - } + // Strict bounds are weakened to inclusive ones through casts which are not + // injective: f64 = 51.00000001 satisfies f64 > 51 but casts to 51f. + Comparison::type cmp = casts.ApplyToComparison(inequality->cmp); ARROW_ASSIGN_OR_RAISE( int64_t lower, internal::SearchSortedValueSet(*sorted, *bound, /*upper=*/false)); @@ -1655,14 +1696,26 @@ 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. + // The inequalities on fields implied by the guarantee. A field which is also + // guaranteed valid cannot be null, whatever its inequalities admit. std::vector inequalities; for (const auto& guarantee : conjunction_members) { for (const auto& inequality : Inequality::Extract(guarantee)) { inequalities.push_back(inequality); } } + for (const auto& guarantee : conjunction_members) { + auto call = guarantee.call(); + if (!call || call->function_name != "is_valid") continue; + auto target = call->arguments[0].field_ref(); + if (!target) continue; + for (auto& inequality : inequalities) { + if (inequality.target == *target) inequality.nullable = false; + } + } + + // `is_in` calls with large value sets are simplified against all inequalities + // on their argument at once, see SimplifyIsInWithInequalities. if (!inequalities.empty()) { Expression before = expr; ARROW_ASSIGN_OR_RAISE(expr, ModifyExpression( @@ -1677,33 +1730,30 @@ Result SimplifyWithGuarantee(Expression expr, } } - for (const auto& guarantee : conjunction_members) { - if (!guarantee.call()) continue; + for (const Inequality& inequality : inequalities) { + ARROW_ASSIGN_OR_RAISE(auto simplified, + ModifyExpression( + std::move(expr), [](Expression expr) { return expr; }, + [&](Expression expr, ...) -> Result { + return inequality.Simplify(std::move(expr)); + })); - if (auto inequality = Inequality::ExtractOne(guarantee)) { - ARROW_ASSIGN_OR_RAISE(auto simplified, - ModifyExpression( - std::move(expr), [](Expression expr) { return expr; }, - [&](Expression expr, ...) -> Result { - return inequality->Simplify(std::move(expr)); - })); + if (Expression::Identical(simplified, expr)) continue; - if (Expression::Identical(simplified, expr)) continue; + expr = std::move(simplified); + RETURN_NOT_OK(CanonicalizeAndFoldConstants()); + } - expr = std::move(simplified); - RETURN_NOT_OK(CanonicalizeAndFoldConstants()); - } + for (const auto& guarantee : conjunction_members) { + if (!guarantee.call() || guarantee.call()->function_name != "is_valid") continue; - if (guarantee.call()->function_name == "is_valid") { - ARROW_ASSIGN_OR_RAISE( - auto simplified, - SimplifyIsValidGuarantee(std::move(expr), *CallNotNull(guarantee))); + ARROW_ASSIGN_OR_RAISE(auto simplified, SimplifyIsValidGuarantee( + std::move(expr), *CallNotNull(guarantee))); - if (Expression::Identical(simplified, expr)) continue; + if (Expression::Identical(simplified, expr)) continue; - expr = std::move(simplified); - RETURN_NOT_OK(CanonicalizeAndFoldConstants()); - } + expr = std::move(simplified); + RETURN_NOT_OK(CanonicalizeAndFoldConstants()); } return expr; diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 0762eeb5d093..41e077a70a6f 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2021,6 +2021,132 @@ TEST(Expression, SimplifyWithComparisonAndNullableCaveat) { true_unless_null(field_ref("i32")))); // not satisfiable, will drop row group } +TEST(Expression, SimplifyWithNullableRangeGuarantee) { + // The shape of Parquet statistics for a row group which may contain nulls: both + // bounds hold for every non-null value. + auto i32 = field_ref("i32"); + auto u32 = field_ref("u32"); + auto range = and_(greater_equal(i32, literal(1)), less_equal(i32, literal(5))); + auto nullable_range = or_(range, is_null(i32)); + + Simplify{greater(i32, literal(10))} + .WithGuarantee(nullable_range) + .Expect(not_(true_unless_null(i32))); + Simplify{less(i32, literal(0))} + .WithGuarantee(nullable_range) + .Expect(not_(true_unless_null(i32))); + Simplify{less(i32, literal(10))} + .WithGuarantee(nullable_range) + .Expect(true_unless_null(i32)); + Simplify{less(i32, literal(3))}.WithGuarantee(nullable_range).ExpectUnchanged(); + Simplify{is_null(i32)}.WithGuarantee(nullable_range).ExpectUnchanged(); + Simplify{is_valid(i32)}.WithGuarantee(nullable_range).ExpectUnchanged(); + + // Small value sets are filtered by both bounds + auto is_in = [&](std::string json_array) { + return call( + "is_in", {i32}, + SetLookupOptions{ArrayFromJSON(int32(), json_array), SetLookupOptions::MATCH}); + }; + Simplify{is_in("[10, 11]")}.WithGuarantee(nullable_range).Expect(false); + Simplify{is_in("[0, 3, 11]")}.WithGuarantee(nullable_range).Expect(is_in("[3]")); + + // Bounds are not inferred for a field other than the one which may be null + Simplify{greater(i32, literal(10))} + .WithGuarantee(or_(range, is_null(u32))) + .ExpectUnchanged(); + + // is_valid elsewhere in the guarantee rules out the nulls + auto valid_range = and_(nullable_range, is_valid(i32)); + Simplify{greater(i32, literal(10))}.WithGuarantee(valid_range).Expect(false); + Simplify{less(i32, literal(10))}.WithGuarantee(valid_range).Expect(true); + Simplify{is_null(i32)}.WithGuarantee(valid_range).Expect(false); +} + +TEST(Expression, SimplifyComparisonThroughNonInjectiveCast) { + auto i32 = field_ref("i32"); + auto i64 = field_ref("i64"); + auto f32 = field_ref("f32"); + auto f64 = field_ref("f64"); + + // Casts which are not injective only bound the cast field inclusively: + // f64 = 51.00000001 satisfies f64 > 51 and casts to 51f. + auto f64_as_f32 = cast(f64, float32()); + Simplify{equal(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater(f64, literal(51.0))) + .ExpectUnchanged(); + Simplify{less_equal(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater(f64, literal(51.0))) + .ExpectUnchanged(); + Simplify{greater(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater(f64, literal(51.0))) + .ExpectUnchanged(); + Simplify{less(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater(f64, literal(51.0))) + .Expect(false); + Simplify{greater_equal(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater(f64, literal(51.0))) + .Expect(true); + Simplify{greater(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater(f64, literal(51.5))) + .Expect(true); + // The bound of the guarantee is rounded through the cast before being compared: + // f64 = 50.99999999 satisfies f64 <= 50.99999999 and casts to 51f. + Simplify{equal(f64_as_f32, literal(51.0F))} + .WithGuarantee(less_equal(f64, literal(50.99999999))) + .ExpectUnchanged(); + Simplify{less_equal(f64_as_f32, literal(51.0F))} + .WithGuarantee(greater_equal(f64, literal(51.00000001))) + .ExpectUnchanged(); + Simplify{equal(f64_as_f32, literal(51.0F))} + .WithGuarantee(less_equal(f64, literal(50.5))) + .Expect(false); + // Small is_in value sets are filtered the same way + auto is_in_f32 = [&](std::string json_array) { + return call( + "is_in", {f64_as_f32}, + SetLookupOptions{ArrayFromJSON(float32(), json_array), SetLookupOptions::MATCH}); + }; + Simplify{is_in_f32("[50, 51]")} + .WithGuarantee(greater(f64, literal(51.0))) + .Expect(is_in_f32("[51]")); + Simplify{is_in_f32("[51]")} + .WithGuarantee(greater(f64, literal(51.0))) + .ExpectUnchanged(); + Simplify{is_in_f32("[51]")} + .WithGuarantee(less_equal(f64, literal(50.99999999))) + .ExpectUnchanged(); + Simplify{is_in_f32("[51]")}.WithGuarantee(greater(f64, literal(51.5))).Expect(false); + + // Nor is int64 to float64: 2^53 + 1 casts to 2^53 (when the cast is allowed to + // truncate; a safe cast fails on values outside [-2^53, 2^53] and so does any + // bound outside that range, which is then ignored) + auto i64_as_f64 = call("cast", {i64}, CastOptions::Unsafe(float64())); + Simplify{equal(i64_as_f64, literal(9007199254740992.0))} + .WithGuarantee(greater(i64, literal(9007199254740992))) + .ExpectUnchanged(); + Simplify{equal(i64_as_f64, literal(9007199254740992.0))} + .WithGuarantee(greater_equal(i64, literal(9007199254740994))) + .Expect(false); + Simplify{equal(cast(i64, float64()), literal(9007199254740992.0))} + .WithGuarantee(greater_equal(i64, literal(9007199254740994))) + .ExpectUnchanged(); + Simplify{greater(cast(i32, float32()), literal(5.0F))} + .WithGuarantee(greater(i32, literal(5))) + .ExpectUnchanged(); + + // Widening casts are injective, so strict bounds stay strict + Simplify{greater(cast(i32, int64()), literal(5))} + .WithGuarantee(greater(i32, literal(5))) + .Expect(true); + Simplify{greater(cast(i32, float64()), literal(5.0))} + .WithGuarantee(greater(i32, literal(5))) + .Expect(true); + Simplify{greater(cast(f32, float64()), literal(5.0))} + .WithGuarantee(greater(f32, literal(5.0F))) + .Expect(true); +} + TEST(Expression, SimplifyIsIn) { auto is_in = [](Expression field, std::shared_ptr value_set_type, std::string json_array, @@ -2573,6 +2699,13 @@ TEST(Expression, SimplifyIsInLargeValueSetNullableRange) { } else { non_nullable.Expect(false); } + // ... as can is_valid elsewhere in the guarantee. + auto valid = Simplify{filter}.WithGuarantee(and_(nullable_range, is_valid(i32))); + if (null_matching == SetLookupOptions::INCONCLUSIVE && value_set_has_null) { + valid.ExpectUnchanged(); + } else { + valid.Expect(false); + } // Verify that simplification also preserves the values on null input. ASSERT_OK_AND_ASSIGN(auto bound_filter, filter.Bind(*kBoringSchema)); diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 71f30b6b5683..f38ea5f036ce 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -826,7 +826,7 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsIn) { {0, 1, 4}); } -TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsInWithNulls) { +TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsWithNulls) { // Each row group contains ten distinct values and a null, so its statistics // yield (i64 >= min and i64 <= max) or is_null(i64). auto dataset_schema = schema({field("i64", int64())}); @@ -852,10 +852,24 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsInWit std::iota(values.begin(), values.begin() + 10, 0); std::iota(values.begin() + 10, values.end(), 20); ASSERT_GT(values.size(), 50); - auto expect_row_groups = [&](bool value_set_has_null, - compute::SetLookupOptions::NullMatchingBehavior behavior, - const std::vector& expected, - std::optional skip_nulls = std::nullopt) { + auto expect_row_groups = [&](compute::Expression filter, + const std::vector& expected) { + 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); + }; + auto expect_is_in = [&](bool value_set_has_null, + compute::SetLookupOptions::NullMatchingBehavior behavior, + const std::vector& expected, + std::optional skip_nulls = std::nullopt) { SCOPED_TRACE(value_set_has_null); SCOPED_TRACE(behavior); auto set_values = values; @@ -868,27 +882,30 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsInWit ArrayFromVector(valid, set_values, &value_set); compute::SetLookupOptions options{std::move(value_set), behavior}; options.skip_nulls = skip_nulls; - SetFilter(call("is_in", {field_ref("i64")}, options)); - 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); + expect_row_groups(call("is_in", {field_ref("i64")}, options), expected); }; - expect_row_groups(false, compute::SetLookupOptions::MATCH, {0, 2}); + expect_is_in(false, compute::SetLookupOptions::MATCH, {0, 2}); // A null in a MATCH set matches the nulls in every row group. - expect_row_groups(true, compute::SetLookupOptions::MATCH, {0, 1, 2}); - expect_row_groups(false, compute::SetLookupOptions::SKIP, {0, 2}); - expect_row_groups(true, compute::SetLookupOptions::SKIP, {0, 2}); + expect_is_in(true, compute::SetLookupOptions::MATCH, {0, 1, 2}); + expect_is_in(false, compute::SetLookupOptions::SKIP, {0, 2}); + expect_is_in(true, compute::SetLookupOptions::SKIP, {0, 2}); // The deprecated flag takes precedence over a conflicting enum value. - expect_row_groups(true, compute::SetLookupOptions::SKIP, {0, 1, 2}, false); - expect_row_groups(true, compute::SetLookupOptions::MATCH, {0, 2}, true); + expect_is_in(true, compute::SetLookupOptions::SKIP, {0, 1, 2}, false); + expect_is_in(true, compute::SetLookupOptions::MATCH, {0, 2}, true); + + // Comparisons and small value sets are pruned against the same statistics. + expect_row_groups(greater(field_ref("i64"), literal(100)), {}); + expect_row_groups(less(field_ref("i64"), literal(10)), {0}); + expect_row_groups(is_null(field_ref("i64")), {0, 1, 2}); + std::shared_ptr small_set; + ArrayFromVector({5, 25}, &small_set); + expect_row_groups( + call("is_in", {field_ref("i64")}, compute::SetLookupOptions{small_set}), {0, 2}); + expect_row_groups( + and_(call("is_in", {field_ref("i64")}, compute::SetLookupOptions{small_set}), + is_valid(field_ref("i64"))), + {0, 2}); } TEST_P(TestParquetFileFormatScan, ExplicitRowGroupSelection) {