From 38e7658c053aa6cc3f84adc25d2711961279e350 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Tue, 27 Oct 2026 04:00:09 +0300 Subject: [PATCH] Route null-safe comparisons through binary expression planners --- datafusion/sql/src/expr/mod.rs | 54 ++++++++-------- datafusion/sql/tests/sql_integration.rs | 85 +++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 28 deletions(-) diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index b2842bf1552d4..525884aba5e52 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/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 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( diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 167cfafb7ae53..eb8dd55e5cce6 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/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> { + 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> { + 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(