Route IS DISTINCT FROM through binary expression planners

datafusion/sql/src/expr/mod.rs

@@ -40,8 +40,8 @@ use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::expr::SetQuantifier;
 use datafusion_expr::expr::{InList, WildcardOptions};
 use datafusion_expr::{
-    Between, BinaryExpr, Cast, Expr, ExprSchemable, GetFieldAccess, Like, Literal,
-    Operator, TryCast, lit, when,
+    Between, Cast, Expr, ExprSchemable, GetFieldAccess, Like, Literal, Operator, TryCast,
+    lit, when,
 };
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
@@ -422,35 +422,33 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             ))),
 
             SQLExpr::IsDistinctFrom(left, right) => {
-                Ok(Expr::BinaryExpr(BinaryExpr::new(
-                    Box::new(self.sql_expr_to_logical_expr(
-                        *left,
-                        schema,
-                        planner_context,
-                    )?),
-                    Operator::IsDistinctFrom,
-                    Box::new(self.sql_expr_to_logical_expr(
-                        *right,
-                        schema,
-                        planner_context,
-                    )?),
-                )))
+                let left =
+                    self.sql_expr_to_logical_expr(*left, schema, planner_context)?;
+                let right =
+                    self.sql_expr_to_logical_expr(*right, schema, planner_context)?;
+                let planned = self.build_logical_expr(
+                    BinaryOperator::Spaceship,
+                    left,
+                    right,
+                    schema,
+                )?;
+                Ok(match planned {
+                    Expr::BinaryExpr(mut binary)
+                        if binary.op == Operator::IsNotDistinctFrom =>
+                    {
+                        binary.op = Operator::IsDistinctFrom;
+                        Expr::BinaryExpr(binary)
+                    }
+                    other => Expr::Not(Box::new(other)),
+                })
             }
 
             SQLExpr::IsNotDistinctFrom(left, right) => {
-                Ok(Expr::BinaryExpr(BinaryExpr::new(
-                    Box::new(self.sql_expr_to_logical_expr(
-                        *left,
-                        schema,
-                        planner_context,
-                    )?),
-                    Operator::IsNotDistinctFrom,
-                    Box::new(self.sql_expr_to_logical_expr(
-                        *right,
-                        schema,
-                        planner_context,
-                    )?),
-                )))
+                let left =
+                    self.sql_expr_to_logical_expr(*left, schema, planner_context)?;
+                let right =
+                    self.sql_expr_to_logical_expr(*right, schema, planner_context)?;
+                self.build_logical_expr(BinaryOperator::Spaceship, left, right, schema)
             }
 
             SQLExpr::IsTrue(expr) => Ok(Expr::IsTrue(Box::new(

datafusion/sql/tests/sql_integration.rs

@@ -30,6 +30,7 @@ use datafusion_common::{
     DFSchema, DataFusionError, Result, ScalarValue, assert_contains,
 };
 use datafusion_expr::expr_fn::create_udwf;
+use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
 use datafusion_expr::{
     ColumnarValue, CreateIndex, DdlStatement, Expr, HigherOrderFunctionArgs,
     HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl,
@@ -64,6 +65,7 @@ use datafusion_functions_window::{
 };
 use insta::{allow_duplicates, assert_snapshot};
 use rstest::rstest;
+use sqlparser::ast::BinaryOperator;
 use sqlparser::dialect::{
     DatabricksDialect, Dialect, GenericDialect, HiveDialect, MsSqlDialect, MySqlDialect,
     SnowflakeDialect,
@@ -73,6 +75,89 @@ use sqlparser::parser::Parser;
 mod cases;
 mod common;
 
+#[derive(Debug)]
+struct NullSafeComparisonPlanner;
+
+impl ExprPlanner for NullSafeComparisonPlanner {
+    fn plan_binary_op(
+        &self,
+        mut expr: RawBinaryExpr,
+        _schema: &DFSchema,
+    ) -> Result<PlannerResult<RawBinaryExpr>> {
+        if expr.op == BinaryOperator::Spaceship {
+            expr.left = Expr::Literal(ScalarValue::Int64(Some(42)), None);
+        }
+        Ok(PlannerResult::Original(expr))
+    }
+}
+
+#[derive(Debug)]
+struct PlannedNullSafeComparisonPlanner;
+
+impl ExprPlanner for PlannedNullSafeComparisonPlanner {
+    fn plan_binary_op(
+        &self,
+        expr: RawBinaryExpr,
+        _schema: &DFSchema,
+    ) -> Result<PlannerResult<RawBinaryExpr>> {
+        if expr.op == BinaryOperator::Spaceship {
+            return Ok(PlannerResult::Planned(Expr::Literal(
+                ScalarValue::Boolean(Some(true)),
+                None,
+            )));
+        }
+        Ok(PlannerResult::Original(expr))
+    }
+}
+
+#[test]
+fn null_safe_comparisons_visit_binary_expr_planners() {
+    let unchanged = logical_plan("SELECT 1 IS DISTINCT FROM 2, 1 IS NOT DISTINCT FROM 2")
+        .unwrap()
+        .display_indent()
+        .to_string();
+    assert!(
+        unchanged.contains("Int64(1) IS DISTINCT FROM Int64(2)"),
+        "{unchanged}"
+    );
+    assert!(
+        unchanged.contains("Int64(1) IS NOT DISTINCT FROM Int64(2)"),
+        "{unchanged}"
+    );
+
+    let state = MockSessionState::default()
+        .with_expr_planner(Arc::new(NullSafeComparisonPlanner));
+    let plan = logical_plan_from_state(
+        "SELECT 1 IS DISTINCT FROM 2, 1 IS NOT DISTINCT FROM 2",
+        &GenericDialect {},
+        ParserOptions::default(),
+        state,
+    )
+    .unwrap();
+    let display = plan.display_indent().to_string();
+    assert!(
+        display.contains("Int64(42) IS DISTINCT FROM Int64(2)"),
+        "{display}"
+    );
+    assert!(
+        display.contains("Int64(42) IS NOT DISTINCT FROM Int64(2)"),
+        "{display}"
+    );
+
+    let state = MockSessionState::default()
+        .with_expr_planner(Arc::new(PlannedNullSafeComparisonPlanner));
+    let plan = logical_plan_from_state(
+        "SELECT 1 IS DISTINCT FROM 2, 1 IS NOT DISTINCT FROM 2",
+        &GenericDialect {},
+        ParserOptions::default(),
+        state,
+    )
+    .unwrap();
+    let display = plan.display_indent().to_string();
+    assert!(display.contains("NOT Boolean(true)"), "{display}");
+    assert!(display.contains("Boolean(true)"), "{display}");
+}
+
 #[test]
 fn pivot_uses_literal_text_or_explicit_alias_for_column_names() {
     let plan = logical_plan_with_dialect(