Support CASE WHEN and BETWEEN

dask_planner/Cargo.toml

@@ -12,8 +12,8 @@ rust-version = "1.59"
 tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "fs", "parking_lot"] }
 rand = "0.7"
 pyo3 = { version = "0.16", features = ["extension-module", "abi3", "abi3-py38"] }
-datafusion = { git="https://github.com/apache/arrow-datafusion/", rev = "ef49d2858c2aba1ea7cd5fed3b1e5feb77fc2233" }
-datafusion-expr = { git="https://github.com/apache/arrow-datafusion/", rev = "ef49d2858c2aba1ea7cd5fed3b1e5feb77fc2233" }
+datafusion = { git="https://github.com/apache/arrow-datafusion/", rev = "5927bfceeba3a4eab7988289c674d925cc82ac05" }
+datafusion-expr = { git="https://github.com/apache/arrow-datafusion/", rev = "5927bfceeba3a4eab7988289c674d925cc82ac05" }
 uuid = { version = "0.8", features = ["v4"] }
 mimalloc = { version = "*", default-features = false }
 sqlparser = "0.14.0"

dask_planner/src/expression.rs

@@ -2,7 +2,7 @@ use crate::sql::logical;
 use crate::sql::types::RexType;
 
 use pyo3::prelude::*;
-use std::convert::{From, Into};
+use std::convert::From;
 
 use datafusion::error::DataFusionError;
 
@@ -86,50 +86,47 @@ impl PyExpr {
                                                             "AGGREGATE COLUMN IS {}",
                                                             col.name
                                                         );
-                                                        col.name.clone()
+                                                        col.name.clone().to_ascii_uppercase()
                                                     }
-                                                    _ => name.clone(),
+                                                    _ => name.clone().to_ascii_uppercase(),
                                                 }
                                             }
-                                            _ => name.clone(),
+                                            _ => name.clone().to_ascii_uppercase(),
                                         }
                                     }
                                     _ => {
                                         println!("Encountered a non-Aggregate type");
-
-                                        name.clone()
+                                        name.clone().to_ascii_uppercase()
                                     }
                                 }
                             }
-                            _ => name.clone(),
+                            _ => name.clone().to_ascii_uppercase(),
                         }
                     }
-                    _ => {
-                        println!("Encountered a non Expr::Column instance");
-                        name.clone()
-                    }
+                    _ => name.clone().to_ascii_uppercase(),
                 }
             }
             Expr::Column(column) => column.name.clone(),
             Expr::ScalarVariable(..) => unimplemented!("ScalarVariable!!!"),
-            Expr::Literal(..) => unimplemented!("Literal!!!"),
-            Expr::BinaryExpr {
-                left: _,
-                op: _,
-                right: _,
-            } => {
-                // /// TODO: Examine this more deeply about whether name comes from the left or right
-                // self.column_name(left)
-                unimplemented!("BinaryExpr HERE!!!")
+            Expr::Literal(scalar_value) => {
+                println!("Scalar Value: {:?}", scalar_value);
+                unimplemented!("Literal!!!")
+            }
+            Expr::BinaryExpr { .. } => {
+                // If the BinaryExpr does not have an Alias
+                // Ex: `df.a - Int64(1)` then use the String
+                // representation of the Expr to match what is
+                // in the DFSchemaRef instance
+                format!("{}", &self.expr)
             }
             Expr::Not(..) => unimplemented!("Not!!!"),
             Expr::IsNotNull(..) => unimplemented!("IsNotNull!!!"),
             Expr::Negative(..) => unimplemented!("Negative!!!"),
             Expr::GetIndexedField { .. } => unimplemented!("GetIndexedField!!!"),
             Expr::IsNull(..) => unimplemented!("IsNull!!!"),
             Expr::Between { .. } => unimplemented!("Between!!!"),
-            Expr::Case { .. } => unimplemented!("Case!!!"),
-            Expr::Cast { .. } => unimplemented!("Cast!!!"),
+            Expr::Case { .. } => format!("{}", &self.expr),
+            Expr::Cast { .. } => format!("{}", &self.expr),
             Expr::TryCast { .. } => unimplemented!("TryCast!!!"),
             Expr::Sort { .. } => unimplemented!("Sort!!!"),
             Expr::ScalarFunction { .. } => unimplemented!("ScalarFunction!!!"),
@@ -141,6 +138,33 @@ impl PyExpr {
             _ => panic!("Nothing found!!!"),
         }
     }
+
+    fn _rex_type(&self, expr: Expr) -> RexType {
+        match &expr {
+            Expr::Alias(expr, name) => RexType::Reference,
+            Expr::Column(..) => RexType::Reference,
+            Expr::ScalarVariable(..) => RexType::Literal,
+            Expr::Literal(..) => RexType::Literal,
+            Expr::BinaryExpr { .. } => RexType::Call,
+            Expr::Not(..) => RexType::Call,
+            Expr::IsNotNull(..) => RexType::Call,
+            Expr::Negative(..) => RexType::Call,
+            Expr::GetIndexedField { .. } => RexType::Reference,
+            Expr::IsNull(..) => RexType::Call,
+            Expr::Between { .. } => RexType::Call,
+            Expr::Case { .. } => RexType::Call,
+            Expr::Cast { .. } => RexType::Call,
+            Expr::TryCast { .. } => RexType::Call,
+            Expr::Sort { .. } => RexType::Call,
+            Expr::ScalarFunction { .. } => RexType::Call,
+            Expr::AggregateFunction { .. } => RexType::Call,
+            Expr::WindowFunction { .. } => RexType::Call,
+            Expr::AggregateUDF { .. } => RexType::Call,
+            Expr::InList { .. } => RexType::Call,
+            Expr::Wildcard => RexType::Call,
+            _ => RexType::Other,
+        }
+    }
 }
 
 #[pymethods]
@@ -222,7 +246,7 @@ impl PyExpr {
     #[pyo3(name = "getRexType")]
     pub fn rex_type(&self) -> RexType {
         match &self.expr {
-            Expr::Alias(expr, name) => RexType::Reference,
+            Expr::Alias(..) => RexType::Reference,
             Expr::Column(..) => RexType::Reference,
             Expr::ScalarVariable(..) => RexType::Literal,
             Expr::Literal(..) => RexType::Literal,
@@ -281,9 +305,44 @@ impl PyExpr {
                 operands.push(py_ex);
                 Ok(operands)
             }
-            _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
-                "unknown Expr type encountered",
-            )),
+            Expr::Case {
+                expr,
+                when_then_expr,
+                else_expr,
+            } => {
+                let mut operands: Vec<PyExpr> = Vec::new();
+                match expr {
+                    Some(e) => operands.push(PyExpr::from(*e.clone(), self.input_plan.clone())),
+                    None => (),
+                };
+
+                for (when, then) in when_then_expr {
+                    operands.push(PyExpr::from(*when.clone(), self.input_plan.clone()));
+                    operands.push(PyExpr::from(*then.clone(), self.input_plan.clone()));
+                }
+
+                match else_expr {
+                    Some(e) => operands.push(PyExpr::from(*e.clone(), self.input_plan.clone())),
+                    None => (),
+                }
+                Ok(operands)
+            }
+            Expr::Between {
+                expr,
+                negated: _,
+                low,
+                high,
+            } => {
+                let mut operands: Vec<PyExpr> = Vec::new();
+                operands.push(PyExpr::from(*expr.clone(), self.input_plan.clone()));
+                operands.push(PyExpr::from(*low.clone(), self.input_plan.clone()));
+                operands.push(PyExpr::from(*high.clone(), self.input_plan.clone()));
+                Ok(operands)
+            }
+            _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
+                "unknown Expr type {:?} encountered",
+                &self.expr
+            ))),
         }
     }
 
@@ -300,9 +359,21 @@ impl PyExpr {
                 expr: _,
                 data_type: _,
             } => Ok(String::from("cast")),
-            _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
-                "Catch all triggered ....",
-            )),
+            Expr::Between {
+                expr: _,
+                negated: _,
+                low: _,
+                high: _,
+            } => Ok(String::from("between")),
+            Expr::Case {
+                expr: _,
+                when_then_expr: _,
+                else_expr: _,
+            } => Ok(String::from("case")),
+            _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
+                "Catch all triggered for get_operator_name: {:?}",
+                &self.expr
+            ))),
         }
     }
 

dask_planner/src/sql.rs

@@ -10,7 +10,7 @@ pub mod types;
 use crate::sql::exceptions::ParsingException;
 
 use datafusion::arrow::datatypes::{Field, Schema};
-use datafusion::catalog::TableReference;
+use datafusion::catalog::{ResolvedTableReference, TableReference};
 use datafusion::error::DataFusionError;
 use datafusion::physical_plan::udaf::AggregateUDF;
 use datafusion::physical_plan::udf::ScalarUDF;
@@ -46,12 +46,18 @@ use pyo3::prelude::*;
 #[pyclass(name = "DaskSQLContext", module = "dask_planner", subclass)]
 #[derive(Debug, Clone)]
 pub struct DaskSQLContext {
+    default_catalog_name: String,
     default_schema_name: String,
     schemas: HashMap<String, schema::DaskSchema>,
 }
 
 impl ContextProvider for DaskSQLContext {
-    fn get_table_provider(&self, name: TableReference) -> Option<Arc<dyn table::TableProvider>> {
+    fn get_table_provider(
+        &self,
+        name: TableReference,
+    ) -> Result<Arc<dyn table::TableProvider>, DataFusionError> {
+        let reference: ResolvedTableReference =
+            name.resolve(&self.default_catalog_name, &self.default_schema_name);
         match self.schemas.get(&self.default_schema_name) {
             Some(schema) => {
                 let mut resp = None;
@@ -67,17 +73,22 @@ impl ContextProvider for DaskSQLContext {
                         resp = Some(Schema::new(fields));
                     }
                 }
-                Some(Arc::new(DaskTableProvider::new(Arc::new(
-                    table::DaskTableSource::new(Arc::new(resp.unwrap())),
-                ))))
-            }
-            None => {
-                DataFusionError::Execution(format!(
-                    "Schema with name {} not found",
-                    &self.default_schema_name
-                ));
-                None
+
+                // If the Table is not found return None. DataFusion will handle the error propagation
+                match resp {
+                    Some(e) => Ok(Arc::new(DaskTableProvider::new(Arc::new(
+                        table::DaskTableSource::new(Arc::new(e)),
+                    )))),
+                    None => Err(DataFusionError::Plan(format!(
+                        "Table '{}.{}.{}' not found",
+                        reference.catalog, reference.schema, reference.table
+                    ))),
+                }
             }
+            None => Err(DataFusionError::Plan(format!(
+                "Unable to located Schema: '{}.{}'",
+                reference.catalog, reference.schema
+            ))),
         }
     }
 
@@ -99,8 +110,9 @@ impl ContextProvider for DaskSQLContext {
 #[pymethods]
 impl DaskSQLContext {
     #[new]
-    pub fn new(default_schema_name: String) -> Self {
+    pub fn new(default_catalog_name: String, default_schema_name: String) -> Self {
         Self {
+            default_catalog_name,
             default_schema_name,
             schemas: HashMap::new(),
         }
@@ -159,13 +171,10 @@ impl DaskSQLContext {
     ) -> PyResult<logical::PyLogicalPlan> {
         let planner = SqlToRel::new(self);
         match planner.statement_to_plan(statement.statement) {
-            Ok(k) => {
-                println!("\nLogicalPlan: {:?}\n\n", k);
-                Ok(logical::PyLogicalPlan {
-                    original_plan: k,
-                    current_node: None,
-                })
-            }
+            Ok(k) => Ok(logical::PyLogicalPlan {
+                original_plan: k,
+                current_node: None,
+            }),
             Err(e) => Err(PyErr::new::<ParsingException, _>(format!("{}", e))),
         }
     }

dask_planner/src/sql/logical.rs

@@ -1,7 +1,6 @@
 use crate::sql::table;
 use crate::sql::types::rel_data_type::RelDataType;
 use crate::sql::types::rel_data_type_field::RelDataTypeField;
-use datafusion::logical_plan::DFField;
 
 mod aggregate;
 mod filter;
@@ -12,7 +11,6 @@ pub use datafusion_expr::LogicalPlan;
 
 use datafusion::common::Result;
 use datafusion::prelude::Column;
-use pyo3::ffi::Py_FatalError;
 
 use crate::sql::exceptions::py_type_err;
 use pyo3::prelude::*;

dask_planner/src/sql/logical/projection.rs

@@ -14,12 +14,15 @@ pub struct PyProjection {
 impl PyProjection {
     /// Projection: Gets the names of the fields that should be projected
     fn projected_expressions(&mut self, local_expr: &PyExpr) -> Vec<PyExpr> {
+        println!("Exprs: {:?}", &self.projection.expr);
+        println!("Input: {:?}", &self.projection.input);
+        println!("Schema: {:?}", &self.projection.schema);
+        println!("Alias: {:?}", &self.projection.alias);
         let mut projs: Vec<PyExpr> = Vec::new();
         match &local_expr.expr {
             Expr::Alias(expr, _name) => {
-                let ex: Expr = *expr.clone();
-                let mut py_expr: PyExpr = PyExpr::from(ex, Some(self.projection.input.clone()));
-                py_expr.input_plan = local_expr.input_plan.clone();
+                let py_expr: PyExpr =
+                    PyExpr::from(*expr.clone(), Some(self.projection.input.clone()));
                 projs.extend_from_slice(self.projected_expressions(&py_expr).as_slice());
             }
             _ => projs.push(local_expr.clone()),
@@ -34,7 +37,7 @@ impl PyProjection {
     fn column_name(&mut self, expr: PyExpr) -> PyResult<String> {
         let mut val: String = String::from("OK");
         match expr.expr {
-            Expr::Alias(expr, _alias) => match expr.as_ref() {
+            Expr::Alias(expr, name) => match expr.as_ref() {
                 Expr::Column(col) => {
                     let index = self.projection.input.schema().index_of_column(col).unwrap();
                     match self.projection.input.as_ref() {
@@ -63,7 +66,7 @@ impl PyProjection {
                     val = self.column_name(py_type).unwrap();
                     println!("Setting col name to: {:?}", val);
                 }
-                _ => panic!("not supported: {:?}", expr),
+                _ => val = name.clone().to_ascii_uppercase(),
             },
             Expr::Column(col) => val = col.name.clone(),
             Expr::Cast { expr, data_type: _ } => {

dask_sql/context.py

@@ -69,6 +69,7 @@ class Context:
 
     """
 
+    DEFAULT_CATALOG_NAME = "dask_sql"
     DEFAULT_SCHEMA_NAME = "root"
 
     def __init__(self, logging_level=logging.INFO):
@@ -79,6 +80,8 @@ def __init__(self, logging_level=logging.INFO):
         # Set the logging level for this SQL context
         logging.basicConfig(level=logging_level)
 
+        # Name of the root catalog
+        self.catalog_name = self.DEFAULT_CATALOG_NAME
         # Name of the root schema
         self.schema_name = self.DEFAULT_SCHEMA_NAME
         # All schema information
@@ -87,7 +90,7 @@ def __init__(self, logging_level=logging.INFO):
         self.sql_server = None
 
         # Create the `DaskSQLContext` Rust context
-        self.context = DaskSQLContext(self.schema_name)
+        self.context = DaskSQLContext(self.catalog_name, self.schema_name)
         self.context.register_schema(self.schema_name, DaskSchema(self.schema_name))
 
         # # Register any default plugins, if nothing was registered before.

dask_sql/physical/rel/logical/limit.py

@@ -9,7 +9,7 @@
 
 if TYPE_CHECKING:
     import dask_sql
-    from dask_sql.java import org
+    from dask_planner.rust import LogicalPlan
 
 
 class DaskLimitPlugin(BaseRelPlugin):
@@ -18,11 +18,9 @@ class DaskLimitPlugin(BaseRelPlugin):
     (LIMIT).
     """
 
-    class_name = "com.dask.sql.nodes.DaskLimit"
+    class_name = "Limit"
 
-    def convert(
-        self, rel: "org.apache.calcite.rel.RelNode", context: "dask_sql.Context"
-    ) -> DataContainer:
+    def convert(self, rel: "LogicalPlan", context: "dask_sql.Context") -> DataContainer:
         (dc,) = self.assert_inputs(rel, 1, context)
         df = dc.df
         cc = dc.column_container