bvolpato commented on code in PR #24952:
URL: https://github.com/apache/datafusion/pull/24952#discussion_r3998888308
##########
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:
good catch. the test now covers four-item lists for Float16, Float32, and
Float64. I verified that removing the floating-point guard makes it fail. the
short-list execution coverage is still there.
--
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]
