kosiew commented on code in PR #24952:
URL: https://github.com/apache/datafusion/pull/24952#discussion_r3977999183
##########
datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs:
##########
@@ -4890,6 +4921,186 @@ mod tests {
assert_eq!(simplify(expr.clone()), expr);
}
+ fn assert_inlist_simplification_result(
+ expr: Expr,
+ batch: &RecordBatch,
+ expected: Vec<bool>,
+ ) -> Result<()> {
+ let schema = batch.schema().to_dfschema_ref()?;
+ let simplifier = ExprSimplifier::new(
+ SimplifyContext::builder()
+ .with_schema(Arc::clone(&schema))
+ .build(),
+ );
+ let original = simplifier.coerce(expr, &schema)?;
+ let simplified = simplifier.simplify(original.clone())?;
+ let props = ExecutionProps::new();
+ let evaluate = |expr: &Expr| {
+ create_physical_expr(
+ expr,
+ &schema,
+ &props,
+ &PhysicalPlanningContext::default(),
+ )?
+ .evaluate(batch)?
+ .into_array(batch.num_rows())
+ };
+ let original_result = evaluate(&original)?;
+ let actual = evaluate(&simplified)?;
+ assert_eq!(original_result.as_ref(), actual.as_ref());
+ assert_eq!(actual.as_boolean(), &BooleanArray::from(expected));
+ Ok(())
+ }
+
+ #[test]
+ fn simplify_inlist_runtime_set_operations() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("x", DataType::Int32, false),
+ Field::new("a", DataType::Int32, false),
+ Field::new("b", DataType::Int32, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2, 3])),
+ Arc::new(Int32Array::from(vec![1, 1, 3])),
+ Arc::new(Int32Array::from(vec![1, 2, 4])),
+ ],
+ )?;
+ let left = |negated| {
+ in_list(col("x"), vec![col("a"), lit(2), lit(10), lit(11)],
negated)
+ };
+ let right = |negated| {
+ in_list(col("x"), vec![col("b"), lit(5), lit(12), lit(13)],
negated)
+ };
+
+ // Structurally different, non-null columns can hold equal values.
+ for (expr, expected) in [
+ (left(false).and(right(false)), vec![true, true, false]),
+ (left(false).and(right(true)), vec![false, false, true]),
+ (left(true).and(right(false)), vec![false, false, false]),
+ (left(true).or(right(true)), vec![false, false, true]),
+ // Union remains valid for nonvolatile runtime expressions.
+ (left(false).or(right(false)), vec![true, true, true]),
+ (left(true).and(right(true)), vec![false, false, false]),
+ ] {
+ assert_inlist_simplification_result(expr, &batch, expected)?;
+ }
+
+ // Coercion must precede structural comparison of mixed integer
literals.
+ let expr = in_list(col("x"), vec![lit(1i32), lit(2), lit(10),
lit(11)], false)
+ .and(in_list(
+ col("x"),
+ vec![lit(1i64), lit(3i64), lit(12i64), lit(13i64)],
+ false,
+ ));
+ assert_inlist_simplification_result(expr, &batch, vec![true, false,
false])
+ }
+
+ #[test]
+ fn simplify_inlist_signed_zero() -> Result<()> {
+ for (data_type, values, positive, negative, one, two) in [
+ (
+ DataType::Float32,
+ Arc::new(Float32Array::from(vec![0.0, -0.0])) as
arrow::array::ArrayRef,
+ lit(0.0f32),
+ lit(-0.0f32),
+ lit(1.0f32),
+ lit(2.0f32),
+ ),
+ (
+ DataType::Float64,
+ Arc::new(Float64Array::from(vec![0.0, -0.0])) as
arrow::array::ArrayRef,
+ lit(0.0f64),
+ lit(-0.0f64),
+ lit(1.0f64),
+ lit(2.0f64),
+ ),
+ ] {
+ let schema = Arc::new(Schema::new(vec![Field::new("x", data_type,
false)]));
+ let batch = RecordBatch::try_new(schema, vec![values])?;
+ let schema = batch.schema().to_dfschema_ref()?;
+ let simplifier = ExprSimplifier::new(
+ SimplifyContext::builder()
+ .with_schema(Arc::clone(&schema))
+ .build(),
Review Comment:
Could we make the signed-zero lists exceed `THRESHOLD_INLINE_INLIST`, and
ideally add coverage for `Float16` too? Right now the two-item lists get
expanded into comparisons before reaching the set-rewrite arms, so this test
would still pass if `inlists_have_set_comparable_literals` were removed. Using
a four-item case would exercise the new guard directly.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
