Implement is [not] distinct from

ballista/rust/core/Cargo.toml

@@ -37,7 +37,7 @@ hashbrown = "0.11"
 log = "0.4"
 prost = "0.8"
 serde = {version = "1", features = ["derive"]}
-sqlparser = "0.11.0"
+sqlparser = "0.12.0"
 tokio = "1.0"
 tonic = "0.5"
 uuid = { version = "0.8", features = ["v4"] }

datafusion/Cargo.toml

@@ -52,7 +52,7 @@ ahash = "0.7"
 hashbrown = { version = "0.11", features = ["raw"] }
 arrow = { version = "^5.3", features = ["prettyprint"] }
 parquet = { version = "^5.3", features = ["arrow"] }
-sqlparser = "0.11"
+sqlparser = "0.12"
 paste = "^1.0"
 num_cpus = "1.13.0"
 chrono = "0.4"

datafusion/src/logical_plan/operators.rs

@@ -52,6 +52,10 @@ pub enum Operator {
     Like,
     /// Does not match a wildcard pattern
     NotLike,
+    /// IS DISTINCT FROM
+    IsDistinctFrom,
+    /// IS NOT DISTINCT FROM
+    IsNotDistinctFrom,
     /// Case sensitive regex match
     RegexMatch,
     /// Case insensitive regex match
@@ -84,6 +88,8 @@ impl fmt::Display for Operator {
             Operator::RegexIMatch => "~*",
             Operator::RegexNotMatch => "!~",
             Operator::RegexNotIMatch => "!~*",
+            Operator::IsDistinctFrom => "IS DISTINCT FROM",
+            Operator::IsNotDistinctFrom => "IS NOT DISTINCT FROM",
         };
         write!(f, "{}", display)
     }

datafusion/src/physical_plan/expressions/binary.rs

@@ -36,7 +36,7 @@ use arrow::compute::kernels::comparison::{
     lt_eq_utf8_scalar, lt_utf8_scalar, neq_utf8_scalar, nlike_utf8_scalar,
     regexp_is_match_utf8_scalar,
 };
-use arrow::datatypes::{DataType, Schema, TimeUnit};
+use arrow::datatypes::{ArrowNumericType, DataType, Schema, TimeUnit};
 use arrow::record_batch::RecordBatch;
 
 use crate::error::{DataFusionError, Result};
@@ -460,6 +460,9 @@ fn common_binary_type(
         | Operator::RegexIMatch
         | Operator::RegexNotMatch
         | Operator::RegexNotIMatch => string_coercion(lhs_type, rhs_type),
+        Operator::IsDistinctFrom | Operator::IsNotDistinctFrom => {
+            eq_coercion(lhs_type, rhs_type)
+        }
     };
 
     // re-write the error message of failed coercions to include the operator's information
@@ -502,7 +505,9 @@ pub fn binary_operator_data_type(
         | Operator::RegexMatch
         | Operator::RegexIMatch
         | Operator::RegexNotMatch
-        | Operator::RegexNotIMatch => Ok(DataType::Boolean),
+        | Operator::RegexNotIMatch
+        | Operator::IsDistinctFrom
+        | Operator::IsNotDistinctFrom => Ok(DataType::Boolean),
         // math operations return the same value as the common coerced type
         Operator::Plus
         | Operator::Minus
@@ -680,6 +685,10 @@ impl BinaryExpr {
             Operator::GtEq => binary_array_op!(left, right, gt_eq),
             Operator::Eq => binary_array_op!(left, right, eq),
             Operator::NotEq => binary_array_op!(left, right, neq),
+            Operator::IsDistinctFrom => binary_array_op!(left, right, is_distinct_from),
+            Operator::IsNotDistinctFrom => {
+                binary_array_op!(left, right, is_not_distinct_from)
+            }
             Operator::Plus => binary_primitive_array_op!(left, right, add),
             Operator::Minus => binary_primitive_array_op!(left, right, subtract),
             Operator::Multiply => binary_primitive_array_op!(left, right, multiply),
@@ -723,6 +732,56 @@ impl BinaryExpr {
     }
 }
 
+fn is_distinct_from<T>(
+    left: &PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> Result<BooleanArray>
+where
+    T: ArrowNumericType,
+{
+    Ok(left
+        .iter()
+        .zip(right.iter())
+        .map(|(x, y)| Some(x != y))
+        .collect())
+}
+
+fn is_distinct_from_utf8<OffsetSize: StringOffsetSizeTrait>(
+    left: &GenericStringArray<OffsetSize>,
+    right: &GenericStringArray<OffsetSize>,
+) -> Result<BooleanArray> {
+    Ok(left
+        .iter()
+        .zip(right.iter())
+        .map(|(x, y)| Some(x != y))
+        .collect())
+}
+
+fn is_not_distinct_from<T>(
+    left: &PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> Result<BooleanArray>
+where
+    T: ArrowNumericType,
+{
+    Ok(left
+        .iter()
+        .zip(right.iter())
+        .map(|(x, y)| Some(x == y))
+        .collect())
+}
+
+fn is_not_distinct_from_utf8<OffsetSize: StringOffsetSizeTrait>(
+    left: &GenericStringArray<OffsetSize>,
+    right: &GenericStringArray<OffsetSize>,
+) -> Result<BooleanArray> {
+    Ok(left
+        .iter()
+        .zip(right.iter())
+        .map(|(x, y)| Some(x == y))
+        .collect())
+}
+
 /// return two physical expressions that are optionally coerced to a
 /// common type that the binary operator supports.
 fn binary_cast(

datafusion/src/sql/planner.rs

@@ -1202,6 +1202,18 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 self.sql_expr_to_logical_expr(expr, schema)?,
             ))),
 
+            SQLExpr::IsDistinctFrom(left, right) => Ok(Expr::BinaryExpr {
+                left: Box::new(self.sql_expr_to_logical_expr(left, schema)?),
+                op: Operator::IsDistinctFrom,
+                right: Box::new(self.sql_expr_to_logical_expr(right, schema)?),
+            }),
+
+            SQLExpr::IsNotDistinctFrom(left, right) => Ok(Expr::BinaryExpr {
+                left: Box::new(self.sql_expr_to_logical_expr(left, schema)?),
+                op: Operator::IsNotDistinctFrom,
+                right: Box::new(self.sql_expr_to_logical_expr(right, schema)?),
+            }),
+
             SQLExpr::UnaryOp { ref op, ref expr } => match op {
                 UnaryOperator::Not => Ok(Expr::Not(Box::new(
                     self.sql_expr_to_logical_expr(expr, schema)?,

datafusion/tests/sql.rs

@@ -594,6 +594,50 @@ async fn select_distinct_simple_4() {
     assert_batches_sorted_eq!(expected, &actual);
 }
 
+#[tokio::test]
+async fn select_distinct_from() {
+    let mut ctx = ExecutionContext::new();
+
+    let sql = "select 
+        1 IS DISTINCT FROM CAST(NULL as INT) as a,
+        1 IS DISTINCT FROM 1 as b,
+        1 IS NOT DISTINCT FROM CAST(NULL as INT) as c,
+        1 IS NOT DISTINCT FROM 1 as d,
+        NULL IS DISTINCT FROM NULL as e,
+        NULL IS NOT DISTINCT FROM NULL as f
+    ";
+    let actual = execute_to_batches(&mut ctx, sql).await;
+    let expected = vec![
+        "+------+-------+-------+------+-------+------+",
+        "| a    | b     | c     | d    | e     | f    |",
+        "+------+-------+-------+------+-------+------+",
+        "| true | false | false | true | false | true |",
+        "+------+-------+-------+------+-------+------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+}
+
+#[tokio::test]
+async fn select_distinct_from_utf8() {
+    let mut ctx = ExecutionContext::new();
+
+    let sql = "select 
+        'x' IS DISTINCT FROM NULL as a,
+        'x' IS DISTINCT FROM 'x' as b,
+        'x' IS NOT DISTINCT FROM NULL as c,
+        'x' IS NOT DISTINCT FROM 'x' as d
+    ";
+    let actual = execute_to_batches(&mut ctx, sql).await;
+    let expected = vec![
+        "+------+-------+-------+------+",
+        "| a    | b     | c     | d    |",
+        "+------+-------+-------+------+",
+        "| true | false | false | true |",
+        "+------+-------+-------+------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+}
+
 #[tokio::test]
 async fn projection_same_fields() -> Result<()> {
     let mut ctx = ExecutionContext::new();