Skip to content

fix(datafusion): treat signed zeros as equal in array_has #9326

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

Closed
jonasdedden wants to merge 1 commit into from
Closed

fix(datafusion): treat signed zeros as equal in array_has #9326

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
284 changes: 278 additions & 6 deletions rust/lance-datafusion/src/signed_zero.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! Rewrites of comparisons against a floating point zero literal.

Check warning on line 4 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs

use datafusion::error::Result as DFResult;
use datafusion::logical_expr::{BinaryExpr, Operator, expr::Between, expr::InList};
use datafusion::logical_expr::{BinaryExpr, Operator, expr::Between, expr::InList, expr::ScalarFunction};
use datafusion::prelude::Expr;
use datafusion::scalar::ScalarValue::{self, Float16, Float32, Float64};
use datafusion_common::tree_node::{Transformed, TreeNode};
Expand All @@ -26,12 +26,16 @@
/// | `x != 0` | `x NOT IN (-0.0, 0.0)` |
/// | `x IN (0, ..)` | the missing encoding is added |
/// | `0 IN (a, b)` | `a IN (-0.0, 0.0) OR b IN (-0.0, 0.0)` |
/// | `array_has(l, 0)` | `array_has(l, -0.0) OR array_has(l, 0.0)` |
/// | `x IS NOT DISTINCT FROM 0` | `x IS NOT NULL AND x IN (-0.0, 0.0)` |
/// | `x IS DISTINCT FROM 0` | `x IS NULL OR x NOT IN (-0.0, 0.0)` |
///
/// Equality has to name both encodings because a scalar index keys on the bit
/// pattern: the btree and bitmap indices order candidates by `total_cmp`, and the
/// bloom filter hashes the value.
/// bloom filter hashes the value. `array_has` needs the same disjunction because
/// its kernel compares list elements with Arrow's total-order `eq`, and a scalar
/// index never accelerates it into a different path: the rewritten filter still
/// plans as a plain filter.
///
/// Runs as the last step of [`crate::planner::Planner::optimize_expr`], after
/// coercion has given the literal the column's type and the simplifier has
Expand Down Expand Up @@ -68,8 +72,8 @@
)
}

/// Fold each zero-sensitive comparison's own operands and rewrite it, bottom-up,
/// before anything above it has a chance to fold.
/// Fold each zero-sensitive comparison's (or `array_has` probe's) own operands and
/// rewrite it, bottom-up, before anything above it has a chance to fold.
///
/// [`rewrite_signed_zero_comparisons`] alone cannot reach a comparison whose zero
/// does not exist yet. `ExprSimplifier::simplify` folds an operand and everything
Expand Down Expand Up @@ -121,6 +125,17 @@
.collect::<DFResult<Vec<_>>>()?,
negated: in_list.negated,
}),
Expr::ScalarFunction(func) if is_array_has(func.name()) => {
// Fold the needle (and haystack) before rewriting, so a
// computed zero such as `-1.0 * 0.0` presents as a zero
// literal. Folding the whole `array_has` here would decide
// it by Arrow's total order before the rewrite runs.
let mut args = Vec::with_capacity(func.args.len());
for arg in func.args {

Check warning on line 134 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
args.push(fold(arg)?);
}
Expr::ScalarFunction(ScalarFunction { func: func.func, args })
}
other => return Ok(Transformed::no(other)),
};
Ok(match rewrite_node(&folded) {
Expand Down Expand Up @@ -210,13 +225,17 @@
// over a bare column back into an OR chain of equalities, so a second
// `optimize_expr` splits this rewrite's own output and re-runs it on each
// half. Both halves then produce the same list, and dropping the repeat is
// what makes the rewrite survive that round trip.
// what makes the rewrite survive that round trip. The `array_has`
// rewrite below has the same shape: each half of its own disjunction
// expands again, so identical zero probes are dropped here too.
Expr::BinaryExpr(BinaryExpr { op, .. }) if matches!(op, Operator::Or | Operator::And) => {
let mut kept: Vec<&Expr> = Vec::new();
flatten_chain(expr, *op, &mut kept);
let mut deduped: Vec<&Expr> = Vec::with_capacity(kept.len());
for term in kept.iter() {
if is_zero_pair_over_column(term) && deduped.contains(term) {
if (is_zero_pair_over_column(term) || is_array_has_zero_probe(term))
&& deduped.contains(term)
{
continue;
}
deduped.push(term);
Expand Down Expand Up @@ -353,10 +372,74 @@
negated: *negated,
}))
}
Expr::ScalarFunction(func) if is_array_has(func.name()) => rewrite_array_has(func),
_ => None,
}
}

/// Whether `name` is the list-membership function, under any of DataFusion's
/// aliases. Matching aliases here keeps `list_has`, `array_contains`, and
/// `list_contains` on the same IEEE semantics as `array_has`: they resolve to
/// the same kernel, so leaving one spelling sign-sensitive would be a dialect
/// accident rather than a rule.
fn is_array_has(name: &str) -> bool {
matches!(
name.to_ascii_lowercase().as_str(),
"array_has" | "list_has" | "array_contains" | "list_contains"
)
}

/// True for a single `array_has` probe against a floating point zero literal,
/// either encoding. Only these terms are deduplicated, so an expression the
/// caller wrote twice for another needle is left alone.
fn is_array_has_zero_probe(expr: &Expr) -> bool {
let Expr::ScalarFunction(func) = expr else {
return false;
};
if !is_array_has(func.name()) || func.args.len() != 2 {
return false;
}
matches!(&func.args[1], Expr::Literal(value, _) if zero_encodings(value).is_some())
}

/// Rewrite `array_has(haystack, zero)` into the disjunction Arrow's total-order
/// `eq` kernel answers the way IEEE 754 defines it.
///
/// Returns `None` for anything but a floating point zero needle, including
/// NULL, NaN, integer zero, and non-literal needles. An integer needle against
/// a float list is still rewritten: coercion has already given it the column's
/// floating type by the time this runs, so it presents as a floating zero. An
/// integer list keeps integer semantics because its needle never presents as a
/// floating zero. A column needle stays sign-sensitive; that is the
/// column-to-column case (issue #9316, first part), not this literal rewrite.
fn rewrite_array_has(func: &ScalarFunction) -> Option<Expr> {
if func.args.len() != 2 {
return None;
}
let haystack = &func.args[0];
let Expr::Literal(value, metadata) = &func.args[1] else {
return None;
};
let (negative, positive) = zero_encodings(value)?;
// Negative first, so both encodings map to the same canonical disjunction
// and a second pass deduplicates rather than growing.

Check warning on line 425 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
let matches_negative = Expr::ScalarFunction(ScalarFunction {
func: func.func.clone(),
args: vec![
haystack.clone(),
Expr::Literal(negative, metadata.clone()),
],
});
let matches_positive = Expr::ScalarFunction(ScalarFunction {
func: func.func.clone(),
args: vec![

Check warning on line 435 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
haystack.clone(),

@lance-gatekeeper lance-gatekeeper Bot Sep 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cloning the haystack into both OR arms changes the semantics of supported volatile expressions: DataFusion permits a Volatility::Volatile function to return a different value on each evaluation. I ran a volatile list UDF that returns [+0.0] on its first call and [-0.0] on its second. Optimizing array_has(volatile_zeros(), 0.0) produced the two calls shown below and returned false for all 10,000 rows even though every UDF result contained a zero. Preserve a single haystack evaluation—for example, lower this to one array_has_any(haystack, [-0.0, 0.0]) probe (the same experiment then made one call and returned all 10,000 rows), or normalize equality in the kernel.

Reproducer run on this head 
let schema = Arc::new(Schema::empty());
let planner = Planner::new(schema.clone());
let calls = Arc::new(AtomicUsize::new(0));
let calls_in_udf = Arc::clone(&calls);
let list_type = DataType::List(Arc::new(Field::new("item", DataType::Float64, true)));
let volatile_zeros = create_udf(
    "volatile_zeros",
    vec![],
    list_type,
    Volatility::Volatile,
    Arc::new(move |_| {
        let value = if calls_in_udf.fetch_add(1, Ordering::SeqCst) % 2 == 0 {
            0.0
        } else {
            -0.0
        };
        let list = ListArray::from_iter_primitive::<Float64Type, _, _>(
            [Some(vec![Some(value)])],
        );
        Ok(ColumnarValue::Scalar(ScalarValue::List(Arc::new(list))))
    }),
);
let optimized = planner.optimize_expr(array_has(volatile_zeros.call(vec![]), lit(0.0)))?;
let batch = RecordBatch::try_new_with_options(
    schema,
    vec![],
    &RecordBatchOptions::new().with_row_count(Some(10_000)),
)?;
let values = planner
    .create_physical_expr(&optimized)?
    .evaluate(&batch)?
    .into_array(10_000)?;
let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
assert_eq!(calls.load(Ordering::SeqCst), 1); // observed: 2
assert_eq!(values.true_count(), 10_000); // observed: 0

Observed optimized expression:

array_has(volatile_zeros(), Float64(-0)) OR array_has(volatile_zeros(), Float64(0))
calls=2 true=0 false=10000

Expr::Literal(positive, metadata.clone()),
],
});
Some(matches_negative.or(matches_positive))
}

/// Add the missing encoding next to every floating point zero in an `IN` list.
///
/// Returns `None` when the list holds no zero, or already spells out both
Expand Down Expand Up @@ -607,6 +690,195 @@
);
}

fn list_schema(inner: arrow_schema::DataType) -> std::sync::Arc<arrow_schema::Schema> {
std::sync::Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new(
"l",
arrow_schema::DataType::List(std::sync::Arc::new(arrow_schema::Field::new(
"item", inner, true,
))),
true,
)]))
}

fn unalias(mut expr: &Expr) -> &Expr {
while let Expr::Alias(alias) = expr {
expr = &alias.expr;
}
expr
}

fn is_array_has_or_of_zero_pair(expr: &Expr) -> bool {
let Expr::BinaryExpr(BinaryExpr { left, op, right }) = unalias(expr) else {
return false;
};
if *op != Operator::Or {
return false;
}
is_array_has_zero_probe(left) && is_array_has_zero_probe(right) && left != right
}

#[test]
fn array_has_zero_needle_covers_both_encodings() {
let planner = crate::planner::Planner::new(list_schema(arrow_schema::DataType::Float64));
let expected = planner
.parse_filter("array_has(l, -0.0) OR array_has(l, 0.0)")
.unwrap();
// Both encodings map to the same canonical disjunction, negative first.
assert_eq!(
rewrite(planner.parse_filter("array_has(l, 0.0)").unwrap()),
expected
);
assert_eq!(
rewrite(planner.parse_filter("array_has(l, -0.0)").unwrap()),
expected
);
}

#[test]
fn array_has_leaves_non_zero_needles_alone() {
let planner = crate::planner::Planner::new(list_schema(arrow_schema::DataType::Float64));
// At the post-simplify layer an integer `0` is still an integer zero,
// so there is nothing to widen. After coercion against a float list it
// presents as a floating zero and is rewritten; that is covered by the
// optimizer test below.
for filter in ["array_has(l, 1.0)", "array_has(l, 0)", "array_has(l, NULL)"] {
let parsed = planner.parse_filter(filter).unwrap();
assert_eq!(rewrite(parsed.clone()), parsed, "filter: {filter}");
}
// NaN is out of scope, like everywhere else in this module.
let parsed = planner
.parse_filter("array_has(l, CAST('NaN' AS DOUBLE))")
.unwrap();
assert_eq!(rewrite(parsed.clone()), parsed);
}

#[test]
fn array_has_leaves_column_needles_alone() {
// A column needle is the column-to-column case (issue #9316, first
// part), not this literal rewrite.
let schema = std::sync::Arc::new(arrow_schema::Schema::new(vec![
arrow_schema::Field::new(
"l",
arrow_schema::DataType::List(std::sync::Arc::new(arrow_schema::Field::new(
"item",
arrow_schema::DataType::Float64,
true,
))),
true,
),
arrow_schema::Field::new("m", arrow_schema::DataType::Float64, true),
]));
let planner = crate::planner::Planner::new(schema);
let parsed = planner.parse_filter("array_has(l, m)").unwrap();
assert_eq!(rewrite(parsed.clone()), parsed);
}

#[test]
fn array_has_integer_list_keeps_integer_semantics() {
let planner = crate::planner::Planner::new(list_schema(arrow_schema::DataType::Int64));
let optimized = planner
.optimize_expr(planner.parse_filter("array_has(l, 0)").unwrap())
.unwrap();
assert!(
!is_array_has_or_of_zero_pair(&optimized),
"integer needle must stay a single probe: {optimized:?}"
);
}

#[test]
fn array_has_float_lists_cover_both_encodings_after_coercion() {
for inner in [
arrow_schema::DataType::Float64,
arrow_schema::DataType::Float32,
] {

Check warning on line 793 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
let planner = crate::planner::Planner::new(list_schema(inner.clone()));
for filter in ["array_has(l, 0.0)", "array_has(l, -0.0)", "array_has(l, 0)"] {
let optimized = planner.optimize_expr(planner.parse_filter(filter).unwrap()).unwrap();
assert!(
is_array_has_or_of_zero_pair(&optimized),
"filter {filter} on {inner:?} must be a zero-pair OR, got: {optimized:?}"
);
// A second optimization is a fixed point.
assert_eq!(
planner.optimize_expr(optimized.clone()).unwrap(),
optimized,
"filter: {filter} on {inner:?}"
);
}
}
}

#[test]
fn array_has_rewrites_nested_and_computed_haystacks() {
let schema = std::sync::Arc::new(arrow_schema::Schema::new(vec![
arrow_schema::Field::new(
"l",
arrow_schema::DataType::List(std::sync::Arc::new(arrow_schema::Field::new(
"item",
arrow_schema::DataType::Float64,
true,
))),
true,
),
arrow_schema::Field::new(
"s",
arrow_schema::DataType::Struct(arrow_schema::Fields::from(vec![
arrow_schema::Field::new(

Check warning on line 826 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
"x",
arrow_schema::DataType::List(std::sync::Arc::new(
arrow_schema::Field::new(
"item",
arrow_schema::DataType::Float64,
true,
),
)),
true,
),
])),
true,
),
]));
let planner = crate::planner::Planner::new(schema);
for filter in [
"array_has(s.x, 0.0)",

Check warning on line 843 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
"array_has(array_append(l, 1.0), 0.0)",
] {
let optimized = planner.optimize_expr(planner.parse_filter(filter).unwrap()).unwrap();
assert!(
is_array_has_or_of_zero_pair(&optimized),
"filter {filter} must be a zero-pair OR, got: {optimized:?}"
);
}
}

#[test]
fn array_has_rewrites_computed_zero_needles() {
let planner = crate::planner::Planner::new(list_schema(arrow_schema::DataType::Float64));
let canonical = planner
.optimize_expr(planner.parse_filter("array_has(l, 0.0)").unwrap())

Check warning on line 858 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
.unwrap();
for filter in ["array_has(l, -1.0 * 0.0)", "array_has(l, 1.0 - 1.0)"] {
let optimized = planner.optimize_expr(planner.parse_filter(filter).unwrap()).unwrap();
assert_eq!(optimized, canonical, "filter: {filter}");
}
}

#[test]
fn array_has_aliases_rewrite_like_array_has() {
let planner = crate::planner::Planner::new(list_schema(arrow_schema::DataType::Float64));
for filter in [
"list_has(l, 0.0)",
"array_contains(l, 0.0)",

Check warning on line 871 in rust/lance-datafusion/src/signed_zero.rs

View workflow job for this annotation

GitHub Actions / format 

Diff in /home/runner/work/lance/lance/rust/lance-datafusion/src/signed_zero.rs
"list_contains(l, 0.0)",
] {
let optimized = planner.optimize_expr(planner.parse_filter(filter).unwrap()).unwrap();
assert!(
is_array_has_or_of_zero_pair(&optimized),
"filter {filter} must be a zero-pair OR, got: {optimized:?}"
);
}
}

#[test]
fn scalar_value_keeps_the_two_zero_encodings_apart() {
// The `IN` list widening decides "already listed" with this comparison. A
Expand Down
Loading