From 56d1309d572f936e889bf2f95ac12e44c2afcfb8 Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Mon, 20 Jul 2020 22:00:31 +0200 Subject: [PATCH 1/3] Added support for lit to all supported rust types. --- rust/datafusion/src/logicalplan.rs | 45 ++++++++++++++++++++++++++---- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/rust/datafusion/src/logicalplan.rs b/rust/datafusion/src/logicalplan.rs index fe711eeadca..8a125371bb3 100644 --- a/rust/datafusion/src/logicalplan.rs +++ b/rust/datafusion/src/logicalplan.rs @@ -378,9 +378,44 @@ pub fn col(name: &str) -> Expr { Expr::UnresolvedColumn(name.to_owned()) } -/// Create a literal string expression -pub fn lit_str(str: &str) -> Expr { - Expr::Literal(ScalarValue::Utf8(str.to_owned())) +/// Whether it can be represented as a literal expression +pub trait Literal { + /// convert the value to a Literal expression + fn lit(&self) -> Expr; +} + +impl Literal for &str { + fn lit(&self) -> Expr { + Expr::Literal(ScalarValue::Utf8((*self).to_owned())) + } +} + +macro_rules! make_literal { + ($TYPE:ty, $SCALAR:ident) => { + #[allow(missing_docs)] + impl Literal for $TYPE { + fn lit(&self) -> Expr { + Expr::Literal(ScalarValue::$SCALAR(self.clone())) + } + } + }; +} + +make_literal!(bool, Boolean); +make_literal!(f32, Float32); +make_literal!(f64, Float64); +make_literal!(i8, Int8); +make_literal!(i16, Int16); +make_literal!(i32, Int32); +make_literal!(i64, Int64); +make_literal!(u8, UInt8); +make_literal!(u16, UInt16); +make_literal!(u32, UInt32); +make_literal!(u64, UInt64); + +/// Create a literal expression +pub fn lit(n: T) -> Expr { + n.lit() } /// Create an convenience function representing a unary scalar function @@ -965,7 +1000,7 @@ mod tests { &employee_schema(), Some(vec![0, 3]), )? - .filter(col("state").eq(&lit_str("CO")))? + .filter(col("state").eq(&lit("CO")))? .project(vec![col("id")])? .build()?; @@ -985,7 +1020,7 @@ mod tests { CsvReadOptions::new().schema(&employee_schema()), Some(vec![0, 3]), )? - .filter(col("state").eq(&lit_str("CO")))? + .filter(col("state").eq(&lit("CO")))? .project(vec![col("id")])? .build()?; From 4b230626cf2ab36821e30cdd1caa34e2c15db751 Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Tue, 21 Jul 2020 07:59:49 +0200 Subject: [PATCH 2/3] Added lit for String. --- rust/datafusion/src/logicalplan.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/rust/datafusion/src/logicalplan.rs b/rust/datafusion/src/logicalplan.rs index 8a125371bb3..032bfb906d8 100644 --- a/rust/datafusion/src/logicalplan.rs +++ b/rust/datafusion/src/logicalplan.rs @@ -390,6 +390,12 @@ impl Literal for &str { } } +impl Literal for String { + fn lit(&self) -> Expr { + Expr::Literal(ScalarValue::Utf8((*self).to_owned())) + } +} + macro_rules! make_literal { ($TYPE:ty, $SCALAR:ident) => { #[allow(missing_docs)] From 63b3d57c25ff851e3b182abffa3a4a1ce0835407 Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Tue, 21 Jul 2020 07:34:25 +0200 Subject: [PATCH 3/3] Simplified calls using Literal. --- rust/datafusion/examples/memory_table_api.rs | 4 ++-- .../src/optimizer/projection_push_down.rs | 7 ++----- rust/datafusion/src/sql/planner.rs | 18 ++++++------------ 3 files changed, 10 insertions(+), 19 deletions(-) diff --git a/rust/datafusion/examples/memory_table_api.rs b/rust/datafusion/examples/memory_table_api.rs index bfa861266c6..937d80eeb3e 100644 --- a/rust/datafusion/examples/memory_table_api.rs +++ b/rust/datafusion/examples/memory_table_api.rs @@ -26,7 +26,7 @@ use arrow::util::pretty; use datafusion::datasource::MemTable; use datafusion::error::Result; use datafusion::execution::context::ExecutionContext; -use datafusion::logicalplan::{Expr, ScalarValue}; +use datafusion::logicalplan::lit; /// This example demonstrates basic uses of the Table API on an in-memory table fn main() -> Result<()> { @@ -54,7 +54,7 @@ fn main() -> Result<()> { let t = ctx.table("t")?; // construct an expression corresponding to "SELECT a, b FROM t WHERE b = 10" in SQL - let filter = t.col("b")?.eq(&Expr::Literal(ScalarValue::Int32(10))); + let filter = t.col("b")?.eq(&lit(10)); let t = t.select_columns(vec!["a", "b"])?.filter(filter)?; diff --git a/rust/datafusion/src/optimizer/projection_push_down.rs b/rust/datafusion/src/optimizer/projection_push_down.rs index 8fc203e6025..d1bba6edbdb 100644 --- a/rust/datafusion/src/optimizer/projection_push_down.rs +++ b/rust/datafusion/src/optimizer/projection_push_down.rs @@ -368,8 +368,8 @@ fn get_projected_schema( mod tests { use super::*; + use crate::logicalplan::lit; use crate::logicalplan::Expr::*; - use crate::logicalplan::ScalarValue; use crate::test::*; use arrow::datatypes::DataType; @@ -498,10 +498,7 @@ mod tests { fn table_scan_with_literal_projection() -> Result<()> { let table_scan = test_table_scan()?; let plan = LogicalPlanBuilder::from(&table_scan) - .project(vec![ - Expr::Literal(ScalarValue::Int64(1)), - Expr::Literal(ScalarValue::Int64(2)), - ])? + .project(vec![lit(1_i64), lit(2_i64)])? .build()?; let expected = "Projection: Int64(1), Int64(2)\ \n TableScan: test projection=Some([0])"; diff --git a/rust/datafusion/src/sql/planner.rs b/rust/datafusion/src/sql/planner.rs index d6bd2965746..8c74c5a8ad8 100644 --- a/rust/datafusion/src/sql/planner.rs +++ b/rust/datafusion/src/sql/planner.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use crate::error::{ExecutionError, Result}; use crate::logicalplan::{ - Expr, FunctionMeta, LogicalPlan, LogicalPlanBuilder, Operator, ScalarValue, + lit, Expr, FunctionMeta, LogicalPlan, LogicalPlanBuilder, Operator, ScalarValue, }; use arrow::datatypes::*; @@ -262,14 +262,10 @@ impl SqlToRel { /// Generate a relational expression from a SQL expression pub fn sql_to_rex(&self, sql: &ASTNode, schema: &Schema) -> Result { match *sql { - ASTNode::SQLValue(sqlparser::sqlast::Value::Long(n)) => { - Ok(Expr::Literal(ScalarValue::Int64(n))) - } - ASTNode::SQLValue(sqlparser::sqlast::Value::Double(n)) => { - Ok(Expr::Literal(ScalarValue::Float64(n))) - } + ASTNode::SQLValue(sqlparser::sqlast::Value::Long(n)) => Ok(lit(n)), + ASTNode::SQLValue(sqlparser::sqlast::Value::Double(n)) => Ok(lit(n)), ASTNode::SQLValue(sqlparser::sqlast::Value::SingleQuotedString(ref s)) => { - Ok(Expr::Literal(ScalarValue::Utf8(s.clone()))) + Ok(lit(s.clone())) } ASTNode::SQLAliasedExpr(ref expr, ref alias) => Ok(Alias( @@ -382,11 +378,9 @@ impl SqlToRel { .iter() .map(|a| match a { ASTNode::SQLValue(sqlparser::sqlast::Value::Long(_)) => { - Ok(Expr::Literal(ScalarValue::UInt8(1))) - } - ASTNode::SQLWildcard => { - Ok(Expr::Literal(ScalarValue::UInt8(1))) + Ok(lit(1_u8)) } + ASTNode::SQLWildcard => Ok(lit(1_u8)), _ => self.sql_to_rex(a, schema), }) .collect::>>()?;