From 230a82daca335cf08184f2da6592132453f36e29 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Fri, 4 Sep 2026 23:11:25 -0400 Subject: [PATCH] fix: avoid unsafe set rewrites for compound IN lists --- .../simplify_expressions/expr_simplifier.rs | 215 +++++++++++++++++- .../sqllogictest/test_files/predicates.slt | 21 ++ 2 files changed, 234 insertions(+), 2 deletions(-) diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index ec01becb8f7a5..e48ecfe658dfd 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -1822,6 +1822,9 @@ impl TreeNodeRewriter for Simplifier<'_> { right.as_ref(), false, false, + ) && inlists_have_set_comparable_literals( + left.as_ref(), + right.as_ref(), ) => { match (*left, *right) { @@ -1862,6 +1865,9 @@ impl TreeNodeRewriter for Simplifier<'_> { right.as_ref(), false, true, + ) && inlists_have_set_comparable_literals( + left.as_ref(), + right.as_ref(), ) => { match (*left, *right) { @@ -1882,6 +1888,9 @@ impl TreeNodeRewriter for Simplifier<'_> { right.as_ref(), true, false, + ) && inlists_have_set_comparable_literals( + left.as_ref(), + right.as_ref(), ) => { match (*left, *right) { @@ -1902,6 +1911,9 @@ impl TreeNodeRewriter for Simplifier<'_> { right.as_ref(), true, true, + ) && inlists_have_set_comparable_literals( + left.as_ref(), + right.as_ref(), ) => { match (*left, *right) { @@ -2185,12 +2197,30 @@ fn are_inlist_and_eq_and_match_neg( ) -> bool { match (left, right) { (Expr::InList(l), Expr::InList(r)) => { - l.expr == r.expr && l.negated == is_left_neg && r.negated == is_right_neg + l.expr == r.expr + && !l.expr.is_volatile() + && !l.list.iter().chain(&r.list).any(Expr::is_volatile) + && l.negated == is_left_neg + && r.negated == is_right_neg } _ => false, } } +/// Structural equality can determine set membership only for non-null literals +/// whose equality agrees with runtime comparisons. Different runtime expressions +/// can evaluate to the same value, and floating-point equality can differ from +/// literal equality (for example, positive and negative zero). +fn inlists_have_set_comparable_literals(left: &Expr, right: &Expr) -> bool { + match (left, right) { + (Expr::InList(l), Expr::InList(r)) => l.list.iter().chain(&r.list).all(|item| { + item.as_literal() + .is_some_and(|value| !value.is_null() && !value.data_type().is_floating()) + }), + _ => false, + } +} + // TODO: We might not need this after defer pattern for Box is stabilized. https://github.com/rust-lang/rust/issues/87121 fn are_inlist_and_eq(left: &Expr, right: &Expr) -> bool { let left = as_inlist(left); @@ -2199,6 +2229,7 @@ fn are_inlist_and_eq(left: &Expr, right: &Expr) -> bool { matches!(lhs.expr.as_ref(), Expr::Column(_)) && matches!(rhs.expr.as_ref(), Expr::Column(_)) && lhs.expr == rhs.expr + && !lhs.list.iter().chain(&rhs.list).any(Expr::is_volatile) && !lhs.negated && !rhs.negated } else { @@ -2348,7 +2379,7 @@ mod tests { use super::*; use crate::test::test_table_scan_with_name; use arrow::{ - array::{Int32Array, StructArray}, + array::{BooleanArray, Float32Array, Float64Array, Int32Array, StructArray}, datatypes::{FieldRef, Fields}, }; use datafusion_common::{DFSchemaRef, ToDFSchema, assert_contains}; @@ -4672,6 +4703,186 @@ mod tests { assert_eq!(simplify(expr.clone()), expr); } + fn assert_inlist_simplification_result( + expr: Expr, + batch: &RecordBatch, + expected: Vec, + ) -> 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(), + ); + let left = + |negated| in_list(col("x"), vec![positive.clone(), one.clone()], negated); + let right = + |negated| in_list(col("x"), vec![negative.clone(), two.clone()], negated); + for (expr, expected) in [ + (left(false).and(right(false)), vec![true, true]), + (left(false).and(right(true)), vec![false, false]), + (left(true).or(right(true)), vec![false, false]), + ] { + // Short lists are expanded to SQL comparisons, where both zero + // representations compare equal. Set reduction must not run first. + let expr = simplifier.coerce(expr, &schema)?; + let simplified = simplifier.simplify(expr)?; + let actual = create_physical_expr( + &simplified, + &schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )? + .evaluate(&batch)? + .into_array(batch.num_rows())?; + assert_eq!(actual.as_boolean(), &BooleanArray::from(expected)); + } + } + Ok(()) + } + + #[test] + fn simplify_inlist_preserves_volatile_evaluations() { + let fun = Arc::new(ScalarUDF::new_from_impl(VolatileUdf::new())); + let volatile = Expr::ScalarFunction(ScalarFunction::new_udf(fun, vec![])); + let left = |negated| { + in_list( + volatile.clone(), + vec![lit(1), lit(2), lit(3), lit(4)], + negated, + ) + }; + let right = |negated| { + in_list( + volatile.clone(), + vec![lit(3), lit(4), lit(5), lit(6)], + negated, + ) + }; + for expr in [ + left(false).and(right(false)), + left(false).and(right(true)), + left(true).and(right(false)), + left(true).or(right(true)), + left(true).and(right(true)), + ] { + assert_eq!(simplify_no_canonicalize(expr.clone()), expr); + } + + // Union deduplication must not remove a repeated volatile list item. + for negated in [false, true] { + let left = in_list( + col("c1"), + vec![volatile.clone(), lit(1), lit(2), lit(3)], + negated, + ); + let right = in_list( + col("c1"), + vec![volatile.clone(), lit(4), lit(5), lit(6)], + negated, + ); + let expr = if negated { + left.and(right) + } else { + left.or(right) + }; + assert_eq!(simplify_no_canonicalize(expr.clone()), expr); + } + } + #[test] fn simplify_null_in_empty_inlist() { // `NULL::boolean IN ()` == `NULL::boolean IN (SELECT foo FROM empty)` == false diff --git a/datafusion/sqllogictest/test_files/predicates.slt b/datafusion/sqllogictest/test_files/predicates.slt index b4482a3af1beb..c447bb61aed23 100644 --- a/datafusion/sqllogictest/test_files/predicates.slt +++ b/datafusion/sqllogictest/test_files/predicates.slt @@ -1022,6 +1022,27 @@ WHERE NULL NOT IN (SELECT * FROM empty); -- all rows should be returned ---- 1 +# Different non-null list expressions can evaluate to the same value. +statement ok +CREATE TABLE in_list_values(x INT NOT NULL, a INT NOT NULL, b INT NOT NULL) +AS VALUES (1, 1, 1), (2, 1, 2), (3, 3, 4); + +query I rowsort +SELECT x FROM in_list_values +WHERE x IN (a, 2, 10, 11) AND x IN (b, 5, 12, 13); +---- +1 +2 + +query I +SELECT x FROM in_list_values +WHERE x IN (a, 2, 10, 11) AND x NOT IN (b, 5, 12, 13); +---- +3 + +statement ok +DROP TABLE in_list_values; + # Config reset statement ok RESET datafusion.catalog.information_schema;