diff --git a/cpp/src/arrow/compute/expression.cc b/cpp/src/arrow/compute/expression.cc index cd952055666..731ba0d42ee 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 @@ -1160,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 { @@ -1172,36 +1231,35 @@ 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; + // 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 = ExtractOneFromComparison(guarantee)) return {*inequality}; - out->nullable = true; - return out; + 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); + if (!inequality || inequality->target != *target) continue; + inequality->nullable = true; + inequalities.push_back(*inequality); } - - // fall back to a simple comparison with no "is_null" - return ExtractOneFromComparison(guarantee); + return inequalities; } static std::optional ExtractOneFromComparison(const Expression& guarantee) { @@ -1225,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. @@ -1270,21 +1342,26 @@ 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; } - 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->null_matching_behavior) { + switch (options->GetNullMatchingBehavior()) { case SetLookupOptions::MATCH: null_selection = guarantee.nullable ? FilterOptions::EMIT_NULL : FilterOptions::DROP; @@ -1305,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, @@ -1321,7 +1398,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), @@ -1330,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(); @@ -1358,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); @@ -1386,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. @@ -1403,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 @@ -1412,6 +1495,164 @@ 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.GetNullMatchingBehavior()) { + 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, 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(); + + 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. 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; + + // 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; + } + } + + // 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,33 +1696,64 @@ 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. + std::vector inequalities; for (const auto& guarantee : conjunction_members) { - if (!guarantee.call()) continue; - - 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)); - })); + 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; + } + } - if (Expression::Identical(simplified, expr)) continue; + // `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( + std::move(expr), [](Expression expr) { return expr; }, + [&](Expression expr, ...) -> Result { + return SimplifyIsInWithInequalities( + inequalities, std::move(expr)); + })); - expr = std::move(simplified); + if (!Expression::Identical(expr, before)) { RETURN_NOT_OK(CanonicalizeAndFoldConstants()); } + } - if (guarantee.call()->function_name == "is_valid") { - ARROW_ASSIGN_OR_RAISE( - auto simplified, - SimplifyIsValidGuarantee(std::move(expr), *CallNotNull(guarantee))); + 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; + 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; + + ARROW_ASSIGN_OR_RAISE(auto simplified, SimplifyIsValidGuarantee( + std::move(expr), *CallNotNull(guarantee))); + + if (Expression::Identical(simplified, expr)) continue; + + expr = std::move(simplified); + RETURN_NOT_OK(CanonicalizeAndFoldConstants()); } return expr; diff --git a/cpp/src/arrow/compute/expression_internal.h b/cpp/src/arrow/compute/expression_internal.h index 50b08a3d7a0..05b41c8d7f5 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 10c1b348e62..41e077a70a6 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"})} @@ -2003,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, @@ -2063,8 +2207,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 +2295,486 @@ TEST(Expression, SimplifyIsIn) { .ExpectUnchanged(); } +TEST(Expression, SimplifyIsInLargeValueSet) { + // Value sets longer than kIsInSimplificationMaxValueSet are binary searched for + // the range admitted by the inequalities of the guarantee. The call is simplified + // to false if the range is empty and is left unchanged otherwise. + auto is_in = [](Expression field, std::shared_ptr value_set_type, + std::string json_array, + SetLookupOptions::NullMatchingBehavior null_matching_behavior) { + SetLookupOptions options{ArrayFromJSON(value_set_type, json_array), + null_matching_behavior}; + return call("is_in", {field}, options); + }; + + auto u32 = field_ref("u32"); + auto i32 = field_ref("i32"); + auto i64 = field_ref("i64"); + auto f64 = field_ref("f64"); + auto str = field_ref("str"); + + // 100 contiguous values + auto range = make_range_json(1, 100); + // 101 values with a gap: [1, 50] and [100, 150] + auto gapped = concat_json(make_range_json(1, 50), make_range_json(100, 150)); + + for (SetLookupOptions::NullMatchingBehavior null_matching : { + SetLookupOptions::MATCH, + SetLookupOptions::SKIP, + SetLookupOptions::EMIT_NULL, + SetLookupOptions::INCONCLUSIVE, + }) { + // Guarantees overlapping the value set leave the call unchanged + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(greater(u32, literal(3))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(less_equal(u32, literal(1))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(100)), less_equal(u32, literal(200)))) + .ExpectUnchanged(); + + // Guarantees disjoint from the value set's hull + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(greater(u32, literal(100))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(greater_equal(u32, literal(101))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(less(u32, literal(1))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee(less_equal(u32, literal(0))) + .Expect(false); + Simplify{is_in(u32, int64(), range, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(101)), less_equal(u32, literal(200)))) + .Expect(false); + + // Guarantees inside the value set's hull but disjoint from its values. This + // requires considering both bounds of the guarantee at once. + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(and_(greater(u32, literal(50)), less(u32, literal(100)))) + .Expect(false); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(51)), less_equal(u32, literal(99)))) + .Expect(false); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(less_equal(u32, literal(99)), greater_equal(u32, literal(51)))) + .Expect(false); + // ... but a single value of the set in the range is enough to keep the call + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(50)), less_equal(u32, literal(99)))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee( + and_(greater_equal(u32, literal(51)), less_equal(u32, literal(100)))) + .ExpectUnchanged(); + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(greater(u32, literal(50))) + .ExpectUnchanged(); + // Redundant bounds on the same field are all applied + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(and_(and_(greater(u32, literal(10)), greater(u32, literal(50))), + less(u32, literal(100)))) + .Expect(false); + + // Guarantees on other fields are irrelevant + Simplify{is_in(u32, int64(), gapped, null_matching)} + .WithGuarantee(and_(greater(i32, literal(50)), less(i32, literal(100)))) + .ExpectUnchanged(); + + // Casts: the value set is cast to the field type when the call is bound and + // the bounds of the guarantee are cast to the value set type. + Simplify{is_in(i64, int32(), range, null_matching)} + .WithGuarantee(greater(i64, literal(100))) + .Expect(false); + Simplify{is_in(i64, int32(), range, null_matching)} + .WithGuarantee(greater(i64, literal(100))) + .Expect(false); + Simplify{is_in(cast(i32, int64()), int64(), range, null_matching)} + .WithGuarantee(greater(i32, literal(100))) + .Expect(false); + + // Unsorted value sets with duplicates are sorted when the call is bound + std::string unsorted = "["; + for (int i = 100; i >= 1; --i) { + unsorted += std::to_string(i) + "," + std::to_string(i) + ","; + } + unsorted.back() = ']'; + Simplify{is_in(u32, int64(), unsorted, null_matching)} + .WithGuarantee(greater(u32, literal(100))) + .Expect(false); + Simplify{is_in(u32, int64(), unsorted, null_matching)} + .WithGuarantee(and_(greater(u32, literal(50)), less(u32, literal(51)))) + .Expect(false); + Simplify{is_in(u32, int64(), unsorted, null_matching)} + .WithGuarantee(and_(greater_equal(u32, literal(50)), less(u32, literal(51)))) + .ExpectUnchanged(); + + // Strings + std::string strings = "["; + for (int i = 0; i < 60; ++i) { + if (i > 0) strings += ","; + strings += "\"k" + std::to_string(100 + i) + "\""; + } + strings += "]"; + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(greater(str, literal("k159"))) + .Expect(false); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(less(str, literal("k100"))) + .Expect(false); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(and_(greater(str, literal("k110")), less(str, literal("k111")))) + .Expect(false); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee( + and_(greater_equal(str, literal("k110")), less(str, literal("k111")))) + .ExpectUnchanged(); + Simplify{is_in(str, utf8(), strings, null_matching)} + .WithGuarantee(greater(str, literal("k1"))) + .ExpectUnchanged(); + // Unsorted strings with duplicates take the generic sorting path + std::string unsorted_strings = "["; + for (int i = 59; i >= 0; --i) { + unsorted_strings += "\"k" + std::to_string(100 + i) + "\","; + unsorted_strings += "\"k" + std::to_string(100 + i) + "\","; + } + unsorted_strings.back() = ']'; + Simplify{is_in(str, utf8(), unsorted_strings, null_matching)} + .WithGuarantee(and_(greater(str, literal("k110")), less(str, literal("k111")))) + .Expect(false); + Simplify{is_in(str, utf8(), unsorted_strings, null_matching)} + .WithGuarantee( + and_(greater(str, literal("k110")), less_equal(str, literal("k111")))) + .ExpectUnchanged(); + // ... as do value sets with nulls, which INCONCLUSIVE cannot simplify + auto with_null = Simplify{is_in(str, utf8(), concat_json(unsorted_strings, "[null]"), + null_matching)} + .WithGuarantee(and_(greater(str, literal("k110")), + less(str, literal("k111")))); + if (null_matching == SetLookupOptions::INCONCLUSIVE) { + with_null.ExpectUnchanged(); + } else { + with_null.Expect(false); + } + + // Floats + Simplify{is_in(f64, float64(), range, null_matching)} + .WithGuarantee(greater(f64, literal(100.5))) + .Expect(false); + Simplify{is_in(f64, float64(), range, null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.0)), less(f64, literal(51.0)))) + .Expect(false); + Simplify{is_in(f64, float64(), range, null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.0)), less_equal(f64, literal(51.0)))) + .ExpectUnchanged(); + // A double bound which does not round trip through the float32 value set type + // is ignored rather than rounded: 0.1 rounds up to 0.1f, which would wrongly + // exclude the member 0.1f (0.1f > 0.1 holds in double). + auto f32 = field_ref("f32"); + Simplify{is_in(f32, float32(), concat_json("[0.1]", range), null_matching)} + .WithGuarantee(and_(greater(f32, literal(0.1)), less(f32, literal(0.5)))) + .ExpectUnchanged(); + Simplify{is_in(f32, float32(), concat_json("[0.7]", range), null_matching)} + .WithGuarantee(less(f32, literal(0.7))) + .ExpectUnchanged(); + Simplify{is_in(f32, float32(), concat_json("[0.5]", range), null_matching)} + .WithGuarantee(less(f32, literal(0.5))) + .Expect(false); + Simplify{is_in(f32, float32(), concat_json("[0.5]", range), null_matching)} + .WithGuarantee(and_(greater(f32, literal(0.25)), less(f32, literal(0.5)))) + .Expect(false); + // Statistics cannot bound NaN, so a value set containing NaN is never simplified + Simplify{is_in(f64, float64(), concat_json("[NaN]", range), null_matching)} + .WithGuarantee(greater(f64, literal(100.5))) + .ExpectUnchanged(); + Simplify{is_in(f64, float64(), concat_json(range, "[NaN]"), null_matching)} + .WithGuarantee(greater(f64, literal(100.5))) + .ExpectUnchanged(); + + // Strict bounds through a stripped narrowing cast are weakened to inclusive + // ones: f64 = 51.00000001 satisfies f64 > 51 and casts to 51f, a set member. + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(greater(f64, literal(51.0))) + .ExpectUnchanged(); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(less(f64, literal(1.0))) + .ExpectUnchanged(); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(greater(f64, literal(51.5))) + .Expect(false); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.0)), less(f64, literal(51.0)))) + .ExpectUnchanged(); + Simplify{ + is_in(cast(f64, float32()), float32(), make_range_json(1, 51), null_matching)} + .WithGuarantee(and_(greater(f64, literal(50.25)), less(f64, literal(50.75)))) + .Expect(false); + + // 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, 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); + } + // ... 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)); + 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)), @@ -2178,6 +2802,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 c6a3562233c..676661b0287 100644 --- a/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup.cc @@ -16,9 +16,11 @@ // under the License. #include "arrow/array/array_base.h" +#include "arrow/array/concatenate.h" #include "arrow/compute/api_scalar.h" #include "arrow/compute/cast.h" #include "arrow/compute/kernels/common_internal.h" +#include "arrow/compute/kernels/scalar_set_lookup_internal.h" #include "arrow/compute/kernels/util_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/type.h" @@ -36,11 +38,6 @@ using internal::HashTraits; namespace compute::internal { namespace { -// This base class enables non-templated access to the value set type -struct SetLookupStateBase : public KernelState { - std::shared_ptr value_set_type; -}; - template struct SetLookupState : public SetLookupStateBase { explicit SetLookupState(MemoryPool* pool) : memory_pool(pool) {} @@ -69,6 +66,7 @@ struct SetLookupState : public SetLookupStateBase { } else { return Status::Invalid("value_set should be an array or chunked array"); } + value_set_has_null = lookup_table->GetNull() >= 0; if (this->null_matching_behavior != SetLookupOptions::SKIP && lookup_table->GetNull() >= 0) { null_index = memo_index_to_value_index[lookup_table->GetNull()]; @@ -130,13 +128,12 @@ struct SetLookupState : public SetLookupStateBase { Status Init(SetLookupOptions& options) { null_matching_behavior = options.GetNullMatchingBehavior(); - value_set_has_null = (options.value_set.length() > 0) && - this->null_matching_behavior != SetLookupOptions::SKIP; + // The kernels consult null_matching_behavior before using this. + value_set_has_null = options.value_set.length() > 0; value_set_type = null(); return Status::OK(); } - bool value_set_has_null; SetLookupOptions::NullMatchingBehavior null_matching_behavior; }; @@ -171,10 +168,15 @@ struct InitStateVisitor { TypeHolder arg_type; std::unique_ptr result; - InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args) + // Whether to populate SetLookupStateBase::sorted_value_set (is_in only). + bool compute_sorted_value_set; + + InitStateVisitor(KernelContext* ctx, const KernelInitArgs& args, + bool compute_sorted_value_set) : ctx(ctx), options(*checked_cast(args.options)), - arg_type(args.inputs[0]) {} + arg_type(args.inputs[0]), + compute_sorted_value_set(compute_sorted_value_set) {} template Status Init() { @@ -254,18 +256,49 @@ struct InitStateVisitor { } RETURN_NOT_OK(VisitTypeInline(*options.value_set.type(), this)); + + if (compute_sorted_value_set && + options.value_set.length() > kIsInSimplificationMaxValueSet && + IsSearchableValueSetType(*options.value_set.type())) { + // Best effort: a missing sorted copy only disables the simplification of + // `is_in` against guarantees, see SetLookupStateBase::GetSortedValueSet. + auto maybe_values = MaterializedValueSet(); + if (maybe_values.ok()) { + ARROW_UNUSED(checked_cast(result.get()) + ->SetSearchableValueSet(*std::move(maybe_values))); + } + } return std::move(result); } + + Result> MaterializedValueSet() { + if (options.value_set.is_array()) { + return options.value_set.make_array(); + } + return Concatenate(options.value_set.chunked_array()->chunks(), + ctx->exec_context()->memory_pool()); + } }; Result> InitSetLookup(KernelContext* ctx, - const KernelInitArgs& args) { + const KernelInitArgs& args, + bool compute_sorted_value_set) { if (args.options == nullptr) { return Status::Invalid( "Attempted to call a set lookup function without SetLookupOptions"); } - return InitStateVisitor{ctx, args}.GetResult(); + return InitStateVisitor{ctx, args, compute_sorted_value_set}.GetResult(); +} + +Result> InitIsIn(KernelContext* ctx, + const KernelInitArgs& args) { + return InitSetLookup(ctx, args, /*compute_sorted_value_set=*/true); +} + +Result> InitIndexIn(KernelContext* ctx, + const KernelInitArgs& args) { + return InitSetLookup(ctx, args, /*compute_sorted_value_set=*/false); } struct IndexInVisitor { @@ -648,7 +681,7 @@ void RegisterScalarSetLookup(FunctionRegistry* registry) { // IsIn writes its boolean output into preallocated memory { ScalarKernel isin_base; - isin_base.init = InitSetLookup; + isin_base.init = InitIsIn; isin_base.exec = ExecIsIn; isin_base.null_handling = NullHandling::COMPUTED_PREALLOCATE; auto is_in = std::make_shared("is_in", Arity::Unary(), is_in_doc); @@ -665,7 +698,7 @@ void RegisterScalarSetLookup(FunctionRegistry* registry) { // IndexIn writes its int32 output into preallocated memory { ScalarKernel index_in_base; - index_in_base.init = InitSetLookup; + index_in_base.init = InitIndexIn; index_in_base.exec = ExecIndexIn; index_in_base.null_handling = NullHandling::COMPUTED_PREALLOCATE; auto index_in = diff --git a/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h new file mode 100644 index 00000000000..60011e4a676 --- /dev/null +++ b/cpp/src/arrow/compute/kernels/scalar_set_lookup_internal.h @@ -0,0 +1,386 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +// Shared between the set lookup kernels (libarrow_compute) and expression +// simplification (libarrow). Everything here must stay header-only. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array/array_base.h" +#include "arrow/array/array_binary.h" +#include "arrow/array/array_decimal.h" +#include "arrow/array/array_primitive.h" +#include "arrow/array/concatenate.h" +#include "arrow/array/data.h" +#include "arrow/array/util.h" +#include "arrow/buffer.h" +#include "arrow/compute/api_vector.h" +#include "arrow/compute/kernel.h" +#include "arrow/datum.h" +#include "arrow/result.h" +#include "arrow/scalar.h" +#include "arrow/status.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/decimal.h" +#include "arrow/util/logging.h" +#include "arrow/visit_type_inline.h" + +namespace arrow::compute::internal { + +/// Value sets up to this length are simplified against a guarantee by filtering +/// the value set and rebinding the `is_in` call (see SimplifyIsIn in expression.cc). +/// Longer value sets are instead binary searched through SetLookupStateBase:: +/// sorted_value_set, which only ever simplifies the call to `false`, see +/// https://github.com/apache/arrow/issues/46777 +constexpr int64_t kIsInSimplificationMaxValueSet = 50; + +/// Whether a (valid) scalar of a floating point type is NaN. +inline bool IsNaNScalar(const Scalar& value) { + if (value.type->id() == Type::FLOAT) { + return std::isnan(::arrow::internal::checked_cast(value).value); + } + if (value.type->id() == Type::DOUBLE) { + return std::isnan(::arrow::internal::checked_cast(value).value); + } + return false; +} + +inline Result IsSortedValueSet(const Array& values); +inline Result> SortPrimitiveValueSet(const Array& values); + +/// Common state for the is_in and index_in kernels. +struct SetLookupStateBase : public KernelState { + std::shared_ptr value_set_type; + + /// Whether the value set contains a null (regardless of null matching behavior). + bool value_set_has_null = false; + + /// Make the value set available to GetSortedValueSet(). + /// + /// Called once by the `is_in` kernel init, before the state is shared. If the + /// value set is sorted already it is used as is; otherwise sorting is deferred + /// to the first call of GetSortedValueSet() so that only expressions which are + /// simplified against guarantees pay for it. + Status SetSearchableValueSet(std::shared_ptr values) { + if (values->null_count() == 0) { + ARROW_ASSIGN_OR_RAISE(bool sorted, IsSortedValueSet(*values)); + if (sorted) { + sorted_value_set_ = std::move(values); + sorted_value_set_ready_.store(true, std::memory_order_release); + return Status::OK(); + } + } + unsorted_value_set_ = std::move(values); + return Status::OK(); + } + + /// A copy of the value set sorted ascending in the order of the comparison + /// kernels, with nulls removed, or nullptr if there is none. + /// + /// This is only available for `is_in` when the value set is longer than + /// kIsInSimplificationMaxValueSet, has a type for which IsSearchableValueSetType() + /// holds and contains no NaN. It is computed at most once; concurrent calls are + /// safe. + std::shared_ptr GetSortedValueSet() const { + if (!sorted_value_set_ready_.load(std::memory_order_acquire)) { + std::lock_guard lock(sorted_value_set_mutex_); + if (!sorted_value_set_ready_.load(std::memory_order_relaxed)) { + if (unsorted_value_set_) { + // Best effort: a failure only disables the simplification of `is_in` + // against guarantees. + sorted_value_set_ = SortValueSet(*unsorted_value_set_).ValueOr(nullptr); + unsorted_value_set_.reset(); + } + sorted_value_set_ready_.store(true, std::memory_order_release); + } + } + return sorted_value_set_; + } + + private: + /// \return `values` sorted with nulls removed, or nullptr if it contains NaN + static Result> SortValueSet(const Array& values) { + // Sorting a copy of primitive values directly is several times faster than + // the generic sort_indices + take path. + auto maybe_sorted = SortPrimitiveValueSet(values); + if (maybe_sorted.ok()) return maybe_sorted; + if (!maybe_sorted.status().IsNotImplemented()) return maybe_sorted.status(); + + ARROW_ASSIGN_OR_RAISE(auto indices, + SortIndices(values, ArraySortOptions(SortOrder::Ascending, + NullPlacement::AtEnd))); + ARROW_ASSIGN_OR_RAISE(Datum taken, + Take(values, indices, TakeOptions::NoBoundsCheck())); + auto sorted = taken.make_array(); + sorted = sorted->Slice(0, sorted->length() - sorted->null_count()); + // NaN sorts after every other value and before nulls. Statistics based + // guarantees cannot bound NaN, so such a value set is not searchable. + if (sorted->length() > 0 && is_floating(sorted->type_id())) { + ARROW_ASSIGN_OR_RAISE(auto last, sorted->GetScalar(sorted->length() - 1)); + if (IsNaNScalar(*last)) return std::shared_ptr(nullptr); + } + return sorted; + } + + mutable std::mutex sorted_value_set_mutex_; + mutable std::atomic sorted_value_set_ready_{false}; + mutable std::shared_ptr unsorted_value_set_; + mutable std::shared_ptr sorted_value_set_; +}; + +/// Whether value sets of this type can be binary searched with a comparator that +/// agrees with both the sorting kernels and the comparison kernels. +inline bool IsSearchableValueSetType(const DataType& type) { + switch (type.id()) { + case Type::INT8: + case Type::INT16: + case Type::INT32: + case Type::INT64: + case Type::UINT8: + case Type::UINT16: + case Type::UINT32: + case Type::UINT64: + case Type::FLOAT: + case Type::DOUBLE: + case Type::DATE32: + case Type::DATE64: + case Type::TIME32: + case Type::TIME64: + case Type::TIMESTAMP: + case Type::DURATION: + case Type::DECIMAL32: + case Type::DECIMAL64: + case Type::DECIMAL128: + case Type::DECIMAL256: + case Type::STRING: + case Type::LARGE_STRING: + case Type::BINARY: + case Type::LARGE_BINARY: + case Type::FIXED_SIZE_BINARY: + return true; + default: + return false; + } +} + +namespace set_lookup_detail { + +template +struct ValueSetAccessor; + +template +struct ValueSetAccessor< + T, enable_if_t::value && !is_boolean_type::value && + !is_half_float_type::value && !is_interval_type::value>> { + using ArrayType = typename TypeTraits::ArrayType; + using ScalarType = typename TypeTraits::ScalarType; + using ValueType = typename T::c_type; + + static ValueType Get(const Array& values, int64_t i) { + return ::arrow::internal::checked_cast(values).Value(i); + } + static ValueType FromScalar(const Scalar& scalar) { + return ::arrow::internal::checked_cast(scalar).value; + } +}; + +template +struct ValueSetAccessor> { + using ArrayType = typename TypeTraits::ArrayType; + using ScalarType = typename TypeTraits::ScalarType; + using ValueType = typename ScalarType::ValueType; + + static ValueType Get(const Array& values, int64_t i) { + return ValueType( + ::arrow::internal::checked_cast(values).GetValue(i)); + } + static ValueType FromScalar(const Scalar& scalar) { + return ::arrow::internal::checked_cast(scalar).value; + } +}; + +template +struct ValueSetAccessor::value || + std::is_same::value>> { + using ArrayType = typename TypeTraits::ArrayType; + using ValueType = std::string_view; + + static ValueType Get(const Array& values, int64_t i) { + return ::arrow::internal::checked_cast(values).GetView(i); + } + static ValueType FromScalar(const Scalar& scalar) { + return ::arrow::internal::checked_cast(scalar).view(); + } +}; + +template +struct is_searchable_value_set_type : std::false_type {}; + +template +struct is_searchable_value_set_type< + T, std::void_t::Get(std::declval(), 0))>> + : std::true_type {}; + +template +bool IsNaNValue(const ValueType&) { + return false; +} +inline bool IsNaNValue(float value) { return std::isnan(value); } +inline bool IsNaNValue(double value) { return std::isnan(value); } + +struct IsSortedVisitor { + const Array& values; + bool result = false; + + template + enable_if_t::value, Status> Visit(const T&) { + using Accessor = ValueSetAccessor; + if (values.length() == 0) { + result = true; + return Status::OK(); + } + auto previous = Accessor::Get(values, 0); + if (IsNaNValue(previous)) return Status::OK(); + for (int64_t i = 1; i < values.length(); ++i) { + auto current = Accessor::Get(values, i); + if (IsNaNValue(current) || current < previous) return Status::OK(); + previous = current; + } + result = true; + return Status::OK(); + } + + Status Visit(const DataType& type) { + return Status::NotImplemented("Searching value sets of type ", type); + } +}; + +struct SortPrimitiveVisitor { + const Array& values; + std::shared_ptr result; + + template + enable_if_t::value && is_searchable_value_set_type::value, Status> + Visit(const T&) { + using c_type = typename T::c_type; + const auto* data = values.data()->template GetValues(1); + std::vector sorted; + sorted.reserve(values.length() - values.null_count()); + if (values.null_count() == 0) { + sorted.assign(data, data + values.length()); + } else { + for (int64_t i = 0; i < values.length(); ++i) { + if (values.IsValid(i)) sorted.push_back(data[i]); + } + } + // NaN would break the strict weak ordering required by std::sort. Statistics + // based guarantees cannot bound NaN anyway, so such a value set is not + // searchable (nullptr). + for (const auto& value : sorted) { + if (IsNaNValue(value)) return Status::OK(); + } + std::sort(sorted.begin(), sorted.end()); + const auto length = static_cast(sorted.size()); + result = MakeArray(ArrayData::Make(values.type(), length, + {nullptr, Buffer::FromVector(std::move(sorted))}, + /*null_count=*/0)); + return Status::OK(); + } + + Status Visit(const DataType& type) { + return Status::NotImplemented("Sorting value sets of type ", type); + } +}; + +struct SearchSortedVisitor { + const Array& sorted; + const Scalar& bound; + bool upper; + int64_t result = 0; + + template + enable_if_t::value, Status> Visit(const T&) { + using Accessor = ValueSetAccessor; + const auto target = Accessor::FromScalar(bound); + int64_t lo = 0, hi = sorted.length(); + while (lo < hi) { + const int64_t mid = lo + (hi - lo) / 2; + const auto value = Accessor::Get(sorted, mid); + const bool go_right = upper ? !(target < value) : (value < target); + if (go_right) { + lo = mid + 1; + } else { + hi = mid; + } + } + result = lo; + return Status::OK(); + } + + Status Visit(const DataType& type) { + return Status::NotImplemented("Searching value sets of type ", type); + } +}; + +} // namespace set_lookup_detail + +/// Whether `values` (which must contain no nulls) is sorted ascending in the order +/// of the comparison kernels and contains no NaN. +inline Result IsSortedValueSet(const Array& values) { + ARROW_DCHECK_EQ(values.null_count(), 0); + set_lookup_detail::IsSortedVisitor visitor{values, /*result=*/false}; + RETURN_NOT_OK(VisitTypeInline(*values.type(), &visitor)); + return visitor.result; +} + +/// Sort a value set of a primitive type, dropping nulls. +/// +/// \return the sorted value set, or nullptr if it contains NaN +/// \return NotImplemented for non-primitive types +inline Result> SortPrimitiveValueSet(const Array& values) { + set_lookup_detail::SortPrimitiveVisitor visitor{values, /*result=*/nullptr}; + RETURN_NOT_OK(VisitTypeInline(*values.type(), &visitor)); + return visitor.result; +} + +/// Binary search a sorted value set for `bound`. +/// +/// \pre `sorted` satisfies IsSortedValueSet() +/// \pre `bound` is valid, not NaN, and has the same type as `sorted` +/// \return the index of the first element not less than `bound` (`upper == false`), +/// or of the first element greater than `bound` (`upper == true`) +inline Result SearchSortedValueSet(const Array& sorted, const Scalar& bound, + bool upper) { + ARROW_DCHECK(bound.is_valid); + ARROW_DCHECK(sorted.type()->Equals(*bound.type)); + set_lookup_detail::SearchSortedVisitor visitor{sorted, bound, upper, /*result=*/0}; + RETURN_NOT_OK(VisitTypeInline(*sorted.type(), &visitor)); + return visitor.result; +} + +} // namespace arrow::compute::internal diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 0d86f5eaa0d..f38ea5f036c 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -19,6 +19,8 @@ #include #include +#include +#include #include #include #include @@ -34,6 +36,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 +750,164 @@ 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, 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); + 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}); + // 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); + + // 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) { constexpr int64_t kNumRowGroups = 16; constexpr int64_t kTotalNumRows = kNumRowGroups * (kNumRowGroups + 1) / 2;