adamreeve commented on code in PR #55293:
URL: https://github.com/apache/arrow/pull/55293#discussion_r4533846004
##########
cpp/src/arrow/compute/expression_test.cc:
##########
@@ -2003,6 +2039,181 @@ 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]"));
+
+ // 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)))
+ .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, 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<int64_t>(5))}.WithGuarantee(range).Expect(false);
+ Simplify{greater(i64,
literal<int64_t>(5))}.WithGuarantee(range).Expect(true);
+ Simplify{is_in("[5, 9]")}.WithGuarantee(range).Expect(is_in("[9]"));
+ }
+
+ // ... 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<int64_t>(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 =
+ and_(greater_equal(i64, literal("10")), less_equal(i64, literal("9")));
Review Comment:
What is going on here, is it really valid to compare an integer to a string?
This comment says "as are", but I'm confused about what this is referring
to. The above test cases use rounding, but then these test cases all expect no
change after simplification.
##########
cpp/src/arrow/compute/expression.cc:
##########
@@ -1270,21 +1410,35 @@ struct Inequality {
auto options = checked_pointer_cast<SetLookupOptions>(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;
+ // Large value sets are handled by the kernel instead, see
SimplifyWithRanges.
if (options->value_set.length() > 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 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 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)
{
Review Comment:
I think this should include HALF_FLOAT too.
##########
cpp/src/arrow/compute/expression_test.cc:
##########
@@ -2003,6 +2039,181 @@ 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]"));
+
+ // 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)))
+ .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, 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<int64_t>(5))}.WithGuarantee(range).Expect(false);
+ Simplify{greater(i64,
literal<int64_t>(5))}.WithGuarantee(range).Expect(true);
+ Simplify{is_in("[5, 9]")}.WithGuarantee(range).Expect(is_in("[9]"));
+ }
+
+ // ... 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<int64_t>(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
Review Comment:
Is this something that could happen in practice?
##########
cpp/src/arrow/compute/kernels/scalar_set_lookup.cc:
##########
@@ -36,9 +44,217 @@ using internal::HashTraits;
namespace compute::internal {
namespace {
-// This base class enables non-templated access to the value set type
-struct SetLookupStateBase : public KernelState {
+/// 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<const FloatScalar&>(value).value);
+ case Type::DOUBLE:
+ return std::isnan(checked_cast<const DoubleScalar&>(value).value);
+ default:
Review Comment:
This should include HALF_FLOAT.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
