From 79759f1ef1ffa9150ef7912c89b79c4246885a27 Mon Sep 17 00:00:00 2001 From: Sarah Yurick Date: Tue, 23 Aug 2022 10:33:18 -0700 Subject: [PATCH] functionality and tests --- datafusion/core/tests/sql/expr.rs | 64 +++++++++++++++++++++++++++++++ datafusion/sql/src/planner.rs | 12 ++++++ 2 files changed, 76 insertions(+) diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs index 48a95884cad10..6ccd55a02637d 100644 --- a/datafusion/core/tests/sql/expr.rs +++ b/datafusion/core/tests/sql/expr.rs @@ -317,6 +317,70 @@ async fn query_is_not_null() -> Result<()> { Ok(()) } +#[tokio::test] +async fn query_is_true() -> 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 TRUE as t FROM test"; + let actual = execute_to_batches(&ctx, sql).await; + let expected = vec![ + "+-------+", + "| t |", + "+-------+", + "| true |", + "| false |", + "| false |", + "+-------+", + ]; + assert_batches_eq!(expected, &actual); + Ok(()) +} + +#[tokio::test] +async fn query_is_false() -> 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 FALSE as f FROM test"; + let actual = execute_to_batches(&ctx, sql).await; + let expected = vec![ + "+-------+", + "| f |", + "+-------+", + "| false |", + "| true |", + "| false |", + "+-------+", + ]; + assert_batches_eq!(expected, &actual); + Ok(()) +} + #[tokio::test] async fn query_without_from() -> Result<()> { // Test for SELECT without FROM. diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 583b5cb70f813..aaea826f06ad5 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -1944,6 +1944,18 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { right: Box::new(self.sql_expr_to_logical_expr(*right, schema, ctes)?), }), + SQLExpr::IsTrue(expr) => Ok(Expr::BinaryExpr { + left: Box::new(self.sql_expr_to_logical_expr(*expr, schema, ctes)?), + op: Operator::IsNotDistinctFrom, + right: Box::new(lit(true)), + }), + + SQLExpr::IsFalse(expr) => Ok(Expr::BinaryExpr { + left: Box::new(self.sql_expr_to_logical_expr(*expr, schema, ctes)?), + op: Operator::IsNotDistinctFrom, + right: Box::new(lit(false)), + }), + SQLExpr::UnaryOp { op, expr } => self.parse_sql_unary_op(op, *expr, schema, ctes),