Skip to content

fix: avoid unsafe set rewrites for compound IN lists #24952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
bvolpato wants to merge 1 commit into
base: main
Choose a base branch
from
+234 −2
Draft

fix: avoid unsafe set rewrites for compound IN lists #24952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
215 changes: 213 additions & 2 deletions datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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) {
Expand All @@ -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) {
Expand All @@ -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) {
Expand Down Expand Up @@ -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);
Expand All @@ -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 {
Expand Down Expand Up @@ -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};
Expand Down Expand Up @@ -4672,6 +4703,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(),
);
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
Expand Down
21 changes: 21 additions & 0 deletions datafusion/sqllogictest/test_files/predicates.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down