From fe8f825fdacf045e08ad2573dc456ef47c2b9e4a Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Thu, 17 Sep 2026 02:49:41 +0200 Subject: [PATCH 1/2] feat: support SQL IS [NOT] DISTINCT FROM as null-safe comparison --- rust/lance-datafusion/src/planner.rs | 136 ++++++++++++++++ rust/lance-datafusion/src/signed_zero.rs | 9 +- rust/lance/tests/query/primitives.rs | 195 +++++++++++++++++++++++ 3 files changed, 335 insertions(+), 5 deletions(-) diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 5ee2bbe9a6a..4191ea77b87 100644 --- a/rust/lance-datafusion/src/planner.rs +++ b/rust/lance-datafusion/src/planner.rs @@ -800,6 +800,16 @@ impl Planner { SQLExpr::IsNotTrue(expr) => Ok(Expr::IsNotTrue(Box::new(self.parse_sql_expr(expr)?))), SQLExpr::IsNull(expr) => Ok(Expr::IsNull(Box::new(self.parse_sql_expr(expr)?))), SQLExpr::IsNotNull(expr) => Ok(Expr::IsNotNull(Box::new(self.parse_sql_expr(expr)?))), + SQLExpr::IsDistinctFrom(left, right) => Ok(Expr::BinaryExpr(BinaryExpr::new( + Box::new(self.parse_sql_expr(left)?), + Operator::IsDistinctFrom, + Box::new(self.parse_sql_expr(right)?), + ))), + SQLExpr::IsNotDistinctFrom(left, right) => Ok(Expr::BinaryExpr(BinaryExpr::new( + Box::new(self.parse_sql_expr(left)?), + Operator::IsNotDistinctFrom, + Box::new(self.parse_sql_expr(right)?), + ))), SQLExpr::InList { expr, list, @@ -1688,6 +1698,132 @@ mod tests { ); } + #[test] + fn test_sql_is_distinct_from() { + use arrow_array::Float64Array; + + // Null-safe equality is never NULL: two NULLs are not distinct, and a + // NULL against a value is distinct. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let planner = Planner::new(schema.clone()); + + let expr = planner.parse_filter("a IS DISTINCT FROM b").unwrap(); + assert_eq!( + expr, + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("a")), + Operator::IsDistinctFrom, + Box::new(col("b")), + )) + ); + let expr = planner.parse_filter("a IS NOT DISTINCT FROM b").unwrap(); + assert_eq!( + expr, + Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("a")), + Operator::IsNotDistinctFrom, + Box::new(col("b")), + )) + ); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(1), Some(2), None, None])) as ArrayRef, + ], + ) + .unwrap(); + for (filter, expected) in [ + ( + "a IS DISTINCT FROM b", + BooleanArray::from(vec![false, true, true, false]), + ), + ( + "a IS NOT DISTINCT FROM b", + BooleanArray::from(vec![true, false, false, true]), + ), + ( + "a IS DISTINCT FROM 1", + BooleanArray::from(vec![false, true, true, true]), + ), + ( + "a IS NOT DISTINCT FROM 1", + BooleanArray::from(vec![true, false, false, false]), + ), + ] { + let logical = planner.parse_filter(filter).unwrap(); + let logical = planner.optimize_expr(logical).unwrap(); + let physical = planner.create_physical_expr(&logical).unwrap(); + let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap(); + assert_eq!(result.as_ref(), &expected, "unexpected result for {filter}"); + } + + // Strings use the same null-safe rule. + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + let planner = Planner::new(schema.clone()); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(StringArray::from(vec![Some("a"), None, Some("b")])) as ArrayRef], + ) + .unwrap(); + let logical = planner + .optimize_expr(planner.parse_filter("s IS DISTINCT FROM 'a'").unwrap()) + .unwrap(); + let physical = planner.create_physical_expr(&logical).unwrap(); + let result = physical.evaluate(&batch).unwrap().into_array(3).unwrap(); + assert_eq!( + result.as_ref(), + &BooleanArray::from(vec![false, true, true]) + ); + + // A zero literal is already corrected by `optimize_expr`: both encodings + // match, and a NULL stays decided rather than becoming NULL. Without the + // rewrite Arrow's total order would report `-0.0` as distinct from `0.0`. + // NaN and column-against-column signed zeros keep Arrow's total-order + // answer; this spelling does not fix those. + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Float64, + true, + )])); + let planner = Planner::new(schema.clone()); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Float64Array::from(vec![ + Some(0.0), + Some(-0.0), + Some(1.0), + None, + ])) as ArrayRef], + ) + .unwrap(); + let logical = planner + .parse_filter("value IS NOT DISTINCT FROM 0.0") + .unwrap(); + let optimized = planner.optimize_expr(logical).unwrap(); + let covered = col("value").in_list(vec![lit(-0.0_f64), lit(0.0_f64)], false); + assert_eq!(optimized, covered.is_true()); + let physical = planner.create_physical_expr(&optimized).unwrap(); + let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap(); + assert_eq!( + result.as_ref(), + &BooleanArray::from(vec![true, true, false, false]) + ); + let logical = planner + .optimize_expr(planner.parse_filter("value IS DISTINCT FROM 0.0").unwrap()) + .unwrap(); + let physical = planner.create_physical_expr(&logical).unwrap(); + let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap(); + assert_eq!( + result.as_ref(), + &BooleanArray::from(vec![false, false, true, true]) + ); + } + #[test] fn test_sql_invert() { let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Boolean, true)])); diff --git a/rust/lance-datafusion/src/signed_zero.rs b/rust/lance-datafusion/src/signed_zero.rs index 1b642bc2a28..537cae38567 100644 --- a/rust/lance-datafusion/src/signed_zero.rs +++ b/rust/lance-datafusion/src/signed_zero.rs @@ -26,8 +26,8 @@ use lance_core::Result; /// | `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)` | -/// | `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)` | +/// | `x IS NOT DISTINCT FROM 0` | `(x IN (-0.0, 0.0)) IS TRUE` | +/// | `x IS DISTINCT FROM 0` | `(x IN (-0.0, 0.0)) IS NOT TRUE` | /// /// 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 @@ -631,9 +631,8 @@ mod tests { // bare column, which is also what DataFusion requires before it shortens a // list. This case fails if a release ever relaxes that. #[case::non_column_probe("abs(value) = 0.0")] - // `IS [NOT] DISTINCT FROM` is missing because `Planner::parse_filter` rejects - // it as unsupported SQL; that arm is reachable only from a programmatically - // built expression, and `rewriting_twice_changes_nothing` covers it there. + #[case::is_not_distinct_from("value IS NOT DISTINCT FROM 0.0")] + #[case::is_distinct_from("value IS DISTINCT FROM 0.0")] fn optimizing_twice_changes_nothing(#[case] filter: &str) { let schema = std::sync::Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( diff --git a/rust/lance/tests/query/primitives.rs b/rust/lance/tests/query/primitives.rs index 348d86ec8ce..b8890a3fd3d 100644 --- a/rust/lance/tests/query/primitives.rs +++ b/rust/lance/tests/query/primitives.rs @@ -422,6 +422,201 @@ async fn test_float_zero_predicate_uses_scalar_index() { assert_filter_ids(&ds, "value >= 0.0", &[0, 1, 2, 4, 5, 6]).await; } +/// `IS [NOT] DISTINCT FROM` is the null-safe equality spelling: it is never +/// NULL, two NULLs are not distinct, and a NULL against a value is distinct. +/// +/// A zero literal is already corrected by `optimize_expr`, so both `-0.0` and +/// `+0.0` match `IS NOT DISTINCT FROM 0.0`. That correction intentionally +/// differs from DataFusion 54's native distinct, which keeps the two encodings +/// distinct under Arrow's total order. Column-against-column signed zeros and +/// NaN keep that total-order answer: `-0.0` is distinct from `+0.0`, while two +/// NaNs are not distinct. This spelling does not fix those. +#[tokio::test] +async fn test_query_is_distinct_from() { + let int_batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from((0..4).collect::>())) as ArrayRef, + ), + ( + "value", + Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None])) as ArrayRef, + ), + ( + "other", + Arc::new(Int32Array::from(vec![Some(1), Some(2), None, None])) as ArrayRef, + ), + ]) + .unwrap(); + DatasetTestCases::from_data(int_batch) + .with_index_types( + "value", + [ + None, + Some(IndexType::Bitmap), + Some(IndexType::BTree), + Some(IndexType::BloomFilter), + Some(IndexType::ZoneMap), + ], + ) + .run(|ds: Dataset, original: RecordBatch| async move { + // Integers have no total-order quirk, so DataFusion is the reference. + test_filter(&original, &ds, "value IS DISTINCT FROM other").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM other").await; + test_filter(&original, &ds, "value IS DISTINCT FROM 1").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM 1").await; + test_filter(&original, &ds, "value IS DISTINCT FROM NULL").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM NULL").await; + }) + .await; + + let string_batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from((0..4).collect::>())) as ArrayRef, + ), + ( + "value", + Arc::new(StringArray::from(vec![Some("a"), None, Some("b"), None])) as ArrayRef, + ), + ( + "other", + Arc::new(StringArray::from(vec![Some("a"), Some("b"), None, None])) as ArrayRef, + ), + ]) + .unwrap(); + DatasetTestCases::from_data(string_batch) + .with_index_types( + "value", + [ + None, + Some(IndexType::Bitmap), + Some(IndexType::BTree), + Some(IndexType::BloomFilter), + Some(IndexType::ZoneMap), + ], + ) + .run(|ds: Dataset, original: RecordBatch| async move { + test_filter(&original, &ds, "value IS DISTINCT FROM other").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM other").await; + test_filter(&original, &ds, "value IS DISTINCT FROM 'a'").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM 'a'").await; + }) + .await; + + // Ids are 0: +0.0, 1: -0.0, 2: 1.0, 3: NULL, 4: NaN. `other` matches except + // row 1, which holds +0.0 where `value` holds -0.0. + let float_batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from((0..5).collect::>())) as ArrayRef, + ), + ( + "value", + Arc::new(Float64Array::from(vec![ + Some(0.0), + Some(-0.0), + Some(1.0), + None, + Some(f64::NAN), + ])) as ArrayRef, + ), + ( + "other", + Arc::new(Float64Array::from(vec![ + Some(0.0), + Some(0.0), + Some(1.0), + None, + Some(f64::NAN), + ])) as ArrayRef, + ), + ]) + .unwrap(); + DatasetTestCases::from_data(float_batch) + .with_index_types( + "value", + [ + None, + Some(IndexType::Bitmap), + Some(IndexType::BTree), + Some(IndexType::BloomFilter), + Some(IndexType::ZoneMap), + ], + ) + .run(|ds: Dataset, _original: RecordBatch| async move { + // Both zero encodings match a zero literal; NULL never matches one. + // DataFusion's native distinct is not the reference here because it + // keeps `-0.0` distinct from `0.0`. + assert_filter_ids(&ds, "value IS NOT DISTINCT FROM 0.0", &[0, 1]).await; + assert_filter_ids(&ds, "value IS DISTINCT FROM 0.0", &[2, 3, 4]).await; + assert_filter_ids(&ds, "value IS NOT DISTINCT FROM 1.0", &[2]).await; + assert_filter_ids(&ds, "value IS DISTINCT FROM 1.0", &[0, 1, 3, 4]).await; + // Column against column keeps Arrow's total order: the sign-swapped + // zero pair is distinct, while two NaNs and two NULLs are not. + assert_filter_ids(&ds, "value IS DISTINCT FROM other", &[1]).await; + assert_filter_ids(&ds, "value IS NOT DISTINCT FROM other", &[0, 2, 3, 4]).await; + assert_filter_ids(&ds, "value IS DISTINCT FROM NULL", &[0, 1, 2, 4]).await; + assert_filter_ids(&ds, "value IS NOT DISTINCT FROM NULL", &[3]).await; + }) + .await; +} + +/// A null-safe comparison stays correct when a scalar index exists, even though +/// the index does not accelerate it. Without this, a future index arm for +/// `IS DISTINCT FROM` (or for the `(IN (..)) IS TRUE` shape the zero rewrite +/// emits) could silently return the wrong rows instead of falling back to a +/// filter. +#[tokio::test] +async fn test_is_distinct_from_does_not_use_scalar_index() { + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from_iter_values(0..4)) as ArrayRef, + ), + ( + "value", + Arc::new(Float64Array::from(vec![0.0, -0.0, 1.0, -1.0])) as ArrayRef, + ), + ]) + .unwrap(); + let schema = batch.schema(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut ds = Dataset::write(reader, "memory://distinct_index_fallback", None) + .await + .unwrap(); + ds.create_index( + &["value"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + for predicate in [ + "value IS DISTINCT FROM 0.0", + "value IS NOT DISTINCT FROM 0.0", + "value IS DISTINCT FROM 1.0", + ] { + let plan = ds + .scan() + .filter(predicate) + .unwrap() + .explain_plan(false) + .await + .unwrap(); + assert!( + !plan.contains("ScalarIndexQuery"), + "`{predicate}` should fall back to filtering, got plan:\n{plan}" + ); + } + + assert_filter_ids(&ds, "value IS NOT DISTINCT FROM 0.0", &[0, 1]).await; + assert_filter_ids(&ds, "value IS DISTINCT FROM 0.0", &[2, 3]).await; +} + #[tokio::test] #[rstest::rstest] #[case::date32(DataType::Date32)] From e380f044ba22c3743e5b529c940a843a861a8a36 Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Thu, 17 Sep 2026 16:07:29 +0200 Subject: [PATCH 2/2] test: fold IS [NOT] DISTINCT FROM coverage into existing tests Drop the test asserting the scalar index is not used and the asserts pinning known float total-order behavior (#9332, #9315); cover integers, strings and both zero spellings in the existing query tests instead. --- rust/lance-datafusion/src/planner.rs | 132 ++--------------- rust/lance/tests/query/primitives.rs | 208 ++------------------------- 2 files changed, 24 insertions(+), 316 deletions(-) diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 4191ea77b87..85d9a8f7970 100644 --- a/rust/lance-datafusion/src/planner.rs +++ b/rust/lance-datafusion/src/planner.rs @@ -1698,130 +1698,20 @@ mod tests { ); } - #[test] - fn test_sql_is_distinct_from() { - use arrow_array::Float64Array; - - // Null-safe equality is never NULL: two NULLs are not distinct, and a - // NULL against a value is distinct. - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - ])); - let planner = Planner::new(schema.clone()); - - let expr = planner.parse_filter("a IS DISTINCT FROM b").unwrap(); - assert_eq!( - expr, - Expr::BinaryExpr(BinaryExpr::new( - Box::new(col("a")), - Operator::IsDistinctFrom, - Box::new(col("b")), - )) - ); - let expr = planner.parse_filter("a IS NOT DISTINCT FROM b").unwrap(); - assert_eq!( - expr, - Expr::BinaryExpr(BinaryExpr::new( - Box::new(col("a")), - Operator::IsNotDistinctFrom, - Box::new(col("b")), - )) - ); - - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None])) as ArrayRef, - Arc::new(Int32Array::from(vec![Some(1), Some(2), None, None])) as ArrayRef, - ], - ) - .unwrap(); - for (filter, expected) in [ - ( - "a IS DISTINCT FROM b", - BooleanArray::from(vec![false, true, true, false]), - ), - ( - "a IS NOT DISTINCT FROM b", - BooleanArray::from(vec![true, false, false, true]), - ), - ( - "a IS DISTINCT FROM 1", - BooleanArray::from(vec![false, true, true, true]), - ), - ( - "a IS NOT DISTINCT FROM 1", - BooleanArray::from(vec![true, false, false, false]), - ), - ] { - let logical = planner.parse_filter(filter).unwrap(); - let logical = planner.optimize_expr(logical).unwrap(); - let physical = planner.create_physical_expr(&logical).unwrap(); - let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap(); - assert_eq!(result.as_ref(), &expected, "unexpected result for {filter}"); - } - - // Strings use the same null-safe rule. - let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); - let planner = Planner::new(schema.clone()); - let batch = RecordBatch::try_new( - schema, - vec![Arc::new(StringArray::from(vec![Some("a"), None, Some("b")])) as ArrayRef], - ) - .unwrap(); - let logical = planner - .optimize_expr(planner.parse_filter("s IS DISTINCT FROM 'a'").unwrap()) - .unwrap(); - let physical = planner.create_physical_expr(&logical).unwrap(); - let result = physical.evaluate(&batch).unwrap().into_array(3).unwrap(); - assert_eq!( - result.as_ref(), - &BooleanArray::from(vec![false, true, true]) - ); - - // A zero literal is already corrected by `optimize_expr`: both encodings - // match, and a NULL stays decided rather than becoming NULL. Without the - // rewrite Arrow's total order would report `-0.0` as distinct from `0.0`. - // NaN and column-against-column signed zeros keep Arrow's total-order - // answer; this spelling does not fix those. - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Float64, - true, - )])); - let planner = Planner::new(schema.clone()); - let batch = RecordBatch::try_new( - schema, - vec![Arc::new(Float64Array::from(vec![ - Some(0.0), - Some(-0.0), - Some(1.0), - None, - ])) as ArrayRef], + #[rstest] + #[case::distinct("a IS DISTINCT FROM b", [false, true, true, false])] + #[case::not_distinct("a IS NOT DISTINCT FROM b", [true, false, false, true])] + fn test_sql_is_distinct_from(#[case] filter: &str, #[case] expected: [bool; 4]) { + let batch = arrow_array::record_batch!( + ("a", Int32, [Some(1), None, Some(2), None]), + ("b", Int32, [Some(1), Some(2), None, None]) ) .unwrap(); - let logical = planner - .parse_filter("value IS NOT DISTINCT FROM 0.0") - .unwrap(); - let optimized = planner.optimize_expr(logical).unwrap(); - let covered = col("value").in_list(vec![lit(-0.0_f64), lit(0.0_f64)], false); - assert_eq!(optimized, covered.is_true()); - let physical = planner.create_physical_expr(&optimized).unwrap(); - let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap(); - assert_eq!( - result.as_ref(), - &BooleanArray::from(vec![true, true, false, false]) - ); - let logical = planner - .optimize_expr(planner.parse_filter("value IS DISTINCT FROM 0.0").unwrap()) - .unwrap(); - let physical = planner.create_physical_expr(&logical).unwrap(); + let planner = Planner::new(batch.schema()); + let expr = planner.optimize_expr(planner.parse_filter(filter).unwrap()); + let physical = planner.create_physical_expr(&expr.unwrap()).unwrap(); let result = physical.evaluate(&batch).unwrap().into_array(4).unwrap(); - assert_eq!( - result.as_ref(), - &BooleanArray::from(vec![false, false, true, true]) - ); + assert_eq!(result.as_ref(), &BooleanArray::from(expected.to_vec())); } #[test] diff --git a/rust/lance/tests/query/primitives.rs b/rust/lance/tests/query/primitives.rs index b8890a3fd3d..4e156ca20e0 100644 --- a/rust/lance/tests/query/primitives.rs +++ b/rust/lance/tests/query/primitives.rs @@ -83,6 +83,8 @@ async fn test_query_integer(#[case] data_type: DataType) { test_take(&original, &ds).await; test_filter(&original, &ds, "value > 20").await; test_filter(&original, &ds, "NOT (value > 20)").await; + test_filter(&original, &ds, "value IS DISTINCT FROM 20").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM NULL").await; test_filter(&original, &ds, "value is null").await; test_filter(&original, &ds, "value is not null").await; test_filter(&original, &ds, "(value != 0) OR (value < 20)").await; @@ -299,6 +301,15 @@ async fn test_query_float_special_values(#[case] data_type: DataType) { assert_filter_ids(&ds, &format!("value < {zero}"), &[3, 6, 7]).await; assert_filter_ids(&ds, &format!("value <= {zero}"), &[0, 1, 3, 6, 7]).await; assert_filter_ids(&ds, &format!("value = {zero}"), &[0, 1]).await; + // Null-safe: the NULL row 9 is distinct rather than unknown. + assert_filter_ids(&ds, &format!("value IS NOT DISTINCT FROM {zero}"), &[0, 1]) + .await; + assert_filter_ids( + &ds, + &format!("value IS DISTINCT FROM {zero}"), + &[2, 3, 4, 5, 6, 7, 8, 9], + ) + .await; assert_filter_ids(&ds, &format!("value != {zero}"), &[2, 3, 4, 5, 6, 7, 8]).await; // NaN is row 4. Arrow sorts it above every other value, so it // survives `>` and `>=`, which IEEE would reject. That gap is @@ -422,201 +433,6 @@ async fn test_float_zero_predicate_uses_scalar_index() { assert_filter_ids(&ds, "value >= 0.0", &[0, 1, 2, 4, 5, 6]).await; } -/// `IS [NOT] DISTINCT FROM` is the null-safe equality spelling: it is never -/// NULL, two NULLs are not distinct, and a NULL against a value is distinct. -/// -/// A zero literal is already corrected by `optimize_expr`, so both `-0.0` and -/// `+0.0` match `IS NOT DISTINCT FROM 0.0`. That correction intentionally -/// differs from DataFusion 54's native distinct, which keeps the two encodings -/// distinct under Arrow's total order. Column-against-column signed zeros and -/// NaN keep that total-order answer: `-0.0` is distinct from `+0.0`, while two -/// NaNs are not distinct. This spelling does not fix those. -#[tokio::test] -async fn test_query_is_distinct_from() { - let int_batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from((0..4).collect::>())) as ArrayRef, - ), - ( - "value", - Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None])) as ArrayRef, - ), - ( - "other", - Arc::new(Int32Array::from(vec![Some(1), Some(2), None, None])) as ArrayRef, - ), - ]) - .unwrap(); - DatasetTestCases::from_data(int_batch) - .with_index_types( - "value", - [ - None, - Some(IndexType::Bitmap), - Some(IndexType::BTree), - Some(IndexType::BloomFilter), - Some(IndexType::ZoneMap), - ], - ) - .run(|ds: Dataset, original: RecordBatch| async move { - // Integers have no total-order quirk, so DataFusion is the reference. - test_filter(&original, &ds, "value IS DISTINCT FROM other").await; - test_filter(&original, &ds, "value IS NOT DISTINCT FROM other").await; - test_filter(&original, &ds, "value IS DISTINCT FROM 1").await; - test_filter(&original, &ds, "value IS NOT DISTINCT FROM 1").await; - test_filter(&original, &ds, "value IS DISTINCT FROM NULL").await; - test_filter(&original, &ds, "value IS NOT DISTINCT FROM NULL").await; - }) - .await; - - let string_batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from((0..4).collect::>())) as ArrayRef, - ), - ( - "value", - Arc::new(StringArray::from(vec![Some("a"), None, Some("b"), None])) as ArrayRef, - ), - ( - "other", - Arc::new(StringArray::from(vec![Some("a"), Some("b"), None, None])) as ArrayRef, - ), - ]) - .unwrap(); - DatasetTestCases::from_data(string_batch) - .with_index_types( - "value", - [ - None, - Some(IndexType::Bitmap), - Some(IndexType::BTree), - Some(IndexType::BloomFilter), - Some(IndexType::ZoneMap), - ], - ) - .run(|ds: Dataset, original: RecordBatch| async move { - test_filter(&original, &ds, "value IS DISTINCT FROM other").await; - test_filter(&original, &ds, "value IS NOT DISTINCT FROM other").await; - test_filter(&original, &ds, "value IS DISTINCT FROM 'a'").await; - test_filter(&original, &ds, "value IS NOT DISTINCT FROM 'a'").await; - }) - .await; - - // Ids are 0: +0.0, 1: -0.0, 2: 1.0, 3: NULL, 4: NaN. `other` matches except - // row 1, which holds +0.0 where `value` holds -0.0. - let float_batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from((0..5).collect::>())) as ArrayRef, - ), - ( - "value", - Arc::new(Float64Array::from(vec![ - Some(0.0), - Some(-0.0), - Some(1.0), - None, - Some(f64::NAN), - ])) as ArrayRef, - ), - ( - "other", - Arc::new(Float64Array::from(vec![ - Some(0.0), - Some(0.0), - Some(1.0), - None, - Some(f64::NAN), - ])) as ArrayRef, - ), - ]) - .unwrap(); - DatasetTestCases::from_data(float_batch) - .with_index_types( - "value", - [ - None, - Some(IndexType::Bitmap), - Some(IndexType::BTree), - Some(IndexType::BloomFilter), - Some(IndexType::ZoneMap), - ], - ) - .run(|ds: Dataset, _original: RecordBatch| async move { - // Both zero encodings match a zero literal; NULL never matches one. - // DataFusion's native distinct is not the reference here because it - // keeps `-0.0` distinct from `0.0`. - assert_filter_ids(&ds, "value IS NOT DISTINCT FROM 0.0", &[0, 1]).await; - assert_filter_ids(&ds, "value IS DISTINCT FROM 0.0", &[2, 3, 4]).await; - assert_filter_ids(&ds, "value IS NOT DISTINCT FROM 1.0", &[2]).await; - assert_filter_ids(&ds, "value IS DISTINCT FROM 1.0", &[0, 1, 3, 4]).await; - // Column against column keeps Arrow's total order: the sign-swapped - // zero pair is distinct, while two NaNs and two NULLs are not. - assert_filter_ids(&ds, "value IS DISTINCT FROM other", &[1]).await; - assert_filter_ids(&ds, "value IS NOT DISTINCT FROM other", &[0, 2, 3, 4]).await; - assert_filter_ids(&ds, "value IS DISTINCT FROM NULL", &[0, 1, 2, 4]).await; - assert_filter_ids(&ds, "value IS NOT DISTINCT FROM NULL", &[3]).await; - }) - .await; -} - -/// A null-safe comparison stays correct when a scalar index exists, even though -/// the index does not accelerate it. Without this, a future index arm for -/// `IS DISTINCT FROM` (or for the `(IN (..)) IS TRUE` shape the zero rewrite -/// emits) could silently return the wrong rows instead of falling back to a -/// filter. -#[tokio::test] -async fn test_is_distinct_from_does_not_use_scalar_index() { - let batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..4)) as ArrayRef, - ), - ( - "value", - Arc::new(Float64Array::from(vec![0.0, -0.0, 1.0, -1.0])) as ArrayRef, - ), - ]) - .unwrap(); - let schema = batch.schema(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); - let mut ds = Dataset::write(reader, "memory://distinct_index_fallback", None) - .await - .unwrap(); - ds.create_index( - &["value"], - IndexType::Scalar, - None, - &ScalarIndexParams::default(), - false, - ) - .await - .unwrap(); - - for predicate in [ - "value IS DISTINCT FROM 0.0", - "value IS NOT DISTINCT FROM 0.0", - "value IS DISTINCT FROM 1.0", - ] { - let plan = ds - .scan() - .filter(predicate) - .unwrap() - .explain_plan(false) - .await - .unwrap(); - assert!( - !plan.contains("ScalarIndexQuery"), - "`{predicate}` should fall back to filtering, got plan:\n{plan}" - ); - } - - assert_filter_ids(&ds, "value IS NOT DISTINCT FROM 0.0", &[0, 1]).await; - assert_filter_ids(&ds, "value IS DISTINCT FROM 0.0", &[2, 3]).await; -} - #[tokio::test] #[rstest::rstest] #[case::date32(DataType::Date32)] @@ -738,6 +554,8 @@ async fn test_query_string(#[case] data_type: DataType) { test_take(&original, &ds).await; test_filter(&original, &ds, "value = 'hello'").await; test_filter(&original, &ds, "value != 'hello'").await; + test_filter(&original, &ds, "value IS DISTINCT FROM 'hello'").await; + test_filter(&original, &ds, "value IS NOT DISTINCT FROM 'hello'").await; test_filter(&original, &ds, "value = ''").await; test_filter(&original, &ds, "value > 'hello'").await; test_filter(&original, &ds, "value is null").await;