From 225aebe7e8fe9c83b0ce55f524eebe99e1105b98 Mon Sep 17 00:00:00 2001
From: Batuhan Taskaya <isidentical@gmail.com>
Date: Wed, 24 Aug 2022 16:20:12 +0300
Subject: [PATCH] Implement `IS UNKNOWN`/`IS NOT UNKNOWN` operators

---
 datafusion/core/tests/sql/expr.rs | 64 +++++++++++++++++++++++++++++++
 datafusion/sql/src/planner.rs     | 11 ++++++
 2 files changed, 75 insertions(+)

diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs
index 6ccd55a02637d..cd2dd3a6c27f6 100644
--- a/datafusion/core/tests/sql/expr.rs
+++ b/datafusion/core/tests/sql/expr.rs
@@ -381,6 +381,70 @@ async fn query_is_false() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn query_is_unknown() -> Result<()> {
+    let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Boolean, true)]));
+
+    let data = RecordBatch::try_new(
+        schema.clone(),
+        vec![Arc::new(BooleanArray::from(vec![
+            Some(true),
+            Some(false),
+            None,
+        ]))],
+    )?;
+
+    let table = MemTable::try_new(schema, vec![vec![data]])?;
+
+    let ctx = SessionContext::new();
+    ctx.register_table("test", Arc::new(table))?;
+    let sql = "SELECT c1 IS UNKNOWN as t FROM test";
+    let actual = execute_to_batches(&ctx, sql).await;
+    let expected = vec![
+        "+-------+",
+        "| t     |",
+        "+-------+",
+        "| false |",
+        "| false |",
+        "| true  |",
+        "+-------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+    Ok(())
+}
+
+#[tokio::test]
+async fn query_is_not_unknown() -> Result<()> {
+    let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Boolean, true)]));
+
+    let data = RecordBatch::try_new(
+        schema.clone(),
+        vec![Arc::new(BooleanArray::from(vec![
+            Some(true),
+            Some(false),
+            None,
+        ]))],
+    )?;
+
+    let table = MemTable::try_new(schema, vec![vec![data]])?;
+
+    let ctx = SessionContext::new();
+    ctx.register_table("test", Arc::new(table))?;
+    let sql = "SELECT c1 IS NOT UNKNOWN as t FROM test";
+    let actual = execute_to_batches(&ctx, sql).await;
+    let expected = vec![
+        "+-------+",
+        "| t     |",
+        "+-------+",
+        "| true  |",
+        "| true  |",
+        "| false |",
+        "+-------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+    Ok(())
+}
+
 #[tokio::test]
 async fn query_without_from() -> Result<()> {
     // Test for SELECT <expression> without FROM.
diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs
index aaea826f06ad5..7225b779720d5 100644
--- a/datafusion/sql/src/planner.rs
+++ b/datafusion/sql/src/planner.rs
@@ -1956,6 +1956,17 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 right: Box::new(lit(false)),
             }),
 
+            SQLExpr::IsUnknown(expr) => Ok(Expr::BinaryExpr {
+                left: Box::new(self.sql_expr_to_logical_expr(*expr, schema, ctes)?),
+                op: Operator::IsNotDistinctFrom,
+                right: Box::new(lit(ScalarValue::Boolean(None))),
+            }),
+
+            SQLExpr::IsNotUnknown(expr) => Ok(Expr::BinaryExpr {
+                left: Box::new(self.sql_expr_to_logical_expr(*expr, schema, ctes)?),
+                op: Operator::IsDistinctFrom,
+                right: Box::new(lit(ScalarValue::Boolean(None))),
+            }),
 
             SQLExpr::UnaryOp { op, expr } => self.parse_sql_unary_op(op, *expr, schema, ctes),