diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 5ee2bbe9a6a..85d9a8f7970 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,22 @@ mod tests { ); } + #[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 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(expected.to_vec())); + } + #[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..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 @@ -543,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;