From e210bf4177a3519aa1dbb7a2a0809eb877871b8e Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 27 Oct 2026 20:26:21 +0100 Subject: [PATCH 01/12] 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 0b1410af33c6055707cf81189031fc0330a7a839 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 27 Oct 2026 20:29:36 +0100 Subject: [PATCH 02/12] 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 00985711cfb17ae06a12e0f8a47bf1fe6aa1163b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Thu, 29 Oct 2026 19:16:11 +0100 Subject: [PATCH 03/12] 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 79750a758ed12c08d83bb518875699460d9e1af8 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Thu, 29 Oct 2026 20:43:41 +0100 Subject: [PATCH 04/12] 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) { From 9cbe6ae0b9b0b15a3fed2c19c8efa2df0fb71c3d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 3 Nov 2026 19:23:12 +0100 Subject: [PATCH 05/12] GH-55292: [C++] Sort is_in value sets lazily and prune EMIT_NULL and INCONCLUSIVE calls - Defer concatenating, checking and sorting the value set to the first GetSortedValueSet() call (std::call_once), so plain is_in kernel calls pay nothing extra at init. - Simplify an is_in call without possible matches to its exact result for every null matching behavior: invert(true_unless_null(arg)) for EMIT_NULL and INCONCLUSIVE, or a null literal for INCONCLUSIVE with a null in the value set, so these prune row groups too. - Factor out Comparison::IsOrderPreservingCast, derive IsSearchableValueSetType from the accessor trait, and drop unreachable NaN, type and comparison checks. - Test chunked value sets and the new null matching results. --- cpp/src/arrow/compute/expression.cc | 149 ++++++++-------- cpp/src/arrow/compute/expression_internal.h | 56 +++--- cpp/src/arrow/compute/expression_test.cc | 45 +++-- .../compute/kernels/scalar_set_lookup.cc | 21 +-- .../kernels/scalar_set_lookup_internal.h | 168 ++++++++---------- cpp/src/arrow/dataset/file_parquet_test.cc | 5 + 6 files changed, 214 insertions(+), 230 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 731ba0d42ee3..ca9d59de76a2 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1176,8 +1176,7 @@ struct OrderPreservingCasts { static OrderPreservingCasts Strip(const Expression& expr) { OrderPreservingCasts out; const Expression* lhs = &expr; - while (lhs->call() && lhs->call()->function_name == "cast" && - &Comparison::StripOrderPreservingCasts(*lhs) != lhs) { + while (lhs->call() && Comparison::IsOrderPreservingCast(*lhs->call())) { const Expression& from = lhs->call()->arguments[0]; out.injective = out.injective && IsInjective(from.type()->id(), lhs->type()->id()); out.casts.push_back(lhs->call()); @@ -1219,8 +1218,35 @@ struct OrderPreservingCasts { } }; +// An Expression which evaluates to `value` where `target` is valid and to null where +// it is null: true_unless_null(target) or invert(true_unless_null(target)). +// +// true_unless_null is cheap; it purely reuses the validity bitmap for the values +// buffer. Inversion is less cheap but we expect that term never to be evaluated +// since invert(true_unless_null(x)) is not satisfiable. +Result ValueUnlessNull(const Expression& target, bool value) { + ExecContext exec_context; + + Expression::Call call; + call.function_name = "true_unless_null"; + call.arguments = {target}; + ARROW_ASSIGN_OR_RAISE(auto true_unless_null, + BindNonRecursive(std::move(call), + /*insert_implicit_casts=*/false, &exec_context)); + if (value) return true_unless_null; + + Expression::Call invert; + invert.function_name = "invert"; + invert.arguments = {std::move(true_unless_null)}; + return BindNonRecursive(std::move(invert), + /*insert_implicit_casts=*/false, &exec_context); +} + // 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. +// +// `target` and `bound` refer into the guarantee the inequality was extracted from, +// which must outlive it. struct Inequality { // The inequality type Comparison::type cmp; @@ -1303,27 +1329,8 @@ struct Inequality { Result simplified_to(const Expression& bound_target, bool value) const { if (!nullable) return literal(value); - ExecContext exec_context; - // Data may be null, so comparison will yield `value` - or null IFF the data was null - // - // true_unless_null is cheap; it purely reuses the validity bitmap for the values - // buffer. Inversion is less cheap but we expect that term never to be evaluated - // since invert(true_unless_null(x)) is not satisfiable. - Expression::Call call; - call.function_name = "true_unless_null"; - call.arguments = {bound_target}; - ARROW_ASSIGN_OR_RAISE( - auto true_unless_null, - BindNonRecursive(std::move(call), - /*insert_implicit_casts=*/false, &exec_context)); - if (value) return true_unless_null; - - Expression::Call invert; - invert.function_name = "invert"; - invert.arguments = {std::move(true_unless_null)}; - return BindNonRecursive(std::move(invert), - /*insert_implicit_casts=*/false, &exec_context); + return ValueUnlessNull(bound_target, value); } /// Simplify an `is_in` call against an inequality guarantee. @@ -1495,28 +1502,34 @@ 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. +/// \brief What an `is_in` call evaluates to 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) { +/// \param arg the argument of the call +/// \param nullable whether `arg` may be null +/// \return the equivalent Expression, or nullopt if the call may still be true +Result> SimplifyIsInWithoutMatches( + const Expression& arg, const SetLookupOptions& options, bool value_set_has_null, + bool nullable) { switch (options.GetNullMatchingBehavior()) { case SetLookupOptions::MATCH: // A null in the value set matches nulls in the data. - return !(nullable && value_set_has_null); + if (nullable && value_set_has_null) return std::nullopt; + return literal(false); case SetLookupOptions::SKIP: // Nulls in the data produce false. - return true; + return literal(false); + case SetLookupOptions::INCONCLUSIVE: + // Any non-match produces null if the value set contains a null, as do + // nulls in the data. + if (value_set_has_null) return literal(MakeNullScalar(boolean())); + [[fallthrough]]; 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; + if (!nullable) return literal(false); + return ValueUnlessNull(arg, false); } - return false; + return std::nullopt; } /// \brief Simplify an `is_in` call with a large value set against all inequality @@ -1525,7 +1538,9 @@ bool IsInWithoutMatchesIsFalse(const SetLookupOptions& options, bool value_set_h /// 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 +/// is simplified to what it evaluates to without matches (`false`, or null for +/// the rows the null matching behavior makes null); 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 @@ -1552,11 +1567,7 @@ Result SimplifyIsInWithInequalities( 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 (inequality.target == target) applicable.push_back(&inequality); } if (applicable.empty()) return expr; @@ -1605,41 +1616,30 @@ Result SimplifyIsInWithInequalities( } // 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. + // rounds, so no exactness is needed. The last cast is to the argument type, and + // a valid non-NaN bound stays valid and non-NaN. 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; + DCHECK(bound->type->Equals(*sorted->type())); // 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)); - 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; + DCHECK_NE(cmp, Comparison::NA); + + // Unless the comparison admits values below the bound, skip them (and the + // bound itself unless it is admitted); likewise for values above the bound. + const bool admits_bound = cmp & Comparison::EQUAL; + if (!(cmp & Comparison::LESS)) { + ARROW_ASSIGN_OR_RAISE(int64_t index, internal::SearchSortedValueSet( + *sorted, *bound, /*upper=*/!admits_bound)); + begin = std::max(begin, index); + } + if (!(cmp & Comparison::GREATER)) { + ARROW_ASSIGN_OR_RAISE(int64_t index, internal::SearchSortedValueSet( + *sorted, *bound, /*upper=*/admits_bound)); + end = std::min(end, index); } } @@ -1647,10 +1647,10 @@ Result SimplifyIsInWithInequalities( // 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); + ARROW_ASSIGN_OR_RAISE(auto simplified, + SimplifyIsInWithoutMatches(call->arguments[0], *options, + state.value_set_has_null, nullable)); + return simplified ? *std::move(simplified) : expr; } /// \brief Simplify an expression given a guarantee, if the guarantee @@ -1697,7 +1697,8 @@ Result SimplifyWithGuarantee(Expression expr, RETURN_NOT_OK(CanonicalizeAndFoldConstants()); // The inequalities on fields implied by the guarantee. A field which is also - // guaranteed valid cannot be null, whatever its inequalities admit. + // guaranteed valid cannot be null, whatever its inequalities admit. They refer + // into conjunction_members, which must not change from here on. std::vector inequalities; for (const auto& guarantee : conjunction_members) { for (const auto& inequality : Inequality::Extract(guarantee)) { diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index 05b41c8d7f54..826879646f16 100644 --- a/cpp/src/arrow/compute/expression_internal.h +++ b/cpp/src/arrow/compute/expression_internal.h @@ -110,49 +110,41 @@ struct Comparison { return less.scalar_as().value ? LESS : GREATER; } - // Given an Expression wrapped in casts which preserve ordering - // (for example, cast(field_ref("i16"), to_type=int32())), unwrap the inner Expression. - // This is used to destructure implicitly cast field_refs during Expression - // simplification. - static const Expression& StripOrderPreservingCasts(const Expression& expr) { - auto call = expr.call(); - if (!call) return expr; - if (call->function_name != "cast") return expr; - - const Expression& from = call->arguments[0]; + // Whether `call` is a cast which preserves the ordering of its argument. + static bool IsOrderPreservingCast(const Expression::Call& call) { + if (call.function_name != "cast") return false; - auto from_id = from.type()->id(); - auto to_id = expr.type()->id(); + auto from_id = call.arguments[0].type()->id(); + auto to_id = call.type.id(); if (is_floating(to_id)) { - if (is_integer(from_id) || is_floating(from_id)) { - return StripOrderPreservingCasts(from); - } - return expr; + return is_integer(from_id) || is_floating(from_id); } if (is_unsigned_integer(to_id)) { - if (is_unsigned_integer(from_id) && bit_width(to_id) >= bit_width(from_id)) { - return StripOrderPreservingCasts(from); - } - return expr; + return is_unsigned_integer(from_id) && bit_width(to_id) >= bit_width(from_id); } 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; + if (!is_integer(from_id) || bit_width(to_id) < bit_width(from_id)) return false; + // An unsigned to signed cast of the same width wraps negative when overflow + // is allowed, which does not preserve ordering. + return !(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 false; + } + + // Given an Expression wrapped in casts which preserve ordering + // (for example, cast(field_ref("i16"), to_type=int32())), unwrap the inner Expression. + // This is used to destructure implicitly cast field_refs during Expression + // simplification. + static const Expression& StripOrderPreservingCasts(const Expression& expr) { + auto call = expr.call(); + if (!call || !IsOrderPreservingCast(*call)) return expr; + return StripOrderPreservingCasts(call->arguments[0]); } static type GetFlipped(type op) { diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 41e077a70a6f..ca3c8bda12a8 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2455,13 +2455,13 @@ TEST(Expression, SimplifyIsInLargeValueSet) { .WithGuarantee( and_(greater(str, literal("k110")), less_equal(str, literal("k111")))) .ExpectUnchanged(); - // ... as do value sets with nulls, which INCONCLUSIVE cannot simplify + // ... as do value sets with nulls, which make any non-match null for INCONCLUSIVE 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(); + with_null.Expect(null_literal(boolean())); } else { with_null.Expect(false); } @@ -2599,13 +2599,13 @@ TEST(Expression, SimplifyIsInLargeValueSet) { .WithGuarantee(nullable_guarantee) .Expect(false); - // EMIT_NULL: nulls produce null + // EMIT_NULL: nulls produce null, which is still never true Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::EMIT_NULL)} .WithGuarantee(nullable_guarantee) - .ExpectUnchanged(); + .Expect(not_(true_unless_null(i32))); Simplify{is_in(i32, int32(), range, SetLookupOptions::EMIT_NULL)} .WithGuarantee(nullable_guarantee) - .ExpectUnchanged(); + .Expect(not_(true_unless_null(i32))); Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::EMIT_NULL)} .WithGuarantee(non_nullable_guarantee) .Expect(false); @@ -2617,13 +2617,13 @@ TEST(Expression, SimplifyIsInLargeValueSet) { // contains a null Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::INCONCLUSIVE)} .WithGuarantee(nullable_guarantee) - .ExpectUnchanged(); + .Expect(null_literal(boolean())); Simplify{is_in(i32, int32(), range, SetLookupOptions::INCONCLUSIVE)} .WithGuarantee(nullable_guarantee) - .ExpectUnchanged(); + .Expect(not_(true_unless_null(i32))); Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::INCONCLUSIVE)} .WithGuarantee(non_nullable_guarantee) - .ExpectUnchanged(); + .Expect(null_literal(boolean())); Simplify{is_in(i32, int32(), range, SetLookupOptions::INCONCLUSIVE)} .WithGuarantee(non_nullable_guarantee) .Expect(false); @@ -2639,6 +2639,17 @@ TEST(Expression, SimplifyIsInLargeValueSet) { .WithGuarantee(or_(equal(i32, literal(50)), is_null(i32))) .ExpectUnchanged(); + // Chunked value sets are concatenated before they are sorted + SetLookupOptions chunked_options{ + ChunkedArrayFromJSON(int32(), {make_range_json(100, 150), make_range_json(1, 50)}), + SetLookupOptions::MATCH}; + Simplify{call("is_in", {i32}, chunked_options)} + .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) + .Expect(false); + Simplify{call("is_in", {i32}, chunked_options)} + .WithGuarantee(and_(greater_equal(i32, literal(50)), less(i32, literal(100)))) + .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)))} @@ -2668,11 +2679,17 @@ TEST(Expression, SimplifyIsInLargeValueSetNullableRange) { 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 { + // No member of the value set can occur, so the call is never true unless a + // null in a MATCH value set matches the nulls. + if (null_matching == SetLookupOptions::MATCH && value_set_has_null) { simplified.ExpectUnchanged(); + } else if (null_matching == SetLookupOptions::INCONCLUSIVE && value_set_has_null) { + simplified.Expect(null_literal(boolean())); + } else if (null_matching == SetLookupOptions::EMIT_NULL || + null_matching == SetLookupOptions::INCONCLUSIVE) { + simplified.Expect(not_(true_unless_null(i32))); + } else { + simplified.Expect(false); } // A matching non-null member prevents pruning regardless of null behavior. @@ -2695,14 +2712,14 @@ TEST(Expression, SimplifyIsInLargeValueSetNullableRange) { 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(); + non_nullable.Expect(null_literal(boolean())); } 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(); + valid.Expect(null_literal(boolean())); } else { valid.Expect(false); } diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc index 676661b02877..8665251e1641 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc @@ -16,7 +16,6 @@ // 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" @@ -168,7 +167,8 @@ struct InitStateVisitor { TypeHolder arg_type; std::unique_ptr result; - // Whether to populate SetLookupStateBase::sorted_value_set (is_in only). + // Whether to make the value set available to + // SetLookupStateBase::GetSortedValueSet (is_in only). bool compute_sorted_value_set; InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args, @@ -260,24 +260,11 @@ struct InitStateVisitor { 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))); - } + checked_cast(result.get()) + ->SetSearchableValueSet(options.value_set, ctx->exec_context()->memory_pool()); } 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, 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 60011e4a6760..ecfd9dc0aff8 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -21,7 +21,6 @@ // simplification (libarrow). Everything here must stay header-only. #include -#include #include #include #include @@ -40,6 +39,7 @@ #include "arrow/compute/api_vector.h" #include "arrow/compute/kernel.h" #include "arrow/datum.h" +#include "arrow/memory_pool.h" #include "arrow/result.h" #include "arrow/scalar.h" #include "arrow/status.h" @@ -55,19 +55,28 @@ 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 +/// GetSortedValueSet, which only ever simplifies the call to `false`, see /// https://github.com/apache/arrow/issues/46777 constexpr int64_t kIsInSimplificationMaxValueSet = 50; +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); } + /// 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); + switch (value.type->id()) { + case Type::FLOAT: + return IsNaNValue(::arrow::internal::checked_cast(value).value); + case Type::DOUBLE: + return IsNaNValue( + ::arrow::internal::checked_cast(value).value); + default: + return false; } - return false; } inline Result IsSortedValueSet(const Array& values); @@ -82,21 +91,12 @@ struct SetLookupStateBase : public KernelState { /// 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(); + /// Called once by the `is_in` kernel init, before the state is shared. Nothing + /// is computed until the first call of GetSortedValueSet(), so that only + /// expressions which are simplified against guarantees pay for it. + void SetSearchableValueSet(Datum value_set, MemoryPool* pool) { + searchable_value_set_ = std::move(value_set); + searchable_value_set_pool_ = pool; } /// A copy of the value set sorted ascending in the order of the comparison @@ -107,87 +107,58 @@ struct SetLookupStateBase : public KernelState { /// 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); + std::call_once(sorted_value_set_once_, [this] { + if (searchable_value_set_.is_arraylike()) { + // Best effort: a failure only disables the simplification of `is_in` + // against guarantees. + sorted_value_set_ = + SortValueSet(searchable_value_set_, searchable_value_set_pool_) + .ValueOr(nullptr); + searchable_value_set_ = Datum(); } - } + }); return sorted_value_set_; } private: - /// \return `values` sorted with nulls removed, or nullptr if it contains NaN - static Result> SortValueSet(const Array& values) { + /// \return `value_set` sorted with nulls removed, or nullptr if it contains NaN + static Result> SortValueSet(const Datum& value_set, + MemoryPool* pool) { + std::shared_ptr values; + if (value_set.is_array()) { + values = value_set.make_array(); + } else { + ARROW_ASSIGN_OR_RAISE(values, + Concatenate(value_set.chunked_array()->chunks(), pool)); + } + // An already sorted value set is used as is. + if (values->null_count() == 0) { + ARROW_ASSIGN_OR_RAISE(bool sorted, IsSortedValueSet(*values)); + if (sorted) return 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(); + auto maybe_sorted = SortPrimitiveValueSet(*values); + if (!maybe_sorted.status().IsNotImplemented()) return maybe_sorted; + // Only decimal and binary types get here, neither of which has NaN. ARROW_ASSIGN_OR_RAISE(auto indices, - SortIndices(values, ArraySortOptions(SortOrder::Ascending, - NullPlacement::AtEnd))); + 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; + return sorted->Slice(0, sorted->length() - sorted->null_count()); } - mutable std::mutex sorted_value_set_mutex_; - mutable std::atomic sorted_value_set_ready_{false}; - mutable std::shared_ptr unsorted_value_set_; + mutable std::once_flag sorted_value_set_once_; + // Released once sorted_value_set_ is computed. + mutable Datum searchable_value_set_; + MemoryPool* searchable_value_set_pool_ = default_memory_pool(); 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 @@ -246,12 +217,15 @@ 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 IsSearchableTypeVisitor { + bool result = false; + + template + Status Visit(const T&) { + result = is_searchable_value_set_type::value; + return Status::OK(); + } +}; struct IsSortedVisitor { const Array& values; @@ -349,6 +323,14 @@ struct SearchSortedVisitor { } // namespace set_lookup_detail +/// 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) { + set_lookup_detail::IsSearchableTypeVisitor visitor; + ARROW_UNUSED(VisitTypeInline(type, &visitor)); + return visitor.result; +} + /// 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) { diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index f38ea5f036ce..f9e072495bb9 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -890,6 +890,11 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsWithNulls) { 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}); + // Nulls which produce null rather than false are pruned too. + expect_is_in(false, compute::SetLookupOptions::EMIT_NULL, {0, 2}); + expect_is_in(true, compute::SetLookupOptions::EMIT_NULL, {0, 2}); + expect_is_in(false, compute::SetLookupOptions::INCONCLUSIVE, {0, 2}); + expect_is_in(true, compute::SetLookupOptions::INCONCLUSIVE, {0, 2}); // The deprecated flag takes precedence over a conflicting enum value. expect_is_in(true, compute::SetLookupOptions::SKIP, {0, 1, 2}, false); expect_is_in(true, compute::SetLookupOptions::MATCH, {0, 2}, true); From a7418705f4ea022e102fac4965971049d8dbb175 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 3 Nov 2026 19:45:06 +0100 Subject: [PATCH 06/12] GH-55292: [C++] Reuse sort_indices and search_sorted for is_in pruning and condense tests - Sort the value set with sort_indices + take and search it with search_sorted instead of hand-written visitors. Decimal and fixed size binary value sets are no longer searchable since search_sorted does not support them. - Make the large is_in simplification tests table-driven, merge the null matching tests, and share a helper which checks that simplified filters evaluate like the originals. - Merge the two Parquet row group pruning tests into one file whose row groups contain nulls. --- cpp/src/arrow/compute/expression.cc | 30 +- cpp/src/arrow/compute/expression_test.cc | 904 +++++++----------- .../kernels/scalar_set_lookup_internal.h | 270 +----- cpp/src/arrow/dataset/file_parquet_test.cc | 172 ++-- 4 files changed, 431 insertions(+), 945 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index ca9d59de76a2..62236b219f69 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1535,13 +1535,13 @@ Result> SimplifyIsInWithoutMatches( /// \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 +/// Rather than filtering the value set (see Inequality::SimplifyIsIn), this /// 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 what it evaluates to without matches (`false`, or null for -/// the rows the null matching behavior makes null); otherwise it is left -/// unchanged. This is exact for any -/// number of values and costs O(log N) per guarantee, with no rebinding. +/// range of values admitted by the inequalities (`search_sorted`). If that range is +/// empty the call is simplified to what it evaluates to without matches (`false`, +/// or null for the rows the null matching behavior makes null); 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 @@ -1628,17 +1628,27 @@ Result SimplifyIsInWithInequalities( Comparison::type cmp = casts.ApplyToComparison(inequality->cmp); DCHECK_NE(cmp, Comparison::NA); + // The index of the first value not less than the bound (Left) or greater than + // the bound (Right). + auto search = [&](SearchSortedOptions::Side side) -> Result { + ARROW_ASSIGN_OR_RAISE( + Datum index, + SearchSorted(sorted, bound, SearchSortedOptions(side), &exec_context)); + return static_cast(index.scalar_as().value); + }; // Unless the comparison admits values below the bound, skip them (and the // bound itself unless it is admitted); likewise for values above the bound. const bool admits_bound = cmp & Comparison::EQUAL; if (!(cmp & Comparison::LESS)) { - ARROW_ASSIGN_OR_RAISE(int64_t index, internal::SearchSortedValueSet( - *sorted, *bound, /*upper=*/!admits_bound)); + ARROW_ASSIGN_OR_RAISE( + int64_t index, + search(admits_bound ? SearchSortedOptions::Left : SearchSortedOptions::Right)); begin = std::max(begin, index); } if (!(cmp & Comparison::GREATER)) { - ARROW_ASSIGN_OR_RAISE(int64_t index, internal::SearchSortedValueSet( - *sorted, *bound, /*upper=*/admits_bound)); + ARROW_ASSIGN_OR_RAISE( + int64_t index, + search(admits_bound ? SearchSortedOptions::Right : SearchSortedOptions::Left)); end = std::min(end, index); } } diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index ca3c8bda12a8..f6eeb3a65e3b 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -1707,6 +1707,24 @@ struct Simplify { Expectable WithGuarantee(Expression guarantee) { return {expr, guarantee}; } }; +// Rows satisfying `guarantee` (as JSON) must evaluate identically under `filter` and +// under `filter` simplified against the guarantee. +void ExpectSimplifiedEvaluatesEqually(Expression filter, const Expression& guarantee, + const std::string& rows) { + ASSERT_OK_AND_ASSIGN(filter, filter.Bind(*kBoringSchema)); + ASSERT_OK_AND_ASSIGN(auto simplified, SimplifyWithGuarantee(filter, guarantee)); + auto input = RecordBatchFromJSON(kBoringSchema, rows); + Datum evaluated, simplified_evaluated; + ExpectExecute(filter, input, &evaluated); + ExpectExecute(simplified, input, &simplified_evaluated); + if (simplified_evaluated.is_scalar()) { + ASSERT_OK_AND_ASSIGN( + simplified_evaluated, + MakeArrayFromScalar(*simplified_evaluated.scalar(), input->num_rows())); + } + AssertDatumsEqual(evaluated, simplified_evaluated, /*verbose=*/true); +} + TEST(Expression, SingleComparisonGuarantees) { auto i32 = field_ref("i32"); @@ -2068,83 +2086,67 @@ TEST(Expression, SimplifyComparisonThroughNonInjectiveCast) { 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 i64_as_f64 = call("cast", {i64}, CastOptions::Unsafe(float64())); 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); + struct Case { + Expression filter; + Expression guarantee; + std::optional expected; // nullopt: unchanged + }; + for (const auto& c : std::vector{ + // Casts which are not injective only bound the cast field inclusively: + // f64 = 51.00000001 satisfies f64 > 51 and casts to 51f. + {equal(f64_as_f32, literal(51.0F)), greater(f64, literal(51.0)), {}}, + {less_equal(f64_as_f32, literal(51.0F)), greater(f64, literal(51.0)), {}}, + {greater(f64_as_f32, literal(51.0F)), greater(f64, literal(51.0)), {}}, + {less(f64_as_f32, literal(51.0F)), greater(f64, literal(51.0)), + literal(false)}, + {greater_equal(f64_as_f32, literal(51.0F)), greater(f64, literal(51.0)), + literal(true)}, + {greater(f64_as_f32, literal(51.0F)), greater(f64, literal(51.5)), + literal(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. + {equal(f64_as_f32, literal(51.0F)), less_equal(f64, literal(50.99999999)), {}}, + {less_equal(f64_as_f32, literal(51.0F)), + greater_equal(f64, literal(51.00000001)), + {}}, + {equal(f64_as_f32, literal(51.0F)), less_equal(f64, literal(50.5)), + literal(false)}, + // Small is_in value sets are filtered the same way + {is_in_f32("[50, 51]"), greater(f64, literal(51.0)), is_in_f32("[51]")}, + {is_in_f32("[51]"), greater(f64, literal(51.0)), {}}, + {is_in_f32("[51]"), less_equal(f64, literal(50.99999999)), {}}, + {is_in_f32("[51]"), greater(f64, literal(51.5)), literal(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) + {equal(i64_as_f64, literal(9007199254740992.0)), + greater(i64, literal(9007199254740992)), + {}}, + {equal(i64_as_f64, literal(9007199254740992.0)), + greater_equal(i64, literal(9007199254740994)), literal(false)}, + {equal(cast(i64, float64()), literal(9007199254740992.0)), + greater_equal(i64, literal(9007199254740994)), + {}}, + {greater(cast(i32, float32()), literal(5.0F)), greater(i32, literal(5)), {}}, + // Widening casts are injective, so strict bounds stay strict + {greater(cast(i32, int64()), literal(5)), greater(i32, literal(5)), + literal(true)}, + {greater(cast(i32, float64()), literal(5.0)), greater(i32, literal(5)), + literal(true)}, + {greater(cast(f32, float64()), literal(5.0)), greater(f32, literal(5.0F)), + literal(true)}, + }) { + auto expect = Simplify{c.filter}.WithGuarantee(c.guarantee); + c.expected ? expect.Expect(*c.expected) : expect.ExpectUnchanged(); + } } TEST(Expression, SimplifyIsIn) { @@ -2296,459 +2298,261 @@ TEST(Expression, SimplifyIsIn) { } 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); - }; - + // Value sets longer than kIsInSimplificationMaxValueSet are 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 u32 = field_ref("u32"); auto i32 = field_ref("i32"); auto i64 = field_ref("i64"); + auto f32 = field_ref("f32"); auto f64 = field_ref("f64"); auto str = field_ref("str"); + auto binary = field_ref("binary"); - // 100 contiguous values + // 100 contiguous values, and 101 values with a gap: [1, 50] and [100, 150] 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)); + auto in_gap = and_(greater_equal(i32, literal(51)), less_equal(i32, literal(99))); + // 100 to 1, each twice + std::string unsorted = "["; + for (int i = 100; i >= 1; --i) { + unsorted += std::to_string(i) + "," + std::to_string(i) + ","; + } + unsorted.back() = ']'; + // "k100" to "k159" + std::string strings = "["; + for (int i = 0; i < 60; ++i) { + strings += (i > 0 ? ",\"k" : "\"k") + std::to_string(100 + i) + "\""; + } + strings += "]"; + auto f32_roundtrip = cast(cast(f64, float32()), float64()); + auto roundtrip_set = concat_json("[0.10000000149011612]", make_range_json(1, 50)); + auto binary_literal = [](const std::string& value) { + return literal(std::make_shared(Buffer::FromString(value))); + }; - for (SetLookupOptions::NullMatchingBehavior null_matching : { - SetLookupOptions::MATCH, - SetLookupOptions::SKIP, - SetLookupOptions::EMIT_NULL, - SetLookupOptions::INCONCLUSIVE, - }) { - // Guarantees overlapping the value set leave the call unchanged - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee(greater(u32, literal(3))) - .ExpectUnchanged(); - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee(less_equal(u32, literal(1))) - .ExpectUnchanged(); - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee( - and_(greater_equal(u32, literal(100)), less_equal(u32, literal(200)))) - .ExpectUnchanged(); - - // Guarantees disjoint from the value set's hull - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee(greater(u32, literal(100))) - .Expect(false); - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee(greater_equal(u32, literal(101))) - .Expect(false); - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee(less(u32, literal(1))) - .Expect(false); - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee(less_equal(u32, literal(0))) - .Expect(false); - Simplify{is_in(u32, int64(), range, null_matching)} - .WithGuarantee( - and_(greater_equal(u32, literal(101)), less_equal(u32, literal(200)))) - .Expect(false); - - // Guarantees inside the value set's hull but disjoint from its values. This - // requires considering both bounds of the guarantee at once. - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee(and_(greater(u32, literal(50)), less(u32, literal(100)))) - .Expect(false); - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee( - and_(greater_equal(u32, literal(51)), less_equal(u32, literal(99)))) - .Expect(false); - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee( - and_(less_equal(u32, literal(99)), greater_equal(u32, literal(51)))) - .Expect(false); - // ... but a single value of the set in the range is enough to keep the call - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee( - and_(greater_equal(u32, literal(50)), less_equal(u32, literal(99)))) - .ExpectUnchanged(); - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee( - and_(greater_equal(u32, literal(51)), less_equal(u32, literal(100)))) - .ExpectUnchanged(); - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee(greater(u32, literal(50))) - .ExpectUnchanged(); - // Redundant bounds on the same field are all applied - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee(and_(and_(greater(u32, literal(10)), greater(u32, literal(50))), - less(u32, literal(100)))) - .Expect(false); - - // Guarantees on other fields are irrelevant - Simplify{is_in(u32, int64(), gapped, null_matching)} - .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) - .ExpectUnchanged(); - - // Casts: the value set is cast to the field type when the call is bound and - // the bounds of the guarantee are cast to the value set type. - Simplify{is_in(i64, int32(), range, null_matching)} - .WithGuarantee(greater(i64, literal(100))) - .Expect(false); - Simplify{is_in(i64, int32(), range, null_matching)} - .WithGuarantee(greater(i64, literal(100))) - .Expect(false); - Simplify{is_in(cast(i32, int64()), int64(), range, null_matching)} - .WithGuarantee(greater(i32, literal(100))) - .Expect(false); - - // Unsorted value sets with duplicates are sorted when the call is bound - std::string unsorted = "["; - for (int i = 100; i >= 1; --i) { - unsorted += std::to_string(i) + "," + std::to_string(i) + ","; - } - unsorted.back() = ']'; - Simplify{is_in(u32, int64(), unsorted, null_matching)} - .WithGuarantee(greater(u32, literal(100))) - .Expect(false); - Simplify{is_in(u32, int64(), unsorted, null_matching)} - .WithGuarantee(and_(greater(u32, literal(50)), less(u32, literal(51)))) - .Expect(false); - Simplify{is_in(u32, int64(), unsorted, null_matching)} - .WithGuarantee(and_(greater_equal(u32, literal(50)), less(u32, literal(51)))) - .ExpectUnchanged(); - - // Strings - std::string strings = "["; - for (int i = 0; i < 60; ++i) { - if (i > 0) strings += ","; - strings += "\"k" + std::to_string(100 + i) + "\""; - } - strings += "]"; - Simplify{is_in(str, utf8(), strings, null_matching)} - .WithGuarantee(greater(str, literal("k159"))) - .Expect(false); - Simplify{is_in(str, utf8(), strings, null_matching)} - .WithGuarantee(less(str, literal("k100"))) - .Expect(false); - Simplify{is_in(str, utf8(), strings, null_matching)} - .WithGuarantee(and_(greater(str, literal("k110")), less(str, literal("k111")))) - .Expect(false); - Simplify{is_in(str, utf8(), strings, null_matching)} - .WithGuarantee( - and_(greater_equal(str, literal("k110")), less(str, literal("k111")))) - .ExpectUnchanged(); - Simplify{is_in(str, utf8(), strings, null_matching)} - .WithGuarantee(greater(str, literal("k1"))) - .ExpectUnchanged(); - // Unsorted strings with duplicates take the generic sorting path - std::string unsorted_strings = "["; - for (int i = 59; i >= 0; --i) { - unsorted_strings += "\"k" + std::to_string(100 + i) + "\","; - unsorted_strings += "\"k" + std::to_string(100 + i) + "\","; - } - unsorted_strings.back() = ']'; - Simplify{is_in(str, utf8(), unsorted_strings, null_matching)} - .WithGuarantee(and_(greater(str, literal("k110")), less(str, literal("k111")))) - .Expect(false); - Simplify{is_in(str, utf8(), unsorted_strings, null_matching)} - .WithGuarantee( - and_(greater(str, literal("k110")), less_equal(str, literal("k111")))) - .ExpectUnchanged(); - // ... as do value sets with nulls, which make any non-match null for INCONCLUSIVE - 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.Expect(null_literal(boolean())); - } else { - with_null.Expect(false); + enum Expected { kUnchanged, kFalse }; + struct Case { + Expression arg; + std::shared_ptr value_set_type; + std::string value_set; + Expression guarantee; + Expected expected; + }; + const std::vector cases = { + // Guarantees overlapping the value set + {u32, int64(), range, greater(u32, literal(3)), kUnchanged}, + {u32, int64(), range, less_equal(u32, literal(1)), kUnchanged}, + {u32, int64(), range, + and_(greater_equal(u32, literal(100)), less_equal(u32, literal(200))), kUnchanged}, + // Guarantees disjoint from the value set's hull + {u32, int64(), range, greater(u32, literal(100)), kFalse}, + {u32, int64(), range, greater_equal(u32, literal(101)), kFalse}, + {u32, int64(), range, less(u32, literal(1)), kFalse}, + {u32, int64(), range, less_equal(u32, literal(0)), kFalse}, + {u32, int64(), range, + and_(greater_equal(u32, literal(101)), less_equal(u32, literal(200))), kFalse}, + // Guarantees inside the value set's hull but disjoint from its values. This + // requires considering both bounds of the guarantee at once. + {i32, int32(), gapped, and_(greater(i32, literal(50)), less(i32, literal(100))), + kFalse}, + {i32, int32(), gapped, in_gap, kFalse}, + {i32, int32(), gapped, + and_(less_equal(i32, literal(99)), greater_equal(i32, literal(51))), kFalse}, + // ... but a single value of the set in the range is enough to keep the call + {i32, int32(), gapped, + and_(greater_equal(i32, literal(50)), less_equal(i32, literal(99))), kUnchanged}, + {i32, int32(), gapped, + and_(greater_equal(i32, literal(51)), less_equal(i32, literal(100))), kUnchanged}, + {i32, int32(), gapped, greater(i32, literal(50)), kUnchanged}, + {i32, int32(), gapped, + or_(and_(greater_equal(i32, literal(50)), less_equal(i32, literal(99))), + is_null(i32)), + kUnchanged}, + {i32, int32(), gapped, or_(equal(i32, literal(50)), is_null(i32)), kUnchanged}, + // Redundant bounds on the same field are all applied + {i32, int32(), gapped, + and_(and_(greater(i32, literal(10)), greater(i32, literal(50))), + less(i32, literal(100))), + kFalse}, + // Bounds on other fields are irrelevant, including those of a range which may + // be null for another field + {i32, int32(), gapped, and_(greater(u32, literal(50)), less(u32, literal(100))), + kUnchanged}, + {i32, int32(), gapped, or_(in_gap, is_null(u32)), kUnchanged}, + {i32, int32(), gapped, + or_(and_(in_gap, greater_equal(u32, literal(51))), is_null(u32)), kUnchanged}, + // 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 + {i64, int32(), range, greater(i64, literal(100)), kFalse}, + {i64, int32(), range, greater(i64, literal(100)), kFalse}, + {cast(i32, int64()), int64(), range, greater(i32, literal(100)), kFalse}, + // Unsorted value sets with duplicates are sorted + {u32, int64(), unsorted, greater(u32, literal(100)), kFalse}, + {u32, int64(), unsorted, and_(greater(u32, literal(50)), less(u32, literal(51))), + kFalse}, + {u32, int64(), unsorted, + and_(greater_equal(u32, literal(50)), less(u32, literal(51))), kUnchanged}, + // Strings + {str, utf8(), strings, greater(str, literal("k159")), kFalse}, + {str, utf8(), strings, less(str, literal("k100")), kFalse}, + {str, utf8(), strings, + and_(greater(str, literal("k110")), less(str, literal("k111"))), kFalse}, + {str, utf8(), strings, + and_(greater_equal(str, literal("k110")), less(str, literal("k111"))), kUnchanged}, + {str, utf8(), strings, greater(str, literal("k1")), kUnchanged}, + // Floats + {f64, float64(), range, greater(f64, literal(100.5)), kFalse}, + {f64, float64(), range, and_(greater(f64, literal(50.0)), less(f64, literal(51.0))), + kFalse}, + {f64, float64(), range, + and_(greater(f64, literal(50.0)), less_equal(f64, literal(51.0))), kUnchanged}, + // 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). + {f32, float32(), concat_json("[0.1]", range), + and_(greater(f32, literal(0.1)), less(f32, literal(0.5))), kUnchanged}, + {f32, float32(), concat_json("[0.7]", range), less(f32, literal(0.7)), kUnchanged}, + {f32, float32(), concat_json("[0.5]", range), less(f32, literal(0.5)), kFalse}, + {f32, float32(), concat_json("[0.5]", range), + and_(greater(f32, literal(0.25)), less(f32, literal(0.5))), kFalse}, + // Statistics cannot bound NaN, so a value set containing NaN is never simplified + {f64, float64(), concat_json("[NaN]", range), greater(f64, literal(100.5)), + kUnchanged}, + // 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. + {cast(f64, float32()), float32(), make_range_json(1, 51), + greater(f64, literal(51.0)), kUnchanged}, + {cast(f64, float32()), float32(), make_range_json(1, 51), less(f64, literal(1.0)), + kUnchanged}, + {cast(f64, float32()), float32(), make_range_json(1, 51), + greater(f64, literal(51.5)), kFalse}, + {cast(f64, float32()), float32(), make_range_json(1, 51), + and_(greater(f64, literal(50.0)), less(f64, literal(51.0))), kUnchanged}, + {cast(f64, float32()), float32(), make_range_json(1, 51), + and_(greater(f64, literal(50.25)), less(f64, literal(50.75))), kFalse}, + // Bounds go through the whole cast chain: 0.1 rounds up to 0.1f and back to + // double(0.1f), a set member, so f64 <= 0.1 cannot be ruled out. Strict bounds + // are weakened: f64 = 50.00000001 round trips to 50.0. + {f32_roundtrip, float64(), roundtrip_set, less_equal(f64, literal(0.1)), + kUnchanged}, + {f32_roundtrip, float64(), roundtrip_set, less_equal(f64, literal(0.05)), kFalse}, + {f32_roundtrip, float64(), roundtrip_set, greater(f64, literal(50.0)), kUnchanged}, + {f32_roundtrip, float64(), roundtrip_set, greater(f64, literal(50.5)), kFalse}, + // An unsafe unsigned to signed cast of the same width wraps and is not order + // preserving: u32 = 4294967221 casts to -75, a set member. + {call("cast", {u32}, CastOptions::Unsafe(int32())), int32(), + make_range_json(-100, -50), greater_equal(u32, literal(1)), kUnchanged}, + {cast(u32, int32()), int32(), make_range_json(-100, -50), + greater_equal(u32, literal(1)), kFalse}, + // 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". + {binary, int64(), make_range_json(10, 60), + and_(greater_equal(binary, binary_literal("10")), + less_equal(binary, binary_literal("2"))), + kUnchanged}, + // Dictionaries are not searchable + {field_ref("dict_i32"), int32(), range, + greater(field_ref("dict_i32"), literal(100)), kUnchanged}, + }; + // The guarantees above rule out nulls or the call is left unchanged, so the null + // matching behavior does not matter (see SimplifyIsInLargeValueSetNullMatching). + for (auto null_matching : + {SetLookupOptions::MATCH, SetLookupOptions::SKIP, SetLookupOptions::EMIT_NULL, + SetLookupOptions::INCONCLUSIVE}) { + SCOPED_TRACE(null_matching); + for (const auto& c : cases) { + SetLookupOptions options{ArrayFromJSON(c.value_set_type, c.value_set), + null_matching}; + auto expect = Simplify{call("is_in", {c.arg}, options)}.WithGuarantee(c.guarantee); + c.expected == kFalse ? expect.Expect(false) : expect.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, which is still never true - Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::EMIT_NULL)} - .WithGuarantee(nullable_guarantee) - .Expect(not_(true_unless_null(i32))); - Simplify{is_in(i32, int32(), range, SetLookupOptions::EMIT_NULL)} - .WithGuarantee(nullable_guarantee) - .Expect(not_(true_unless_null(i32))); - 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) - .Expect(null_literal(boolean())); - Simplify{is_in(i32, int32(), range, SetLookupOptions::INCONCLUSIVE)} - .WithGuarantee(nullable_guarantee) - .Expect(not_(true_unless_null(i32))); - Simplify{is_in(i32, int32(), range_with_null, SetLookupOptions::INCONCLUSIVE)} - .WithGuarantee(non_nullable_guarantee) - .Expect(null_literal(boolean())); - 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(); - // Chunked value sets are concatenated before they are sorted SetLookupOptions chunked_options{ - ChunkedArrayFromJSON(int32(), {make_range_json(100, 150), make_range_json(1, 50)}), - SetLookupOptions::MATCH}; - Simplify{call("is_in", {i32}, chunked_options)} - .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) - .Expect(false); + ChunkedArrayFromJSON(int32(), {make_range_json(100, 150), make_range_json(1, 50)})}; + Simplify{call("is_in", {i32}, chunked_options)}.WithGuarantee(in_gap).Expect(false); Simplify{call("is_in", {i32}, chunked_options)} .WithGuarantee(and_(greater_equal(i32, literal(50)), less(i32, literal(100)))) .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, SimplifyIsInLargeValueSetNullableRange) { + auto is_in_gapped = + call("is_in", {i32}, SetLookupOptions{ArrayFromJSON(int32(), gapped)}); + Simplify{and_(is_in_gapped, greater(u32, literal(5)))}.WithGuarantee(in_gap).Expect( + false); + Simplify{or_(is_in_gapped, greater(u32, literal(5)))}.WithGuarantee(in_gap).Expect( + greater(u32, literal(5))); +} + +TEST(Expression, SimplifyIsInLargeValueSetNullMatching) { + // Without possible matches the call is simplified to what it evaluates to: false, + // or null for the rows which the null matching behavior makes null. The field may + // be null if every inequality on it in the guarantee admits nulls. 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)); + auto in_gap = and_(greater_equal(i32, literal(51)), less_equal(i32, literal(99))); + // The shape of Parquet statistics when a row group may contain nulls + auto nullable_in_gap = or_(in_gap, is_null(i32)); + auto nullable_above = or_(greater(i32, literal(150)), is_null(i32)); + struct Case { + Expression guarantee; + bool nullable; + std::string rows; // rows satisfying the guarantee + }; + const std::vector cases = { + {nullable_in_gap, true, + R"([{"i32": null}, {"i32": 51}, {"i32": 75}, {"i32": 99}])"}, + {nullable_above, true, R"([{"i32": null}, {"i32": 151}])"}, + {or_(equal(i32, literal(75)), is_null(i32)), true, + R"([{"i32": null}, {"i32": 75}])"}, + {greater(i32, literal(150)), false, R"([{"i32": 151}])"}, + // Another inequality or is_valid elsewhere in the guarantee rules out nulls + {and_(nullable_above, less_equal(i32, literal(200))), false, R"([{"i32": 151}])"}, + {and_(nullable_in_gap, greater(i32, literal(0))), false, R"([{"i32": 51}])"}, + {and_(nullable_in_gap, is_valid(i32)), false, R"([{"i32": 51}])"}, + }; 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); - // No member of the value set can occur, so the call is never true unless a - // null in a MATCH value set matches the nulls. - if (null_matching == SetLookupOptions::MATCH && value_set_has_null) { - simplified.ExpectUnchanged(); - } else if (null_matching == SetLookupOptions::INCONCLUSIVE && value_set_has_null) { - simplified.Expect(null_literal(boolean())); - } else if (null_matching == SetLookupOptions::EMIT_NULL || - null_matching == SetLookupOptions::INCONCLUSIVE) { - simplified.Expect(not_(true_unless_null(i32))); - } else { - simplified.Expect(false); - } - - // 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.Expect(null_literal(boolean())); - } 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.Expect(null_literal(boolean())); - } else { - valid.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())); + for (const auto& c : cases) { + SCOPED_TRACE(null_matching); + SCOPED_TRACE(value_set_has_null); + auto value_set = value_set_has_null ? concat_json(gapped, "[null]") : gapped; + auto filter = + call("is_in", {i32}, + SetLookupOptions{ArrayFromJSON(int32(), value_set), null_matching}); + auto expect = Simplify{filter}.WithGuarantee(c.guarantee); + if (null_matching == SetLookupOptions::MATCH && value_set_has_null && + c.nullable) { + // A null in the value set matches the nulls in the data + expect.ExpectUnchanged(); + } else if (null_matching == SetLookupOptions::INCONCLUSIVE && + value_set_has_null) { + // Any non-match is null + expect.Expect(null_literal(boolean())); + } else if (c.nullable && (null_matching == SetLookupOptions::EMIT_NULL || + null_matching == SetLookupOptions::INCONCLUSIVE)) { + // Nulls in the data produce null + expect.Expect(not_(true_unless_null(i32))); + } else { + expect.Expect(false); + } + ExpectSimplifiedEvaluatesEqually(filter, c.guarantee, c.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. + auto i32 = field_ref("i32"); + auto rows = R"([{"i32": null}, {"i32": 1001}, {"i32": 1002}])"; for (int num_values : {10, 100}) { for (auto behavior : {SetLookupOptions::MATCH, SetLookupOptions::SKIP, SetLookupOptions::EMIT_NULL, SetLookupOptions::INCONCLUSIVE}) { @@ -2760,33 +2564,18 @@ TEST(Expression, SimplifyIsInDeprecatedSkipNulls) { 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)); + auto filter = call("is_in", {i32}, options); 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); + ExpectSimplifiedEvaluatesEqually( + filter, or_(greater(i32, literal(bound)), is_null(i32)), rows); } + // Without skip_nulls the null in the value set matches the null row + ASSERT_OK_AND_ASSIGN(auto bound_filter, filter.Bind(*kBoringSchema)); + ASSERT_OK_AND_ASSIGN( + auto simplified, + SimplifyWithGuarantee(bound_filter, + or_(greater(i32, literal(1000)), is_null(i32)))); + EXPECT_EQ(simplified.IsSatisfiable(), !skip_nulls); } } } @@ -2820,78 +2609,43 @@ TEST(Expression, SimplifyThenExecute) { } 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)); + // Rows satisfying the guarantee must evaluate identically under the original and + // the simplified filter, in particular through casts which round or wrap. + auto f64 = field_ref("f64"); + auto f64_as_f32 = + call("is_in", {cast(f64, float32())}, + SetLookupOptions{ArrayFromJSON(float32(), make_range_json(1, 51))}); + auto f32_roundtrip = + call("is_in", {cast(cast(f64, float32()), float64())}, + SetLookupOptions{ArrayFromJSON( + float64(), concat_json("[0.10000000149011612]", make_range_json(1, 50)))}); + auto unsafe_u32_as_i32 = + call("is_in", {call("cast", {field_ref("u32")}, CastOptions::Unsafe(int32()))}, + SetLookupOptions{ArrayFromJSON(int32(), make_range_json(-100, -50))}); struct Case { + Expression filter; Expression guarantee; - std::string input; // rows satisfying the guarantee + std::string rows; // rows satisfying the guarantee }; for ( - const Case& c : std::vector{ - {greater(field_ref("f64"), literal(51.0)), + const auto& c : std::vector{ + {f64_as_f32, greater(f64, literal(51.0)), R"([{"f64": 51.00000001}, {"f64": 51.0000001}, {"f64": 51.5}, {"f64": 1e9}])"}, - {greater_equal(field_ref("f64"), literal(51.0)), + {f64_as_f32, greater_equal(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))), + {f64_as_f32, and_(greater(f64, literal(51.0)), less(f64, literal(52.0))), R"([{"f64": 51.00000001}, {"f64": 51.5}, {"f64": 51.9}])"}, - {less(field_ref("f64"), literal(1.0)), + {f64_as_f32, less(f64, literal(1.0)), R"([{"f64": 0.99999999}, {"f64": 0.5}, {"f64": -1e9}])"}, - {greater(field_ref("f64"), literal(51.5)), + {f64_as_f32, greater(f64, literal(51.5)), R"([{"f64": 51.6}, {"f64": 52.0}, {"f64": 1e9}])"}, + {f32_roundtrip, less_equal(f64, literal(0.1)), + R"([{"f64": 0.09}, {"f64": 0.1}])"}, + {unsafe_u32_as_i32, greater_equal(field_ref("u32"), literal(1)), + R"([{"u32": 1}, {"u32": 4294967221}])"}, }) { - 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); + ExpectSimplifiedEvaluatesEqually(c.filter, c.guarantee, c.rows); } } 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 ecfd9dc0aff8..cd4344687cbd 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -20,67 +20,52 @@ // 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 "arrow/array/array_base.h" -#include "arrow/array/array_binary.h" -#include "arrow/array/array_decimal.h" -#include "arrow/array/array_primitive.h" #include "arrow/array/concatenate.h" -#include "arrow/array/data.h" -#include "arrow/array/util.h" -#include "arrow/buffer.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernel.h" #include "arrow/datum.h" #include "arrow/memory_pool.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:: -/// GetSortedValueSet, which only ever simplifies the call to `false`, see -/// https://github.com/apache/arrow/issues/46777 +/// Longer value sets are instead searched with `search_sorted` in +/// SetLookupStateBase::GetSortedValueSet, which only ever simplifies the call to a +/// result without matches, see https://github.com/apache/arrow/issues/46777 constexpr int64_t kIsInSimplificationMaxValueSet = 50; -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); } - /// Whether a (valid) scalar of a floating point type is NaN. inline bool IsNaNScalar(const Scalar& value) { switch (value.type->id()) { case Type::FLOAT: - return IsNaNValue(::arrow::internal::checked_cast(value).value); + return std::isnan(::arrow::internal::checked_cast(value).value); case Type::DOUBLE: - return IsNaNValue( + return std::isnan( ::arrow::internal::checked_cast(value).value); default: return false; } } -inline Result IsSortedValueSet(const Array& values); -inline Result> SortPrimitiveValueSet(const Array& values); +/// Whether value sets of this type are sorted by `sort_indices` and searched by +/// `search_sorted` in the order of the comparison kernels. +inline bool IsSearchableValueSetType(const DataType& type) { + const auto id = type.id(); + return is_integer(id) || id == Type::FLOAT || id == Type::DOUBLE || is_temporal(id) || + id == Type::DURATION || is_base_binary_like(id); +} /// Common state for the is_in and index_in kernels. struct SetLookupStateBase : public KernelState { @@ -99,8 +84,8 @@ struct SetLookupStateBase : public KernelState { searchable_value_set_pool_ = pool; } - /// A copy of the value set sorted ascending in the order of the comparison - /// kernels, with nulls removed, or nullptr if there is none. + /// A copy of the value set sorted ascending, 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() @@ -131,25 +116,20 @@ struct SetLookupStateBase : public KernelState { ARROW_ASSIGN_OR_RAISE(values, Concatenate(value_set.chunked_array()->chunks(), pool)); } - // An already sorted value set is used as is. - if (values->null_count() == 0) { - ARROW_ASSIGN_OR_RAISE(bool sorted, IsSortedValueSet(*values)); - if (sorted) return 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.status().IsNotImplemented()) return maybe_sorted; - - // Only decimal and binary types get here, neither of which has NaN. 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(); - return sorted->Slice(0, sorted->length() - sorted->null_count()); + 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) { + ARROW_ASSIGN_OR_RAISE(auto last, sorted->GetScalar(sorted->length() - 1)); + if (IsNaNScalar(*last)) return nullptr; + } + return sorted; } mutable std::once_flag sorted_value_set_once_; @@ -159,210 +139,4 @@ struct SetLookupStateBase : public KernelState { mutable std::shared_ptr sorted_value_set_; }; -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 {}; - -struct IsSearchableTypeVisitor { - bool result = false; - - template - Status Visit(const T&) { - result = is_searchable_value_set_type::value; - return Status::OK(); - } -}; - -struct IsSortedVisitor { - const Array& values; - bool result = false; - - template - enable_if_t::value, Status> Visit(const T&) { - using Accessor = ValueSetAccessor; - if (values.length() == 0) { - result = true; - return Status::OK(); - } - auto previous = Accessor::Get(values, 0); - if (IsNaNValue(previous)) return Status::OK(); - for (int64_t i = 1; i < values.length(); ++i) { - auto current = Accessor::Get(values, i); - if (IsNaNValue(current) || current < previous) return Status::OK(); - previous = current; - } - result = true; - return Status::OK(); - } - - Status Visit(const DataType& type) { - return Status::NotImplemented("Searching value sets of type ", type); - } -}; - -struct SortPrimitiveVisitor { - const Array& values; - std::shared_ptr result; - - template - enable_if_t::value && is_searchable_value_set_type::value, Status> - Visit(const T&) { - using c_type = typename T::c_type; - const auto* data = values.data()->template GetValues(1); - std::vector sorted; - sorted.reserve(values.length() - values.null_count()); - if (values.null_count() == 0) { - sorted.assign(data, data + values.length()); - } else { - for (int64_t i = 0; i < values.length(); ++i) { - if (values.IsValid(i)) sorted.push_back(data[i]); - } - } - // NaN would break the strict weak ordering required by std::sort. Statistics - // based guarantees cannot bound NaN anyway, so such a value set is not - // searchable (nullptr). - for (const auto& value : sorted) { - if (IsNaNValue(value)) return Status::OK(); - } - std::sort(sorted.begin(), sorted.end()); - const auto length = static_cast(sorted.size()); - result = MakeArray(ArrayData::Make(values.type(), length, - {nullptr, Buffer::FromVector(std::move(sorted))}, - /*null_count=*/0)); - return Status::OK(); - } - - Status Visit(const DataType& type) { - return Status::NotImplemented("Sorting value sets of type ", type); - } -}; - -struct SearchSortedVisitor { - const Array& sorted; - const Scalar& bound; - bool upper; - int64_t result = 0; - - template - enable_if_t::value, Status> Visit(const T&) { - using Accessor = ValueSetAccessor; - const auto target = Accessor::FromScalar(bound); - int64_t lo = 0, hi = sorted.length(); - while (lo < hi) { - const int64_t mid = lo + (hi - lo) / 2; - const auto value = Accessor::Get(sorted, mid); - const bool go_right = upper ? !(target < value) : (value < target); - if (go_right) { - lo = mid + 1; - } else { - hi = mid; - } - } - result = lo; - return Status::OK(); - } - - Status Visit(const DataType& type) { - return Status::NotImplemented("Searching value sets of type ", type); - } -}; - -} // namespace set_lookup_detail - -/// Whether 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) { - set_lookup_detail::IsSearchableTypeVisitor visitor; - ARROW_UNUSED(VisitTypeInline(type, &visitor)); - return visitor.result; -} - -/// Whether `values` (which must contain no nulls) is sorted ascending in the order -/// of the comparison kernels and contains no NaN. -inline Result IsSortedValueSet(const Array& values) { - ARROW_DCHECK_EQ(values.null_count(), 0); - set_lookup_detail::IsSortedVisitor visitor{values, /*result=*/false}; - RETURN_NOT_OK(VisitTypeInline(*values.type(), &visitor)); - return visitor.result; -} - -/// Sort a value set of a primitive type, dropping nulls. -/// -/// \return the sorted value set, or nullptr if it contains NaN -/// \return NotImplemented for non-primitive types -inline Result> SortPrimitiveValueSet(const Array& values) { - set_lookup_detail::SortPrimitiveVisitor visitor{values, /*result=*/nullptr}; - RETURN_NOT_OK(VisitTypeInline(*values.type(), &visitor)); - return visitor.result; -} - -/// Binary search a sorted value set for `bound`. -/// -/// \pre `sorted` satisfies IsSortedValueSet() -/// \pre `bound` is valid, not NaN, and has the same type as `sorted` -/// \return the index of the first element not less than `bound` (`upper == false`), -/// or of the first element greater than `bound` (`upper == true`) -inline Result SearchSortedValueSet(const Array& sorted, const Scalar& bound, - bool upper) { - ARROW_DCHECK(bound.is_valid); - ARROW_DCHECK(sorted.type()->Equals(*bound.type)); - set_lookup_detail::SearchSortedVisitor visitor{sorted, bound, upper, /*result=*/0}; - RETURN_NOT_OK(VisitTypeInline(*sorted.type(), &visitor)); - return visitor.result; -} - } // namespace arrow::compute::internal diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index f9e072495bb9..365d2c1302f6 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -20,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -751,37 +750,27 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragments) { } 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. + // Row group i holds i64 == 10 * i, ..., 10 * i + 9 and a null, so that its + // statistics yield (i64 >= min and i64 <= max) or is_null(i64). 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})); + 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); - 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)); @@ -796,121 +785,80 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsIn) { } EXPECT_EQ(row_groups, expected_row_groups); }; + using compute::SetLookupOptions; + auto make_set = [](std::vector values, bool with_null = false) { + std::vector valid(values.size(), true); + if (with_null) { + values.push_back(0); + valid.push_back(false); + } + std::shared_ptr array; + ArrayFromVector(valid, values, &array); + return array; + }; + auto is_in = [](std::shared_ptr set, + SetLookupOptions::NullMatchingBehavior behavior = + SetLookupOptions::MATCH) { + return call("is_in", {field_ref("i64")}, SetLookupOptions{std::move(set), behavior}); + }; - // 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}); + // 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]. Only the row + // groups holding a value of the set are read. + std::vector values(60); + std::iota(values.begin(), values.begin() + 20, 0); + std::iota(values.begin() + 20, values.end(), 40); + const std::vector matching = {0, 1, 4, 5, 6, 7}; + const auto all = ::arrow::internal::Iota(kNumRowGroups); + expect_row_groups(is_in(make_set(values)), matching); // 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}); + expect_row_groups(is_in(make_set(shuffled)), matching); // Values beyond the data select nothing. - expect_row_groups(is_in(as_int64_array(iota64(1000, 1100))), {}); + std::vector beyond(100); + std::iota(beyond.begin(), beyond.end(), 1000); + expect_row_groups(is_in(make_set(beyond)), {}); // 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}); + expect_row_groups(is_in(int32_set), matching); // Combined with other predicates. expect_row_groups( - and_(is_in(as_int64_array(values)), less(field_ref("i64"), literal(50))), + and_(is_in(make_set(values)), less(field_ref("i64"), literal(50))), {0, 1, 4}); -} -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())}); - 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 = [&](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); + // Only a null in a MATCH value set matches the nulls in every row group; nulls + // which produce null rather than false are pruned too. + for (auto behavior : {SetLookupOptions::MATCH, SetLookupOptions::SKIP, + SetLookupOptions::EMIT_NULL, SetLookupOptions::INCONCLUSIVE}) { 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; - expect_row_groups(call("is_in", {field_ref("i64")}, options), expected); - }; - - expect_is_in(false, compute::SetLookupOptions::MATCH, {0, 2}); - // A null in a MATCH set matches the nulls in every row group. - 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}); - // Nulls which produce null rather than false are pruned too. - expect_is_in(false, compute::SetLookupOptions::EMIT_NULL, {0, 2}); - expect_is_in(true, compute::SetLookupOptions::EMIT_NULL, {0, 2}); - expect_is_in(false, compute::SetLookupOptions::INCONCLUSIVE, {0, 2}); - expect_is_in(true, compute::SetLookupOptions::INCONCLUSIVE, {0, 2}); + expect_row_groups(is_in(make_set(values), behavior), matching); + expect_row_groups(is_in(make_set(values, /*with_null=*/true), behavior), + behavior == SetLookupOptions::MATCH ? all : matching); + } // The deprecated flag takes precedence over a conflicting enum value. - expect_is_in(true, compute::SetLookupOptions::SKIP, {0, 1, 2}, false); - expect_is_in(true, compute::SetLookupOptions::MATCH, {0, 2}, true); + for (bool skip_nulls : {false, true}) { + SetLookupOptions options{ + make_set(values, /*with_null=*/true), + skip_nulls ? SetLookupOptions::MATCH : SetLookupOptions::SKIP}; + options.skip_nulls = skip_nulls; + expect_row_groups(call("is_in", {field_ref("i64")}, options), + skip_nulls ? matching : all); + } // 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}); + expect_row_groups(is_null(field_ref("i64")), all); + auto small_set = ArrayFromJSON(int64(), "[5, 25]"); + expect_row_groups(is_in(small_set), {0, 2}); + expect_row_groups(and_(is_in(small_set), is_valid(field_ref("i64"))), {0, 2}); } TEST_P(TestParquetFileFormatScan, ExplicitRowGroupSelection) { From 625cf4eaf07bdb4e1345ec217efef493156b00e1 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 3 Nov 2026 20:15:07 +0100 Subject: [PATCH 07/12] GH-55292: [C++] Bound value sets longer than 2^17 by their min and max instead of sorting them Sorting a value set costs 1.5-3x as much as building the is_in lookup table, so for 10M shuffled ids the one-time sort outweighed the row groups it pruned and scans got slower than without pruning. Value sets longer than kIsInMaxSortedValueSet (2^17) are now only bounded by their smallest and largest value (min_max), which is cheap at any size and prunes as well when the values are clustered. --- cpp/src/arrow/compute/expression.cc | 16 +++-- cpp/src/arrow/compute/expression_test.cc | 13 +++++ .../kernels/scalar_set_lookup_internal.h | 58 ++++++++++++++++--- 3 files changed, 74 insertions(+), 13 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 62236b219f69..282b02dacfbd 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1540,8 +1540,9 @@ Result> SimplifyIsInWithoutMatches( /// range of values admitted by the inequalities (`search_sorted`). If that range is /// empty the call is simplified to what it evaluates to without matches (`false`, /// or null for the rows the null matching behavior makes null); otherwise it is -/// left unchanged. This is exact for any number of values and costs O(log N) per -/// guarantee, with no rebinding. +/// left unchanged. This costs O(log N) per guarantee, with no rebinding, and is exact +/// for value sets of up to kIsInMaxSortedValueSet values. Longer value sets are only +/// bounded by their smallest and largest value. /// /// 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 @@ -1579,7 +1580,8 @@ Result SimplifyIsInWithInequalities( 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(); + const auto sorted_value_set = state.GetSortedValueSet(); + const std::shared_ptr& sorted = sorted_value_set.values; if (!sorted) return expr; DCHECK(sorted->type()->Equals(*call->arguments[0].type())); @@ -1654,8 +1656,12 @@ Result SimplifyIsInWithInequalities( } // 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; + // than rebinding it against a narrowed value set. Of a value set known only by its + // smallest and largest value, a member may occur unless all admitted values lie + // below or above both. + const bool may_match = + sorted_value_set.exact ? begin < end : begin < sorted->length() && end > 0; + if (may_match) return expr; ARROW_ASSIGN_OR_RAISE(auto simplified, SimplifyIsInWithoutMatches(call->arguments[0], *options, diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index f6eeb3a65e3b..9bbae21242da 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2470,6 +2470,19 @@ TEST(Expression, SimplifyIsInLargeValueSet) { } } + // Value sets longer than kIsInMaxSortedValueSet are only bounded by their smallest + // and largest value, so a guarantee inside a gap of the value set does not prune + auto huge = call("is_in", {i32}, + SetLookupOptions{ArrayFromJSON( + int32(), concat_json(make_range_json(1, 70000), + make_range_json(100001, 170000)))}); + Simplify{huge}.WithGuarantee(greater(i32, literal(170000))).Expect(false); + Simplify{huge}.WithGuarantee(less(i32, literal(1))).Expect(false); + Simplify{huge} + .WithGuarantee(and_(greater(i32, literal(70000)), less(i32, literal(100001)))) + .ExpectUnchanged(); + Simplify{huge}.WithGuarantee(greater_equal(i32, literal(170000))).ExpectUnchanged(); + // Chunked value sets are concatenated before they are sorted SetLookupOptions chunked_options{ ChunkedArrayFromJSON(int32(), {make_range_json(100, 150), make_range_json(1, 50)})}; 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 cd4344687cbd..5a9bd5fb2c03 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -26,7 +26,10 @@ #include #include "arrow/array/array_base.h" +#include "arrow/array/builder_base.h" #include "arrow/array/concatenate.h" +#include "arrow/compute/api_aggregate.h" +#include "arrow/compute/api_scalar.h" #include "arrow/compute/api_vector.h" #include "arrow/compute/kernel.h" #include "arrow/datum.h" @@ -46,6 +49,12 @@ namespace arrow::compute::internal { /// result without matches, see https://github.com/apache/arrow/issues/46777 constexpr int64_t kIsInSimplificationMaxValueSet = 50; +/// Value sets up to this length are sorted for SetLookupStateBase::GetSortedValueSet. +/// Sorting costs more than building the lookup table of the kernel, so longer value +/// sets are only bounded by their smallest and largest value, which is cheap to find +/// and prunes as well when the values are clustered. +constexpr int64_t kIsInMaxSortedValueSet = 1 << 17; + /// Whether a (valid) scalar of a floating point type is NaN. inline bool IsNaNScalar(const Scalar& value) { switch (value.type->id()) { @@ -84,21 +93,28 @@ struct SetLookupStateBase : public KernelState { searchable_value_set_pool_ = pool; } - /// A copy of the value set sorted ascending, with nulls removed, or nullptr if - /// there is none. + struct SortedValueSet { + /// Ascending values without nulls, or nullptr if the value set cannot be searched + std::shared_ptr values; + /// Whether `values` holds every value of the value set, or only its smallest and + /// largest value (for value sets longer than kIsInMaxSortedValueSet) + bool exact = true; + }; + + /// The value set sorted for searching with `search_sorted`. /// /// 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 { + SortedValueSet GetSortedValueSet() const { std::call_once(sorted_value_set_once_, [this] { if (searchable_value_set_.is_arraylike()) { // Best effort: a failure only disables the simplification of `is_in` // against guarantees. sorted_value_set_ = SortValueSet(searchable_value_set_, searchable_value_set_pool_) - .ValueOr(nullptr); + .ValueOr(SortedValueSet{}); searchable_value_set_ = Datum(); } }); @@ -106,9 +122,7 @@ struct SetLookupStateBase : public KernelState { } private: - /// \return `value_set` sorted with nulls removed, or nullptr if it contains NaN - static Result> SortValueSet(const Datum& value_set, - MemoryPool* pool) { + static Result SortValueSet(const Datum& value_set, MemoryPool* pool) { std::shared_ptr values; if (value_set.is_array()) { values = value_set.make_array(); @@ -116,6 +130,16 @@ struct SetLookupStateBase : public KernelState { ARROW_ASSIGN_OR_RAISE(values, Concatenate(value_set.chunked_array()->chunks(), pool)); } + if (values->length() > kIsInMaxSortedValueSet) { + ARROW_ASSIGN_OR_RAISE(auto bounds, ValueSetBounds(values)); + return SortedValueSet{std::move(bounds), /*exact=*/false}; + } + ARROW_ASSIGN_OR_RAISE(auto sorted, SortValues(values)); + return SortedValueSet{std::move(sorted), /*exact=*/true}; + } + + /// \return `values` sorted with nulls removed, or nullptr if it contains NaN + static Result> SortValues(const std::shared_ptr& values) { ARROW_ASSIGN_OR_RAISE(auto indices, SortIndices(*values, ArraySortOptions(SortOrder::Ascending, NullPlacement::AtEnd))); @@ -132,11 +156,29 @@ struct SetLookupStateBase : public KernelState { return sorted; } + /// \return the smallest and largest value of `values`, or nullptr if it contains NaN + /// or only nulls + static Result> ValueSetBounds( + const std::shared_ptr& values) { + if (is_floating(values->type_id())) { + ARROW_ASSIGN_OR_RAISE(Datum is_nan, IsNan(values)); + ARROW_ASSIGN_OR_RAISE(Datum any_nan, Any(is_nan)); + if (any_nan.scalar_as().value) return nullptr; + } + ARROW_ASSIGN_OR_RAISE(Datum min_max, MinMax(values)); + const auto& bounds = min_max.scalar_as().value; + if (!bounds[0]->is_valid) return nullptr; + ARROW_ASSIGN_OR_RAISE(auto builder, MakeBuilder(values->type())); + RETURN_NOT_OK(builder->AppendScalar(*bounds[0])); + RETURN_NOT_OK(builder->AppendScalar(*bounds[1])); + return builder->Finish(); + } + mutable std::once_flag sorted_value_set_once_; // Released once sorted_value_set_ is computed. mutable Datum searchable_value_set_; MemoryPool* searchable_value_set_pool_ = default_memory_pool(); - mutable std::shared_ptr sorted_value_set_; + mutable SortedValueSet sorted_value_set_; }; } // namespace arrow::compute::internal From e3b1556632e6f2a7138280425154f1fbe0d93578 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 3 Nov 2026 22:30:38 +0100 Subject: [PATCH 08/12] GH-55292: [C++] Build the int32 value set in the Parquet test with a cast MSVC flags constructing a std::vector from int64_t iterators as a narrowing conversion (C4244), which fails the Windows build. --- cpp/src/arrow/dataset/file_parquet_test.cc | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 365d2c1302f6..5877728e551a 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -25,6 +25,7 @@ #include #include "arrow/compute/api_scalar.h" +#include "arrow/compute/cast.h" #include "arrow/dataset/dataset_internal.h" #include "arrow/dataset/parquet_encryption_config.h" #include "arrow/dataset/scanner.h" @@ -823,9 +824,7 @@ TEST_P(TestParquetFileFormatScan, PredicatePushdownRowGroupFragmentsLargeIsIn) { expect_row_groups(is_in(make_set(beyond)), {}); // Values cast from a narrower type are handled too. - std::shared_ptr int32_set; - ArrayFromVector(std::vector(values.begin(), values.end()), - &int32_set); + ASSERT_OK_AND_ASSIGN(auto int32_set, compute::Cast(*make_set(values), int32())); expect_row_groups(is_in(int32_set), matching); // Combined with other predicates. From 1078fe79a67c9d04682569d09ba05c898359624c Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 13:57:28 +0100 Subject: [PATCH 09/12] GH-55292: [C++] Only convert guarantee bounds between types ordered alike and keep NaN in small is_in value sets A guarantee bound of another type was brought into the field's type whenever a safe cast round tripped. Strings and numbers round trip but are ordered differently, so the statistics of a string column read as int64 pruned row groups holding matches of a large is_in. Bounds are now only converted between types which are ordered alike. Value sets of up to 50 values containing NaN are no longer filtered against inequality guarantees: statistics cannot bound NaN, which is_in matches. --- cpp/src/arrow/compute/expression.cc | 37 ++++++++++++++++++++++-- cpp/src/arrow/compute/expression_test.cc | 22 ++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 282b02dacfbd..debb5634301c 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1218,6 +1218,22 @@ struct OrderPreservingCasts { } }; +// Whether values of the two types are ordered alike, so that a guarantee bound of +// one type which casts exactly to the other bounds the same values there. This does +// not hold across families: "10" <= "9" as strings, but both cast exactly to +// integers (and back), for which 10 > 9. +bool HaveSameOrder(const DataType& left, const DataType& right) { + if (left.id() == right.id()) return true; + auto family = [](Type::type id) { + if (is_integer(id) || is_floating(id) || is_decimal(id)) return 1; + if (is_base_binary_like(id)) return 2; + if (id == Type::DATE32 || id == Type::DATE64 || id == Type::TIMESTAMP) return 3; + return 0; + }; + const int left_family = family(left.id()); + return left_family != 0 && left_family == family(right.id()); +} + // An Expression which evaluates to `value` where `target` is valid and to null where // it is null: true_unless_null(target) or invert(true_unless_null(target)). // @@ -1312,11 +1328,14 @@ struct Inequality { /// 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. + /// to a value the expression compares with). Nullopt if the bound cannot be cast, + /// or if it is of a type which is ordered differently from the target's (see + /// HaveSameOrder), in which case casting it would not bound the cast target. std::optional> ThroughCasts( const OrderPreservingCasts& casts) const { Comparison::type cast_cmp = casts.ApplyToComparison(cmp); if (casts.injective) return std::make_pair(cast_cmp, bound); + if (!HaveSameOrder(*bound.type(), *casts.inner->type())) return std::nullopt; ExecContext exec_context; auto cast_bound = casts.ApplyToBound(bound.scalar(), &exec_context); if (!cast_bound.ok()) return std::nullopt; @@ -1367,6 +1386,17 @@ struct Inequality { if (!through_casts) return std::nullopt; const auto& [cmp, bound] = *through_casts; + // Statistics based guarantees cannot bound NaN (a row group's minimum and + // maximum ignore it) whereas `is_in` matches NaN with NaN, so a NaN in the + // value set must not be filtered out. Leave such value sets alone, as + // SimplifyIsInWithInequalities does. + const auto value_set_type_id = options->value_set.type()->id(); + if (value_set_type_id == Type::FLOAT || value_set_type_id == Type::DOUBLE) { + ARROW_ASSIGN_OR_RAISE(Datum is_nan, IsNan(options->value_set)); + ARROW_ASSIGN_OR_RAISE(Datum any_nan, Any(is_nan)); + if (any_nan.scalar_as().value) return std::nullopt; + } + FilterOptions::NullSelectionBehavior null_selection{}; switch (options->GetNullMatchingBehavior()) { case SetLookupOptions::MATCH: @@ -1605,7 +1635,10 @@ Result SimplifyIsInWithInequalities( // 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. + // require the cast to round trip. A bound of a type which is ordered + // differently (a string bound on a field read as an integer, say) says + // nothing about the field's order even if it round trips. + if (!HaveSameOrder(*bound->type, *lhs->type())) continue; auto maybe_cast = Cast(Datum(bound), lhs->type(), CastOptions::Safe(), &exec_context); if (!maybe_cast.ok()) continue; diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 9bbae21242da..9e0f482f692a 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2069,6 +2069,17 @@ TEST(Expression, SimplifyWithNullableRangeGuarantee) { Simplify{is_in("[10, 11]")}.WithGuarantee(nullable_range).Expect(false); Simplify{is_in("[0, 3, 11]")}.WithGuarantee(nullable_range).Expect(is_in("[3]")); + // Statistics cannot bound NaN, which is_in matches, so a small value set + // containing NaN is not filtered either + auto f64 = field_ref("f64"); + auto f64_range = and_(greater_equal(f64, literal(0.0)), less_equal(f64, literal(10.0))); + auto is_in_f64 = [&](std::string json_array) { + return call("is_in", {f64}, SetLookupOptions{ArrayFromJSON(float64(), json_array)}); + }; + Simplify{is_in_f64("[20.0, 30.0]")}.WithGuarantee(f64_range).Expect(false); + Simplify{is_in_f64("[NaN]")}.WithGuarantee(f64_range).ExpectUnchanged(); + Simplify{is_in_f64("[NaN, 5.0, 20.0]")}.WithGuarantee(f64_range).ExpectUnchanged(); + // 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))) @@ -2452,6 +2463,17 @@ TEST(Expression, SimplifyIsInLargeValueSet) { and_(greater_equal(binary, binary_literal("10")), less_equal(binary, binary_literal("2"))), kUnchanged}, + // Likewise for a bound of a type which is ordered differently from the field's + // (the statistics of a string column read as int64, say), even though it casts + // exactly to the field type: "10" <= x <= "9" admits "5". + {i64, int64(), concat_json("[5]", make_range_json(1000, 1059)), + and_(greater_equal(i64, literal("10")), less_equal(i64, literal("9"))), + kUnchanged}, + {str, utf8(), concat_json("[\"10\"]", strings), + and_(greater_equal(str, literal(9)), less_equal(str, literal(10))), kUnchanged}, + // ... whereas bounds of another numeric type are converted + {i64, int64(), concat_json("[5]", make_range_json(1000, 1059)), + and_(greater_equal(i64, literal(9.0)), less_equal(i64, literal(10.0))), kFalse}, // Dictionaries are not searchable {field_ref("dict_i32"), int32(), range, greater(field_ref("dict_i32"), literal(100)), kUnchanged}, From 3d76c38c3b4e299252883046913182cb6926ea3b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 14:29:52 +0100 Subject: [PATCH 10/12] GH-55292: [C++] Let the is_in kernel evaluate itself over a range and simplify against per-field ranges Expression simplification no longer reaches into the is_in kernel state. A kernel state may implement the new internal RangeEvaluator interface to tell what its call evaluates to when the argument lies in a range; is_in implements it with its lazily sorted value set and its null matching rules, and scalar_set_lookup_internal.h is removed. SimplifyWithGuarantee folds the guarantee once into the inequalities and nullability of each field and simplifies every call in a single pass. Guarantee bounds are converted by one rule (Inequality::Restate) for comparisons, small is_in value sets and range evaluation alike: a bound of another type must be ordered like the field's type and convert to it exactly. --- cpp/src/arrow/compute/expression.cc | 411 ++++++++---------- cpp/src/arrow/compute/expression_test.cc | 35 ++ .../compute/kernels/scalar_set_lookup.cc | 247 ++++++++++- .../kernels/scalar_set_lookup_internal.h | 184 -------- .../arrow/compute/range_evaluation_internal.h | 70 +++ 5 files changed, 526 insertions(+), 421 deletions(-) delete mode 100644 cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h create mode 100644 cpp/src/arrow/compute/range_evaluation_internal.h diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index debb5634301c..f72953ad363d 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -20,6 +20,7 @@ #include "arrow/compute/expression.h" #include +#include #include #include #include @@ -30,7 +31,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/range_evaluation_internal.h" #include "arrow/compute/util.h" #include "arrow/io/memory.h" #ifdef ARROW_IPC @@ -1218,6 +1219,25 @@ struct OrderPreservingCasts { } }; +// Value sets up to this length are simplified against a guarantee by filtering the +// value set and rebinding the `is_in` call (see Inequality::SimplifyIsIn). This is +// too expensive for longer value sets (https://github.com/apache/arrow/issues/46777), +// for which the kernel is instead asked what the call evaluates to over the range of +// the guarantee (see internal::RangeEvaluator), which never rebinds. +constexpr int64_t kIsInSimplificationMaxValueSet = 50; + +// Whether a (valid) scalar of a floating point type is NaN. +bool IsNaNScalar(const Scalar& value) { + switch (value.type->id()) { + case Type::FLOAT: + return std::isnan(checked_cast(value).value); + case Type::DOUBLE: + return std::isnan(checked_cast(value).value); + default: + return false; + } +} + // Whether values of the two types are ordered alike, so that a guarantee bound of // one type which casts exactly to the other bounds the same values there. This does // not hold across families: "10" <= "9" as strings, but both cast exactly to @@ -1325,21 +1345,45 @@ 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, - /// or if it is of a type which is ordered differently from the target's (see - /// HaveSameOrder), in which case casting it would not bound the cast target. - std::optional> ThroughCasts( + /// This inequality restated for the target wrapped in `casts`, as a comparison of + /// the cast target with a bound of the cast target's type. Every simplification + /// goes through this, so that they agree on what a guarantee says: + /// + /// - A bound of another type than the target's is brought into the target's type + /// exactly: "safe" casts between floating point types or to decimals may still + /// round, so the cast must round trip, and the types must be ordered alike (see + /// HaveSameOrder). + /// - The bound is then put through the casts the target goes through. Each cast is + /// nondecreasing, so `x <= b` implies `cast(x) <= cast(b)` even if `cast(b)` + /// rounds, and no exactness is needed. + /// - Through casts which are not injective `x > b` only implies + /// `cast(x) >= cast(b)`, so strict comparisons are weakened to inclusive ones. + /// + /// Nullopt if the inequality cannot be restated; ignoring it only loses precision. + std::optional>> Restate( const OrderPreservingCasts& casts) const { - Comparison::type cast_cmp = casts.ApplyToComparison(cmp); - if (casts.injective) return std::make_pair(cast_cmp, bound); - if (!HaveSameOrder(*bound.type(), *casts.inner->type())) return std::nullopt; + std::shared_ptr value = bound.scalar(); + // A comparison against null or NaN cannot be guaranteed true, but be defensive. + if (!value->is_valid || IsNaNScalar(*value)) return std::nullopt; + 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)); + const auto& target_type = casts.inner->type(); + if (!value->type->Equals(*target_type)) { + if (!HaveSameOrder(*value->type, *target_type)) return std::nullopt; + auto maybe_cast = + Cast(Datum(value), target_type, CastOptions::Safe(), &exec_context); + if (!maybe_cast.ok()) return std::nullopt; + auto maybe_round_trip = + Cast(*maybe_cast, value->type, CastOptions::Safe(), &exec_context); + if (!maybe_round_trip.ok() || !maybe_round_trip->scalar()->Equals(*value)) { + return std::nullopt; + } + value = maybe_cast->scalar(); + } + auto cast_value = casts.ApplyToBound(std::move(value), &exec_context); + if (!cast_value.ok()) return std::nullopt; + // A valid non-NaN bound stays valid and non-NaN through order preserving casts. + return std::make_pair(casts.ApplyToComparison(cmp), *std::move(cast_value)); } /// The given expression simplifies to `value` if the inequality @@ -1368,10 +1412,8 @@ struct Inequality { auto options = checked_pointer_cast(is_in_call->options); - // 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) { + // Large value sets are handled by the kernel instead, see SimplifyWithRanges. + if (options->value_set.length() > kIsInSimplificationMaxValueSet) { return std::nullopt; } @@ -1382,14 +1424,14 @@ struct Inequality { // 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; + auto restated = guarantee.Restate(casts); + if (!restated) return std::nullopt; + const auto& [cmp, bound] = *restated; // Statistics based guarantees cannot bound NaN (a row group's minimum and // maximum ignore it) whereas `is_in` matches NaN with NaN, so a NaN in the - // value set must not be filtered out. Leave such value sets alone, as - // SimplifyIsInWithInequalities does. + // value set must not be filtered out. Leave such value sets alone, as the + // kernel does for large ones. const auto value_set_type_id = options->value_set.type()->id(); if (value_set_type_id == Type::FLOAT || value_set_type_id == Type::DOUBLE) { ARROW_ASSIGN_OR_RAISE(Datum is_nan, IsNan(options->value_set)); @@ -1478,9 +1520,9 @@ struct Inequality { 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; + auto restated = guarantee.Restate(casts); + if (!restated) return expr; + const auto& [guarantee_cmp, bound] = *restated; // Whether the RHS of the expression is EQUAL, LESS, or GREATER than the // RHS of the guarantee. N.B. Comparison::type is a bitmask @@ -1532,174 +1574,135 @@ struct Inequality { } }; -/// \brief What an `is_in` call evaluates to once we know that none of its -/// (non-null) value set members can occur in the data. +/// \brief What a guarantee says about the fields it bounds: for each field, the +/// inequalities it satisfies and whether it may be null. /// -/// \param arg the argument of the call -/// \param nullable whether `arg` may be null -/// \return the equivalent Expression, or nullopt if the call may still be true -Result> SimplifyIsInWithoutMatches( - const Expression& arg, const SetLookupOptions& options, bool value_set_has_null, - bool nullable) { - switch (options.GetNullMatchingBehavior()) { - case SetLookupOptions::MATCH: - // A null in the value set matches nulls in the data. - if (nullable && value_set_has_null) return std::nullopt; - return literal(false); - case SetLookupOptions::SKIP: - // Nulls in the data produce false. - return literal(false); - case SetLookupOptions::INCONCLUSIVE: - // Any non-match produces null if the value set contains a null, as do - // nulls in the data. - if (value_set_has_null) return literal(MakeNullScalar(boolean())); - [[fallthrough]]; - case SetLookupOptions::EMIT_NULL: - // Nulls in the data produce null. - if (!nullable) return literal(false); - return ValueUnlessNull(arg, false); +/// The inequalities refer into the guarantee, which must outlive this. +struct FieldRanges { + std::unordered_map, FieldRef::Hash> map; + + /// Collect the inequalities implied by the members of a guarantee. A field can + /// only be null if every inequality on it admits nulls and it is not also + /// guaranteed valid. + static FieldRanges Make(const std::vector& conjunction_members) { + FieldRanges ranges; + for (const auto& guarantee : conjunction_members) { + for (const auto& inequality : Inequality::Extract(guarantee)) { + ranges.map[inequality.target].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; + auto it = ranges.map.find(*target); + if (it != ranges.map.end()) it->second.front().nullable = false; + } + for (auto& [target, inequalities] : ranges.map) { + const bool nullable = + std::all_of(inequalities.begin(), inequalities.end(), + [](const Inequality& inequality) { return inequality.nullable; }); + for (auto& inequality : inequalities) inequality.nullable = nullable; + } + return ranges; } - return std::nullopt; -} -/// \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 -/// searches the sorted copy of the value set kept in the kernel state for the -/// range of values admitted by the inequalities (`search_sorted`). If that range is -/// empty the call is simplified to what it evaluates to without matches (`false`, -/// or null for the rows the null matching behavior makes null); otherwise it is -/// left unchanged. This costs O(log N) per guarantee, with no rebinding, and is exact -/// for value sets of up to kIsInMaxSortedValueSet values. Longer value sets are only -/// bounded by their smallest and largest value. -/// -/// Applying all inequalities on a field together matters: a row group's -/// statistics yield `x >= min and x <= max`, and neither bound on its own can -/// rule out a value set whose hull spans the row group. -Result SimplifyIsInWithInequalities( - const std::vector& inequalities, Expression expr) { - auto call = expr.call(); - if (!call || call->function_name != "is_in") return expr; - if (!call->kernel_state) return expr; - - auto options = checked_pointer_cast(call->options); - if (options->value_set.length() <= internal::kIsInSimplificationMaxValueSet) { - // Handled by Inequality::SimplifyIsIn - return expr; + const std::vector* Find(const FieldRef& target) const { + auto it = map.find(target); + return it != map.end() ? &it->second : nullptr; } +}; - // 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(); +/// \brief The range which the inequalities on a field leave to the field wrapped in +/// `casts`: the tightest lower and upper bound, in the type of the cast field. +/// +/// Considering 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. +internal::ArgumentRange ResolveRange(const std::vector& inequalities, + const OrderPreservingCasts& casts) { + internal::ArgumentRange range; + DCHECK(!inequalities.empty()); + range.nullable = inequalities.front().nullable; + + // Whether a bound is tighter than the current one on that side. Bounds which + // cannot be compared are ignored, which only loses precision. + auto is_tighter = [](const std::shared_ptr& bound, bool inclusive, + const std::shared_ptr& current, bool current_inclusive, + Comparison::type tighter) { + if (!current) return true; + auto cmp = Comparison::Execute(bound, current); + if (!cmp.ok()) return false; + if (*cmp == Comparison::EQUAL) return current_inclusive && !inclusive; + return *cmp == tighter; + }; - std::vector applicable; for (const Inequality& inequality : inequalities) { - if (inequality.target == target) 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. - const auto sorted_value_set = state.GetSortedValueSet(); - const std::shared_ptr& sorted = sorted_value_set.values; - 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. A bound of a type which is ordered - // differently (a string bound on a field read as an integer, say) says - // nothing about the field's order even if it round trips. - if (!HaveSameOrder(*bound->type, *lhs->type())) continue; - 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(); + auto restated = inequality.Restate(casts); + if (!restated) continue; + const auto& [cmp, bound] = *restated; + const bool inclusive = cmp & Comparison::EQUAL; + // Unless the comparison admits values below the bound, it is a lower bound; + // likewise for values above the bound. Equality is both. + if (!(cmp & Comparison::LESS) && is_tighter(bound, inclusive, range.lower, + range.lower_inclusive, + Comparison::GREATER)) { + range.lower = bound; + range.lower_inclusive = inclusive; } - // 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. The last cast is to the argument type, and - // a valid non-NaN bound stays valid and non-NaN. - auto cast_bound = casts.ApplyToBound(std::move(bound), &exec_context); - if (!cast_bound.ok()) continue; - bound = *cast_bound; - DCHECK(bound->type->Equals(*sorted->type())); - - // 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); - DCHECK_NE(cmp, Comparison::NA); - - // The index of the first value not less than the bound (Left) or greater than - // the bound (Right). - auto search = [&](SearchSortedOptions::Side side) -> Result { - ARROW_ASSIGN_OR_RAISE( - Datum index, - SearchSorted(sorted, bound, SearchSortedOptions(side), &exec_context)); - return static_cast(index.scalar_as().value); - }; - // Unless the comparison admits values below the bound, skip them (and the - // bound itself unless it is admitted); likewise for values above the bound. - const bool admits_bound = cmp & Comparison::EQUAL; - if (!(cmp & Comparison::LESS)) { - ARROW_ASSIGN_OR_RAISE( - int64_t index, - search(admits_bound ? SearchSortedOptions::Left : SearchSortedOptions::Right)); - begin = std::max(begin, index); - } - if (!(cmp & Comparison::GREATER)) { - ARROW_ASSIGN_OR_RAISE( - int64_t index, - search(admits_bound ? SearchSortedOptions::Right : SearchSortedOptions::Left)); - end = std::min(end, index); + if (!(cmp & Comparison::GREATER) && is_tighter(bound, inclusive, range.upper, + range.upper_inclusive, + Comparison::LESS)) { + range.upper = bound; + range.upper_inclusive = inclusive; } } + return range; +} - // Some value set member may occur in the data: leave the call unchanged rather - // than rebinding it against a narrowed value set. Of a value set known only by its - // smallest and largest value, a member may occur unless all admitted values lie - // below or above both. - const bool may_match = - sorted_value_set.exact ? begin < end : begin < sorted->length() && end > 0; - if (may_match) return expr; +/// \brief Simplify a call against what the guarantee says about the field it is +/// applied to (possibly through order preserving casts). +/// +/// Kernels which can tell what they evaluate to over a range (currently `is_in`, +/// which searches its sorted value set) are asked; this costs O(log N) per guarantee +/// with no rebinding. Comparisons, validity checks and small `is_in` value sets are +/// simplified against each inequality in turn. +Result SimplifyWithRanges(const FieldRanges& ranges, Expression expr) { + auto call = expr.call(); + if (!call || call->arguments.empty()) return expr; - ARROW_ASSIGN_OR_RAISE(auto simplified, - SimplifyIsInWithoutMatches(call->arguments[0], *options, - state.value_set_has_null, nullable)); - return simplified ? *std::move(simplified) : expr; + auto casts = OrderPreservingCasts::Strip(call->arguments[0]); + if (!casts.inner->field_ref()) return expr; + const auto* inequalities = ranges.Find(*casts.inner->field_ref()); + if (!inequalities) return expr; + + const bool is_small_is_in = + call->function_name == "is_in" && + checked_cast(*call->options).value_set.length() <= + kIsInSimplificationMaxValueSet; + const auto* evaluator = + dynamic_cast(call->kernel_state.get()); + if (evaluator && call->arguments.size() == 1 && !is_small_is_in) { + ARROW_ASSIGN_OR_RAISE( + auto evaluation, evaluator->EvaluateOverRange(ResolveRange(*inequalities, casts))); + switch (evaluation) { + case internal::RangeEvaluation::kUnknown: + return expr; + case internal::RangeEvaluation::kFalse: + return literal(false); + case internal::RangeEvaluation::kNull: + return literal(MakeNullScalar(call->type.GetSharedPtr())); + case internal::RangeEvaluation::kFalseUnlessNull: + return ValueUnlessNull(call->arguments[0], false); + } + return expr; + } + + for (const Inequality& inequality : *inequalities) { + ARROW_ASSIGN_OR_RAISE(expr, inequality.Simplify(std::move(expr))); + } + return expr; } /// \brief Simplify an expression given a guarantee, if the guarantee @@ -1745,34 +1748,16 @@ Result SimplifyWithGuarantee(Expression expr, }; RETURN_NOT_OK(CanonicalizeAndFoldConstants()); - // The inequalities on fields implied by the guarantee. A field which is also - // guaranteed valid cannot be null, whatever its inequalities admit. They refer - // into conjunction_members, which must not change from here on. - 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()) { + // What the guarantee says about each field it bounds. This refers into + // conjunction_members, which must not change from here on. Every call is + // simplified against the range of its field in a single pass. + const FieldRanges ranges = FieldRanges::Make(conjunction_members); + if (!ranges.map.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)); + return SimplifyWithRanges(ranges, std::move(expr)); })); if (!Expression::Identical(expr, before)) { @@ -1780,20 +1765,6 @@ Result SimplifyWithGuarantee(Expression expr, } } - 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 (Expression::Identical(simplified, expr)) continue; - - expr = std::move(simplified); - RETURN_NOT_OK(CanonicalizeAndFoldConstants()); - } - for (const auto& guarantee : conjunction_members) { if (!guarantee.call() || guarantee.call()->function_name != "is_valid") continue; diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 9e0f482f692a..09db9ef597af 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2092,6 +2092,41 @@ TEST(Expression, SimplifyWithNullableRangeGuarantee) { Simplify{is_null(i32)}.WithGuarantee(valid_range).Expect(false); } +TEST(Expression, SimplifyWithBoundOfAnotherType) { + // A guarantee may bound a field with a value of another type, for example the + // statistics of a Parquet column which is read as another type. + auto i64 = field_ref("i64"); + auto is_in = [&](std::string json_array) { + return call("is_in", {i64}, SetLookupOptions{ArrayFromJSON(int64(), json_array)}); + }; + + // Bounds which convert exactly to the field's type are used + auto int32_range = and_(greater_equal(i64, literal(9)), less_equal(i64, literal(10))); + auto double_range = + and_(greater_equal(i64, literal(9.0)), less_equal(i64, literal(10.0))); + for (const auto& range : {int32_range, double_range}) { + Simplify{equal(i64, literal(5))}.WithGuarantee(range).Expect(false); + Simplify{greater(i64, literal(5))}.WithGuarantee(range).Expect(true); + Simplify{is_in("[5, 9]")}.WithGuarantee(range).Expect(is_in("[9]")); + } + + // ... while those which do not are ignored rather than rounded + auto inexact_range = + and_(greater_equal(i64, literal(9.5)), less_equal(i64, literal(10.5))); + Simplify{equal(i64, literal(5))}.WithGuarantee(inexact_range).ExpectUnchanged(); + Simplify{is_in("[5, 9]")}.WithGuarantee(inexact_range).ExpectUnchanged(); + + // ... as are bounds of a type which is ordered differently, even though they cast + // exactly: "10" <= x <= "9" admits "5" + auto string_range = + and_(greater_equal(i64, literal("10")), less_equal(i64, literal("9"))); + Simplify{equal(i64, literal(5))}.WithGuarantee(string_range).ExpectUnchanged(); + Simplify{greater(i64, literal(50))} + .WithGuarantee(string_range) + .ExpectUnchanged(); + Simplify{is_in("[5, 9]")}.WithGuarantee(string_range).ExpectUnchanged(); +} + TEST(Expression, SimplifyComparisonThroughNonInjectiveCast) { auto i32 = field_ref("i32"); auto i64 = field_ref("i64"); diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc index 8665251e1641..097157c42702 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc @@ -15,12 +15,19 @@ // specific language governing permissions and limitations // under the License. +#include +#include + #include "arrow/array/array_base.h" +#include "arrow/array/builder_base.h" +#include "arrow/array/concatenate.h" +#include "arrow/compute/api_aggregate.h" #include "arrow/compute/api_scalar.h" +#include "arrow/compute/api_vector.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/range_evaluation_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/type.h" #include "arrow/util/bit_util.h" @@ -37,6 +44,218 @@ using internal::HashTraits; namespace compute::internal { namespace { +/// Value sets up to this length are sorted for RangeEvaluator::EvaluateOverRange. +/// Sorting costs more than building the lookup table of the kernel, so longer value +/// sets are only bounded by their smallest and largest value, which is cheap to find +/// and prunes as well when the values are clustered. +constexpr int64_t kIsInMaxSortedValueSet = 1 << 17; + +/// Whether a (valid) scalar of a floating point type is NaN. +bool IsNaNScalar(const Scalar& value) { + switch (value.type->id()) { + case Type::FLOAT: + return std::isnan(checked_cast(value).value); + case Type::DOUBLE: + return std::isnan(checked_cast(value).value); + default: + return false; + } +} + +/// Whether value sets of this type are sorted by `sort_indices` and searched by +/// `search_sorted` in the order of the comparison kernels. +bool IsSearchableValueSetType(const DataType& type) { + const auto id = type.id(); + return is_integer(id) || id == Type::FLOAT || id == Type::DOUBLE || is_temporal(id) || + id == Type::DURATION || is_base_binary_like(id); +} + +/// Common state for the is_in and index_in kernels. +/// +/// For `is_in` it also tells what the call evaluates to when the argument is known +/// to lie in a range (RangeEvaluator), which lets expression simplification prune +/// filters against statistics: if no value of the value set lies in the range, the +/// call evaluates to false, or to null for the rows which the null matching behavior +/// makes null. +struct SetLookupStateBase : public KernelState, public RangeEvaluator { + std::shared_ptr value_set_type; + + SetLookupOptions::NullMatchingBehavior null_matching_behavior; + + /// Whether the value set contains a null (regardless of null matching behavior). + bool value_set_has_null = false; + + /// Keep the value set (as matched against the argument) for EvaluateOverRange(). + /// + /// Called once by the `is_in` kernel init, before the state is shared. Nothing + /// is computed until the first call of EvaluateOverRange(), so that only + /// expressions which are simplified against guarantees pay for it. + void SetSearchableValueSet(Datum value_set, MemoryPool* pool) { + searchable_value_set_ = std::move(value_set); + searchable_value_set_pool_ = pool; + } + + Result EvaluateOverRange(const ArgumentRange& range) const override { + if (!range.lower && !range.upper) return RangeEvaluation::kUnknown; + for (const auto* bound : {&range.lower, &range.upper}) { + if (*bound == nullptr) continue; + // 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 range refers to) need not be the value set's order. + if (!(*bound)->type->Equals(*value_set_type)) return RangeEvaluation::kUnknown; + // A range cannot be bounded by null or NaN, but be defensive. + if (!(*bound)->is_valid || IsNaNScalar(**bound)) return RangeEvaluation::kUnknown; + } + + // Only now materialize the (lazily sorted) value set. + const SortedValueSet& sorted_value_set = GetSortedValueSet(); + const std::shared_ptr& sorted = sorted_value_set.values; + if (!sorted) return RangeEvaluation::kUnknown; + + // Indices into `sorted` delimiting the values in the range: O(log N) per bound. + // A failed search only costs the simplification. + ExecContext exec_context(searchable_value_set_pool_); + auto search = [&](const std::shared_ptr& bound, + SearchSortedOptions::Side side) -> std::optional { + auto index = SearchSorted(sorted, bound, SearchSortedOptions(side), &exec_context); + if (!index.ok() || !index->scalar()->is_valid) return std::nullopt; + return static_cast(index->scalar_as().value); + }; + int64_t begin = 0; + int64_t end = sorted->length(); + if (range.lower) { + // The first value not less than (or, for an exclusive bound, greater than) it + auto index = search(range.lower, range.lower_inclusive ? SearchSortedOptions::Left + : SearchSortedOptions::Right); + if (!index) return RangeEvaluation::kUnknown; + begin = *index; + } + if (range.upper) { + auto index = search(range.upper, range.upper_inclusive ? SearchSortedOptions::Right + : SearchSortedOptions::Left); + if (!index) return RangeEvaluation::kUnknown; + end = *index; + } + + // Of a value set known only by its smallest and largest value, a member may lie + // in the range unless the whole range is below or above both. + const bool may_match = + sorted_value_set.exact ? begin < end : begin < sorted->length() && end > 0; + if (may_match) return RangeEvaluation::kUnknown; + + // No (non-null) member of the value set can occur in the data. + switch (null_matching_behavior) { + case SetLookupOptions::MATCH: + // A null in the value set matches nulls in the data. + if (range.nullable && value_set_has_null) return RangeEvaluation::kUnknown; + return RangeEvaluation::kFalse; + case SetLookupOptions::SKIP: + // Nulls in the data produce false. + return RangeEvaluation::kFalse; + case SetLookupOptions::INCONCLUSIVE: + // Any non-match produces null if the value set contains a null, as do + // nulls in the data. + if (value_set_has_null) return RangeEvaluation::kNull; + [[fallthrough]]; + case SetLookupOptions::EMIT_NULL: + // Nulls in the data produce null. + return range.nullable ? RangeEvaluation::kFalseUnlessNull + : RangeEvaluation::kFalse; + } + return RangeEvaluation::kUnknown; + } + + private: + struct SortedValueSet { + /// Ascending values without nulls, or nullptr if the value set cannot be searched + std::shared_ptr values; + /// Whether `values` holds every value of the value set, or only its smallest and + /// largest value (for value sets longer than kIsInMaxSortedValueSet) + bool exact = true; + }; + + /// The value set sorted for searching with `search_sorted`. + /// + /// This is only available for `is_in` when the value set has a type for which + /// IsSearchableValueSetType() holds and contains no NaN. It is computed at most + /// once; concurrent calls are safe. + const SortedValueSet& GetSortedValueSet() const { + std::call_once(sorted_value_set_once_, [this] { + if (searchable_value_set_.is_arraylike()) { + // Best effort: a failure only disables the simplification of `is_in` + // against guarantees. + ExecContext exec_context(searchable_value_set_pool_); + sorted_value_set_ = + SortValueSet(searchable_value_set_, &exec_context).ValueOr(SortedValueSet{}); + searchable_value_set_ = Datum(); + } + }); + return sorted_value_set_; + } + + static Result SortValueSet(const Datum& value_set, ExecContext* ctx) { + std::shared_ptr values; + if (value_set.is_array()) { + values = value_set.make_array(); + } else { + ARROW_ASSIGN_OR_RAISE(values, Concatenate(value_set.chunked_array()->chunks(), + ctx->memory_pool())); + } + if (values->length() > kIsInMaxSortedValueSet) { + ARROW_ASSIGN_OR_RAISE(auto bounds, ValueSetBounds(values, ctx)); + return SortedValueSet{std::move(bounds), /*exact=*/false}; + } + ARROW_ASSIGN_OR_RAISE(auto sorted, SortValues(values, ctx)); + return SortedValueSet{std::move(sorted), /*exact=*/true}; + } + + /// \return `values` sorted with nulls removed, or nullptr if it contains NaN + static Result> SortValues(const std::shared_ptr& values, + ExecContext* ctx) { + ARROW_ASSIGN_OR_RAISE( + auto indices, + SortIndices(*values, ArraySortOptions(SortOrder::Ascending, NullPlacement::AtEnd), + ctx)); + ARROW_ASSIGN_OR_RAISE(Datum taken, + Take(values, indices, TakeOptions::NoBoundsCheck(), ctx)); + 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) { + ARROW_ASSIGN_OR_RAISE(auto last, sorted->GetScalar(sorted->length() - 1)); + if (IsNaNScalar(*last)) return nullptr; + } + return sorted; + } + + /// \return the smallest and largest value of `values`, or nullptr if it contains NaN + /// or only nulls + static Result> ValueSetBounds( + const std::shared_ptr& values, ExecContext* ctx) { + if (is_floating(values->type_id())) { + ARROW_ASSIGN_OR_RAISE(Datum is_nan, IsNan(values, ctx)); + ARROW_ASSIGN_OR_RAISE(Datum any_nan, + Any(is_nan, ScalarAggregateOptions::Defaults(), ctx)); + if (any_nan.scalar_as().value) return nullptr; + } + ARROW_ASSIGN_OR_RAISE(Datum min_max, + MinMax(values, ScalarAggregateOptions::Defaults(), ctx)); + const auto& bounds = min_max.scalar_as().value; + if (!bounds[0]->is_valid) return nullptr; + ARROW_ASSIGN_OR_RAISE(auto builder, MakeBuilder(values->type(), ctx->memory_pool())); + RETURN_NOT_OK(builder->AppendScalar(*bounds[0])); + RETURN_NOT_OK(builder->AppendScalar(*bounds[1])); + return builder->Finish(); + } + + mutable std::once_flag sorted_value_set_once_; + // Released once sorted_value_set_ is computed. + mutable Datum searchable_value_set_; + MemoryPool* searchable_value_set_pool_ = default_memory_pool(); + mutable SortedValueSet sorted_value_set_; +}; + template struct SetLookupState : public SetLookupStateBase { explicit SetLookupState(MemoryPool* pool) : memory_pool(pool) {} @@ -118,7 +337,6 @@ struct SetLookupState : public SetLookupStateBase { // be mapped back to indices in the value_set. std::vector memo_index_to_value_index; int32_t null_index = -1; - SetLookupOptions::NullMatchingBehavior null_matching_behavior; }; template <> @@ -132,8 +350,6 @@ struct SetLookupState : public SetLookupStateBase { value_set_type = null(); return Status::OK(); } - - SetLookupOptions::NullMatchingBehavior null_matching_behavior; }; // TODO: Put this concept somewhere reusable @@ -167,16 +383,15 @@ struct InitStateVisitor { TypeHolder arg_type; std::unique_ptr result; - // Whether to make the value set available to - // SetLookupStateBase::GetSortedValueSet (is_in only). - bool compute_sorted_value_set; + // Whether to keep the value set in the state so that the kernel can be evaluated + // over a range, see SetLookupStateBase::EvaluateOverRange (is_in only). + bool keep_value_set; - InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args, - bool compute_sorted_value_set) + InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args, bool keep_value_set) : ctx(ctx), options(*checked_cast(args.options)), arg_type(args.inputs[0]), - compute_sorted_value_set(compute_sorted_value_set) {} + keep_value_set(keep_value_set) {} template Status Init() { @@ -257,9 +472,7 @@ 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())) { + if (keep_value_set && IsSearchableValueSetType(*options.value_set.type())) { checked_cast(result.get()) ->SetSearchableValueSet(options.value_set, ctx->exec_context()->memory_pool()); } @@ -269,23 +482,23 @@ struct InitStateVisitor { Result> InitSetLookup(KernelContext* ctx, const KernelInitArgs& args, - bool compute_sorted_value_set) { + bool keep_value_set) { if (args.options == nullptr) { return Status::Invalid( "Attempted to call a set lookup function without SetLookupOptions"); } - return InitStateVisitor{ctx, args, compute_sorted_value_set}.GetResult(); + return InitStateVisitor{ctx, args, keep_value_set}.GetResult(); } Result> InitIsIn(KernelContext* ctx, const KernelInitArgs& args) { - return InitSetLookup(ctx, args, /*compute_sorted_value_set=*/true); + return InitSetLookup(ctx, args, /*keep_value_set=*/true); } Result> InitIndexIn(KernelContext* ctx, const KernelInitArgs& args) { - return InitSetLookup(ctx, args, /*compute_sorted_value_set=*/false); + return InitSetLookup(ctx, args, /*keep_value_set=*/false); } struct IndexInVisitor { diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h deleted file mode 100644 index 5a9bd5fb2c03..000000000000 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h +++ /dev/null @@ -1,184 +0,0 @@ -// 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 "arrow/array/array_base.h" -#include "arrow/array/builder_base.h" -#include "arrow/array/concatenate.h" -#include "arrow/compute/api_aggregate.h" -#include "arrow/compute/api_scalar.h" -#include "arrow/compute/api_vector.h" -#include "arrow/compute/kernel.h" -#include "arrow/datum.h" -#include "arrow/memory_pool.h" -#include "arrow/result.h" -#include "arrow/scalar.h" -#include "arrow/type.h" -#include "arrow/type_traits.h" -#include "arrow/util/checked_cast.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 searched with `search_sorted` in -/// SetLookupStateBase::GetSortedValueSet, which only ever simplifies the call to a -/// result without matches, see https://github.com/apache/arrow/issues/46777 -constexpr int64_t kIsInSimplificationMaxValueSet = 50; - -/// Value sets up to this length are sorted for SetLookupStateBase::GetSortedValueSet. -/// Sorting costs more than building the lookup table of the kernel, so longer value -/// sets are only bounded by their smallest and largest value, which is cheap to find -/// and prunes as well when the values are clustered. -constexpr int64_t kIsInMaxSortedValueSet = 1 << 17; - -/// Whether a (valid) scalar of a floating point type is NaN. -inline bool IsNaNScalar(const Scalar& value) { - switch (value.type->id()) { - case Type::FLOAT: - return std::isnan(::arrow::internal::checked_cast(value).value); - case Type::DOUBLE: - return std::isnan( - ::arrow::internal::checked_cast(value).value); - default: - return false; - } -} - -/// Whether value sets of this type are sorted by `sort_indices` and searched by -/// `search_sorted` in the order of the comparison kernels. -inline bool IsSearchableValueSetType(const DataType& type) { - const auto id = type.id(); - return is_integer(id) || id == Type::FLOAT || id == Type::DOUBLE || is_temporal(id) || - id == Type::DURATION || is_base_binary_like(id); -} - -/// 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. Nothing - /// is computed until the first call of GetSortedValueSet(), so that only - /// expressions which are simplified against guarantees pay for it. - void SetSearchableValueSet(Datum value_set, MemoryPool* pool) { - searchable_value_set_ = std::move(value_set); - searchable_value_set_pool_ = pool; - } - - struct SortedValueSet { - /// Ascending values without nulls, or nullptr if the value set cannot be searched - std::shared_ptr values; - /// Whether `values` holds every value of the value set, or only its smallest and - /// largest value (for value sets longer than kIsInMaxSortedValueSet) - bool exact = true; - }; - - /// The value set sorted for searching with `search_sorted`. - /// - /// 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. - SortedValueSet GetSortedValueSet() const { - std::call_once(sorted_value_set_once_, [this] { - if (searchable_value_set_.is_arraylike()) { - // Best effort: a failure only disables the simplification of `is_in` - // against guarantees. - sorted_value_set_ = - SortValueSet(searchable_value_set_, searchable_value_set_pool_) - .ValueOr(SortedValueSet{}); - searchable_value_set_ = Datum(); - } - }); - return sorted_value_set_; - } - - private: - static Result SortValueSet(const Datum& value_set, MemoryPool* pool) { - std::shared_ptr values; - if (value_set.is_array()) { - values = value_set.make_array(); - } else { - ARROW_ASSIGN_OR_RAISE(values, - Concatenate(value_set.chunked_array()->chunks(), pool)); - } - if (values->length() > kIsInMaxSortedValueSet) { - ARROW_ASSIGN_OR_RAISE(auto bounds, ValueSetBounds(values)); - return SortedValueSet{std::move(bounds), /*exact=*/false}; - } - ARROW_ASSIGN_OR_RAISE(auto sorted, SortValues(values)); - return SortedValueSet{std::move(sorted), /*exact=*/true}; - } - - /// \return `values` sorted with nulls removed, or nullptr if it contains NaN - static Result> SortValues(const std::shared_ptr& 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) { - ARROW_ASSIGN_OR_RAISE(auto last, sorted->GetScalar(sorted->length() - 1)); - if (IsNaNScalar(*last)) return nullptr; - } - return sorted; - } - - /// \return the smallest and largest value of `values`, or nullptr if it contains NaN - /// or only nulls - static Result> ValueSetBounds( - const std::shared_ptr& values) { - if (is_floating(values->type_id())) { - ARROW_ASSIGN_OR_RAISE(Datum is_nan, IsNan(values)); - ARROW_ASSIGN_OR_RAISE(Datum any_nan, Any(is_nan)); - if (any_nan.scalar_as().value) return nullptr; - } - ARROW_ASSIGN_OR_RAISE(Datum min_max, MinMax(values)); - const auto& bounds = min_max.scalar_as().value; - if (!bounds[0]->is_valid) return nullptr; - ARROW_ASSIGN_OR_RAISE(auto builder, MakeBuilder(values->type())); - RETURN_NOT_OK(builder->AppendScalar(*bounds[0])); - RETURN_NOT_OK(builder->AppendScalar(*bounds[1])); - return builder->Finish(); - } - - mutable std::once_flag sorted_value_set_once_; - // Released once sorted_value_set_ is computed. - mutable Datum searchable_value_set_; - MemoryPool* searchable_value_set_pool_ = default_memory_pool(); - mutable SortedValueSet sorted_value_set_; -}; - -} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/compute/range_evaluation_internal.h b/cpp/src/arrow/compute/range_evaluation_internal.h new file mode 100644 index 000000000000..9714e8079449 --- /dev/null +++ b/cpp/src/arrow/compute/range_evaluation_internal.h @@ -0,0 +1,70 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include + +#include "arrow/result.h" +#include "arrow/type_fwd.h" +#include "arrow/util/visibility.h" + +namespace arrow::compute::internal { + +/// \brief The values the argument of a unary call is known to take, for example +/// from the statistics of a Parquet row group. +struct ArgumentRange { + /// The smallest and largest value the argument may take, in the argument's + /// type and neither null nor NaN. A null pointer leaves that side unbounded. + std::shared_ptr lower, upper; + /// Whether the argument may be equal to the respective bound + bool lower_inclusive = true, upper_inclusive = true; + /// Whether the argument may be null, in addition to the values in the range + bool nullable = true; +}; + +/// \brief What a call evaluates to for every argument in an ArgumentRange +enum class RangeEvaluation { + /// Nothing is known: the call must be evaluated + kUnknown, + /// false for every argument, null or not + kFalse, + /// null for every argument, null or not + kNull, + /// false for arguments which are not null and null for those which are + kFalseUnlessNull, +}; + +/// \brief Interface for kernel states which can tell what their (unary) kernel +/// evaluates to when its argument lies in a range. +/// +/// A KernelState may implement this in addition to KernelState. Expression +/// simplification (SimplifyWithGuarantee) discovers it with dynamic_cast and +/// replaces calls for which the result is known, which is how filters are +/// pruned against statistics without the expression machinery knowing the +/// semantics of the function. +/// +/// Implementations must be safe to call concurrently: bound expressions and +/// their kernel states are shared between threads. +class ARROW_EXPORT RangeEvaluator { + public: + virtual ~RangeEvaluator() = default; + + virtual Result EvaluateOverRange(const ArgumentRange& range) const = 0; +}; + +} // namespace arrow::compute::internal From 75569514be9a8dbd4203f2d27d33f4a0d4605efa Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 11 Nov 2026 00:27:48 +0100 Subject: [PATCH 11/12] GH-55292: [C++] Drop guards no test or caller depends on and test the ones which matter Removing each piece of the simplification in turn showed that the NaN check on guarantee bounds in expression.cc duplicates the kernel's, and that decimal bounds never converted to an integer or floating point field's type; both are removed. Guards which no test covered now have one: NaN in a sorted value set and in a value set bounded by its smallest and largest value, a strict bound next to an inclusive one at the same value, date bounds on a timestamp field, and index_in never being simplified. --- cpp/src/arrow/compute/expression.cc | 20 +++--------------- cpp/src/arrow/compute/expression_test.cc | 26 ++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index f72953ad363d..66d8ca8cad45 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -20,7 +20,6 @@ #include "arrow/compute/expression.h" #include -#include #include #include #include @@ -1226,18 +1225,6 @@ struct OrderPreservingCasts { // the guarantee (see internal::RangeEvaluator), which never rebinds. constexpr int64_t kIsInSimplificationMaxValueSet = 50; -// Whether a (valid) scalar of a floating point type is NaN. -bool IsNaNScalar(const Scalar& value) { - switch (value.type->id()) { - case Type::FLOAT: - return std::isnan(checked_cast(value).value); - case Type::DOUBLE: - return std::isnan(checked_cast(value).value); - default: - return false; - } -} - // Whether values of the two types are ordered alike, so that a guarantee bound of // one type which casts exactly to the other bounds the same values there. This does // not hold across families: "10" <= "9" as strings, but both cast exactly to @@ -1245,7 +1232,7 @@ bool IsNaNScalar(const Scalar& value) { bool HaveSameOrder(const DataType& left, const DataType& right) { if (left.id() == right.id()) return true; auto family = [](Type::type id) { - if (is_integer(id) || is_floating(id) || is_decimal(id)) return 1; + if (is_integer(id) || is_floating(id)) return 1; if (is_base_binary_like(id)) return 2; if (id == Type::DATE32 || id == Type::DATE64 || id == Type::TIMESTAMP) return 3; return 0; @@ -1363,8 +1350,8 @@ struct Inequality { std::optional>> Restate( const OrderPreservingCasts& casts) const { std::shared_ptr value = bound.scalar(); - // A comparison against null or NaN cannot be guaranteed true, but be defensive. - if (!value->is_valid || IsNaNScalar(*value)) return std::nullopt; + // A comparison against null cannot be guaranteed true, but be defensive. + if (!value->is_valid) return std::nullopt; ExecContext exec_context; const auto& target_type = casts.inner->type(); @@ -1382,7 +1369,6 @@ struct Inequality { } auto cast_value = casts.ApplyToBound(std::move(value), &exec_context); if (!cast_value.ok()) return std::nullopt; - // A valid non-NaN bound stays valid and non-NaN through order preserving casts. return std::make_pair(casts.ApplyToComparison(cmp), *std::move(cast_value)); } diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 09db9ef597af..58277d4767a3 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2421,6 +2421,15 @@ TEST(Expression, SimplifyIsInLargeValueSet) { and_(and_(greater(i32, literal(10)), greater(i32, literal(50))), less(i32, literal(100))), kFalse}, + // ... of which the strict one wins over an inclusive one at the same value + {i32, int32(), gapped, + and_(and_(greater_equal(i32, literal(50)), greater(i32, literal(50))), + less(i32, literal(100))), + kFalse}, + {i32, int32(), gapped, + and_(and_(less(i32, literal(100)), less_equal(i32, literal(100))), + greater(i32, literal(50))), + kFalse}, // Bounds on other fields are irrelevant, including those of a range which may // be null for another field {i32, int32(), gapped, and_(greater(u32, literal(50)), less(u32, literal(100))), @@ -2465,6 +2474,7 @@ TEST(Expression, SimplifyIsInLargeValueSet) { // Statistics cannot bound NaN, so a value set containing NaN is never simplified {f64, float64(), concat_json("[NaN]", range), greater(f64, literal(100.5)), kUnchanged}, + {f64, float64(), concat_json("[NaN]", range), less(f64, literal(0.5)), kUnchanged}, // 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. {cast(f64, float32()), float32(), make_range_json(1, 51), @@ -2509,6 +2519,10 @@ TEST(Expression, SimplifyIsInLargeValueSet) { // ... whereas bounds of another numeric type are converted {i64, int64(), concat_json("[5]", make_range_json(1000, 1059)), and_(greater_equal(i64, literal(9.0)), less_equal(i64, literal(10.0))), kFalse}, + // ... as are dates bounding a timestamp: day 1 is second 86400 + {field_ref("ts_s"), timestamp(TimeUnit::SECOND), range, + greater_equal(field_ref("ts_s"), literal(std::make_shared(1))), + kFalse}, // Dictionaries are not searchable {field_ref("dict_i32"), int32(), range, greater(field_ref("dict_i32"), literal(100)), kUnchanged}, @@ -2527,6 +2541,11 @@ TEST(Expression, SimplifyIsInLargeValueSet) { } } + // index_in shares the kernel state of is_in but is never simplified + Simplify{call("index_in", {i32}, SetLookupOptions{ArrayFromJSON(int32(), gapped)})} + .WithGuarantee(in_gap) + .ExpectUnchanged(); + // Value sets longer than kIsInMaxSortedValueSet are only bounded by their smallest // and largest value, so a guarantee inside a gap of the value set does not prune auto huge = call("is_in", {i32}, @@ -2539,6 +2558,13 @@ TEST(Expression, SimplifyIsInLargeValueSet) { .WithGuarantee(and_(greater(i32, literal(70000)), less(i32, literal(100001)))) .ExpectUnchanged(); Simplify{huge}.WithGuarantee(greater_equal(i32, literal(170000))).ExpectUnchanged(); + // ... unless they contain NaN, which the smallest and largest value ignore + auto huge_nan = + call("is_in", {f64}, + SetLookupOptions{ArrayFromJSON( + float64(), concat_json("[NaN]", make_range_json(1, 140000)))}); + Simplify{huge_nan}.WithGuarantee(less(f64, literal(0.5))).ExpectUnchanged(); + Simplify{huge_nan}.WithGuarantee(greater(f64, literal(140000.5))).ExpectUnchanged(); // Chunked value sets are concatenated before they are sorted SetLookupOptions chunked_options{ From 683b23e2c657913c765a71c21caa8531a2d5d5f0 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 11 Nov 2026 00:50:34 +0100 Subject: [PATCH 12/12] GH-55292: [C++] Round guarantee bounds which a safe cast rounds instead of ignoring them A bound of another type than the field's (the statistics of a float64 column read as float32, say) was ignored unless it converted exactly, so comparisons against such statistics were no longer simplified. The bound is now cast safely and, unless it round trips, strict comparisons are weakened to inclusive ones: 0.2 rounds up to 0.2f, which a stored 0.2 is read as. Also apply clang-format. --- cpp/src/arrow/compute/expression.cc | 46 ++++++++++++------- cpp/src/arrow/compute/expression_test.cc | 34 ++++++++++---- .../compute/kernels/scalar_set_lookup.cc | 9 ++-- 3 files changed, 60 insertions(+), 29 deletions(-) diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index 66d8ca8cad45..005ccc7cd566 100644 --- a/cpp/src/arrow/compute/expression.cc +++ b/cpp/src/arrow/compute/expression.cc @@ -1161,6 +1161,15 @@ Result Canonicalize(Expression expr, compute::ExecContext* exec_cont namespace { +// A strict comparison weakened to an inclusive one. A nondecreasing conversion which +// may map distinct values to the same value only preserves inclusive comparisons: +// `x > b` implies `f(x) >= f(b)`, but not `f(x) > f(b)`. +Comparison::type WeakenStrictComparison(Comparison::type cmp) { + if (cmp == Comparison::GREATER) return Comparison::GREATER_EQUAL; + if (cmp == Comparison::LESS) return Comparison::LESS_EQUAL; + return cmp; +} + // The order preserving casts wrapping an Expression, outermost first, and the // wrapped Expression. See Comparison::StripOrderPreservingCasts. struct OrderPreservingCasts { @@ -1211,10 +1220,7 @@ struct OrderPreservingCasts { // 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; + return injective ? cmp : WeakenStrictComparison(cmp); } }; @@ -1337,9 +1343,14 @@ struct Inequality { /// goes through this, so that they agree on what a guarantee says: /// /// - A bound of another type than the target's is brought into the target's type - /// exactly: "safe" casts between floating point types or to decimals may still - /// round, so the cast must round trip, and the types must be ordered alike (see - /// HaveSameOrder). + /// with a safe cast, provided the types are ordered alike (see HaveSameOrder). + /// Such a bound arises when a column is read as another type than it is stored + /// as: its statistics bound the stored values and the same cast is applied to + /// the data. A safe cast is nondecreasing but may round (between floating point + /// types), so unless the bound round trips, strict comparisons are weakened to + /// inclusive ones. This also holds if the guarantee is read as comparing the + /// target with the bound in their common type: no value of the target's type + /// lies strictly between the bound and its rounding. /// - The bound is then put through the casts the target goes through. Each cast is /// nondecreasing, so `x <= b` implies `cast(x) <= cast(b)` even if `cast(b)` /// rounds, and no exactness is needed. @@ -1354,6 +1365,7 @@ struct Inequality { if (!value->is_valid) return std::nullopt; ExecContext exec_context; + Comparison::type restated_cmp = cmp; const auto& target_type = casts.inner->type(); if (!value->type->Equals(*target_type)) { if (!HaveSameOrder(*value->type, *target_type)) return std::nullopt; @@ -1363,13 +1375,13 @@ struct Inequality { auto maybe_round_trip = Cast(*maybe_cast, value->type, CastOptions::Safe(), &exec_context); if (!maybe_round_trip.ok() || !maybe_round_trip->scalar()->Equals(*value)) { - return std::nullopt; + restated_cmp = WeakenStrictComparison(restated_cmp); } value = maybe_cast->scalar(); } auto cast_value = casts.ApplyToBound(std::move(value), &exec_context); if (!cast_value.ok()) return std::nullopt; - return std::make_pair(casts.ApplyToComparison(cmp), *std::move(cast_value)); + return std::make_pair(casts.ApplyToComparison(restated_cmp), *std::move(cast_value)); } /// The given expression simplifies to `value` if the inequality @@ -1631,15 +1643,15 @@ internal::ArgumentRange ResolveRange(const std::vector& inequalities const bool inclusive = cmp & Comparison::EQUAL; // Unless the comparison admits values below the bound, it is a lower bound; // likewise for values above the bound. Equality is both. - if (!(cmp & Comparison::LESS) && is_tighter(bound, inclusive, range.lower, - range.lower_inclusive, - Comparison::GREATER)) { + if (!(cmp & Comparison::LESS) && + is_tighter(bound, inclusive, range.lower, range.lower_inclusive, + Comparison::GREATER)) { range.lower = bound; range.lower_inclusive = inclusive; } - if (!(cmp & Comparison::GREATER) && is_tighter(bound, inclusive, range.upper, - range.upper_inclusive, - Comparison::LESS)) { + if (!(cmp & Comparison::GREATER) && + is_tighter(bound, inclusive, range.upper, range.upper_inclusive, + Comparison::LESS)) { range.upper = bound; range.upper_inclusive = inclusive; } @@ -1670,8 +1682,8 @@ Result SimplifyWithRanges(const FieldRanges& ranges, Expression expr const auto* evaluator = dynamic_cast(call->kernel_state.get()); if (evaluator && call->arguments.size() == 1 && !is_small_is_in) { - ARROW_ASSIGN_OR_RAISE( - auto evaluation, evaluator->EvaluateOverRange(ResolveRange(*inequalities, casts))); + ARROW_ASSIGN_OR_RAISE(auto evaluation, evaluator->EvaluateOverRange( + ResolveRange(*inequalities, casts))); switch (evaluation) { case internal::RangeEvaluation::kUnknown: return expr; diff --git a/cpp/src/arrow/compute/expression_test.cc b/cpp/src/arrow/compute/expression_test.cc index 58277d4767a3..8c97ef1e13f5 100644 --- a/cpp/src/arrow/compute/expression_test.cc +++ b/cpp/src/arrow/compute/expression_test.cc @@ -2110,12 +2110,31 @@ TEST(Expression, SimplifyWithBoundOfAnotherType) { Simplify{is_in("[5, 9]")}.WithGuarantee(range).Expect(is_in("[9]")); } - // ... while those which do not are ignored rather than rounded + // ... while those which a safe cast rejects are ignored auto inexact_range = and_(greater_equal(i64, literal(9.5)), less_equal(i64, literal(10.5))); - Simplify{equal(i64, literal(5))}.WithGuarantee(inexact_range).ExpectUnchanged(); + Simplify{equal(i64, literal(5))} + .WithGuarantee(inexact_range) + .ExpectUnchanged(); Simplify{is_in("[5, 9]")}.WithGuarantee(inexact_range).ExpectUnchanged(); + // Bounds which a safe cast rounds (the statistics of a float64 column read as + // float32, say) are used rounded, and then only bound inclusively: 0.2 rounds up + // to 0.2f, which a stored 0.2 is read as. + auto f32 = field_ref("f32"); + auto rounded_range = + and_(greater_equal(f32, literal(0.1)), less_equal(f32, literal(0.2))); + Simplify{greater(f32, literal(5.0F))}.WithGuarantee(rounded_range).Expect(false); + Simplify{less(f32, literal(0.05F))}.WithGuarantee(rounded_range).Expect(false); + Simplify{equal(f32, literal(0.2F))}.WithGuarantee(rounded_range).ExpectUnchanged(); + Simplify{greater(f32, literal(0.2F))}.WithGuarantee(rounded_range).Expect(false); + Simplify{equal(f32, literal(0.2F))} + .WithGuarantee(less(f32, literal(0.2))) + .ExpectUnchanged(); + Simplify{equal(f32, literal(0.1F))} + .WithGuarantee(greater(f32, literal(0.1))) + .ExpectUnchanged(); + // ... as are bounds of a type which is ordered differently, even though they cast // exactly: "10" <= x <= "9" admits "5" auto string_range = @@ -2463,8 +2482,8 @@ TEST(Expression, SimplifyIsInLargeValueSet) { {f64, float64(), range, and_(greater(f64, literal(50.0)), less_equal(f64, literal(51.0))), kUnchanged}, // 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). + // is rounded and only bounds inclusively: 0.1 rounds up to 0.1f, and a strict + // bound would wrongly exclude the member 0.1f (0.1f > 0.1 holds in double). {f32, float32(), concat_json("[0.1]", range), and_(greater(f32, literal(0.1)), less(f32, literal(0.5))), kUnchanged}, {f32, float32(), concat_json("[0.7]", range), less(f32, literal(0.7)), kUnchanged}, @@ -2559,10 +2578,9 @@ TEST(Expression, SimplifyIsInLargeValueSet) { .ExpectUnchanged(); Simplify{huge}.WithGuarantee(greater_equal(i32, literal(170000))).ExpectUnchanged(); // ... unless they contain NaN, which the smallest and largest value ignore - auto huge_nan = - call("is_in", {f64}, - SetLookupOptions{ArrayFromJSON( - float64(), concat_json("[NaN]", make_range_json(1, 140000)))}); + auto huge_nan = call("is_in", {f64}, + SetLookupOptions{ArrayFromJSON( + float64(), concat_json("[NaN]", make_range_json(1, 140000)))}); Simplify{huge_nan}.WithGuarantee(less(f64, literal(0.5))).ExpectUnchanged(); Simplify{huge_nan}.WithGuarantee(greater(f64, literal(140000.5))).ExpectUnchanged(); diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc index 097157c42702..a9d993ec5c7c 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc @@ -125,8 +125,9 @@ struct SetLookupStateBase : public KernelState, public RangeEvaluator { int64_t end = sorted->length(); if (range.lower) { // The first value not less than (or, for an exclusive bound, greater than) it - auto index = search(range.lower, range.lower_inclusive ? SearchSortedOptions::Left - : SearchSortedOptions::Right); + auto index = + search(range.lower, range.lower_inclusive ? SearchSortedOptions::Left + : SearchSortedOptions::Right); if (!index) return RangeEvaluation::kUnknown; begin = *index; } @@ -198,8 +199,8 @@ struct SetLookupStateBase : public KernelState, public RangeEvaluator { if (value_set.is_array()) { values = value_set.make_array(); } else { - ARROW_ASSIGN_OR_RAISE(values, Concatenate(value_set.chunked_array()->chunks(), - ctx->memory_pool())); + ARROW_ASSIGN_OR_RAISE( + values, Concatenate(value_set.chunked_array()->chunks(), ctx->memory_pool())); } if (values->length() > kIsInMaxSortedValueSet) { ARROW_ASSIGN_OR_RAISE(auto bounds, ValueSetBounds(values, ctx));