Execute substrait plans

dask_planner/Cargo.toml

@@ -15,6 +15,7 @@ datafusion-common = "17.0.0"
 datafusion-expr = "17.0.0"
 datafusion-optimizer = "17.0.0"
 datafusion-sql = "17.0.0"
+datafusion-substrait = "17.0.0"
 env_logger = "0.10"
 log = "^0.4"
 mimalloc = { version = "*", default-features = false }

dask_planner/src/sql.rs

@@ -9,9 +9,17 @@ pub mod statement;
 pub mod table;
 pub mod types;
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, future::Future, sync::Arc};
 
-use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use datafusion::{
+    arrow::datatypes::{DataType, Field, Schema, TimeUnit},
+    catalog::{
+        catalog::{CatalogProvider, MemoryCatalogProvider},
+        schema::MemorySchemaProvider,
+    },
+    datasource::TableProvider,
+    prelude::SessionContext,
+};
 use datafusion_common::{config::ConfigOptions, DFSchema, DataFusionError};
 use datafusion_expr::{
     logical_plan::Extension,
@@ -34,7 +42,9 @@ use datafusion_sql::{
     ResolvedTableReference,
     TableReference,
 };
+use datafusion_substrait::{consumer, serializer};
 use pyo3::prelude::*;
+use tokio::runtime::Runtime;
 
 use self::logical::{
     create_catalog_schema::CreateCatalogSchemaPlanNode,
@@ -63,6 +73,7 @@ use crate::{
             show_tables::ShowTablesPlanNode,
             PyLogicalPlan,
         },
+        table::DaskTableSource,
     },
 };
 
@@ -86,12 +97,13 @@ use crate::{
 /// # }
 /// ```
 #[pyclass(name = "DaskSQLContext", module = "dask_planner", subclass)]
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct DaskSQLContext {
     current_catalog: String,
     current_schema: String,
     schemas: HashMap<String, schema::DaskSchema>,
     options: ConfigOptions,
+    session_ctx: SessionContext,
 }
 
 impl ContextProvider for DaskSQLContext {
@@ -108,6 +120,7 @@ impl ContextProvider for DaskSQLContext {
                 reference.catalog
             )));
         }
+
         match self.schemas.get(reference.schema) {
             Some(schema) => {
                 let mut resp = None;
@@ -411,6 +424,7 @@ impl DaskSQLContext {
             current_schema: default_schema_name.to_owned(),
             schemas: HashMap::new(),
             options: ConfigOptions::new(),
+            session_ctx: SessionContext::new(),
         }
     }
 
@@ -432,7 +446,27 @@ impl DaskSQLContext {
         schema_name: String,
         schema: schema::DaskSchema,
     ) -> PyResult<bool> {
-        self.schemas.insert(schema_name, schema);
+        self.schemas.insert(schema_name.clone(), schema);
+
+        match self.session_ctx.catalog(&self.current_catalog) {
+            Some(catalog) => {
+                let schema_provider = MemorySchemaProvider::new();
+                let _result = catalog.register_schema(&schema_name, Arc::new(schema_provider));
+
+                self.session_ctx
+                    .register_catalog(self.current_catalog.clone(), catalog);
+            }
+            None => {
+                let mem_catalog = MemoryCatalogProvider::new();
+                let schema_provider = MemorySchemaProvider::new();
+                let _result = mem_catalog.register_schema(&schema_name, Arc::new(schema_provider));
+
+                // Insert the new schema into this newly created catalog
+                self.session_ctx
+                    .register_catalog(self.current_catalog.clone(), Arc::new(mem_catalog));
+            }
+        }
+
         Ok(true)
     }
 
@@ -444,7 +478,30 @@ impl DaskSQLContext {
     ) -> PyResult<bool> {
         match self.schemas.get_mut(&schema_name) {
             Some(schema) => {
-                schema.add_table(table);
+                schema.add_table(table.clone());
+
+                let tbl_ref = TableReference::Partial {
+                    schema: &self.current_schema,
+                    table: table.table_name.as_str(),
+                };
+                let tbl_src = self.get_table_provider(tbl_ref).unwrap();
+                let provider = tbl_src
+                    .as_any()
+                    .downcast_ref::<DaskTableSource>()
+                    .expect("Invalid DefaulTableSource instance");
+                let tbl_provider = provider.provider.clone() as Arc<dyn TableProvider>;
+
+                let catalog = self.session_ctx.catalog(&self.current_catalog).unwrap();
+                let schema = catalog.schema(&table.schema_name.unwrap()).unwrap();
+                let _result = schema.register_table(table.table_name.clone(), tbl_provider.clone());
+
+                let bare_tbl_ref = TableReference::Bare {
+                    table: table.table_name.as_str(),
+                };
+                let _result = self
+                    .session_ctx
+                    .register_table(bare_tbl_ref, tbl_provider.clone());
+
                 Ok(true)
             }
             None => Err(py_runtime_err(format!(
@@ -509,10 +566,39 @@ impl DaskSQLContext {
             Err(e) => Err(py_optimization_exp(e)),
         }
     }
+
+    /// Loads a `LogicalPlan` from a local Substrait protobuf file.
+    pub fn plan_from_substrait(
+        &self,
+        plan_path: String,
+        py: Python,
+    ) -> PyResult<logical::PyLogicalPlan> {
+        let result = serializer::deserialize(plan_path.as_str());
+        let plan = Self::wait_for_future(py, result).unwrap();
+
+        let result = Self::wait_for_future(
+            py,
+            consumer::from_substrait_plan(&mut self.session_ctx.clone(), &plan),
+        )
+        .map_err(DataFusionError::from)
+        .unwrap();
+
+        Ok(PyLogicalPlan::from(result))
+    }
 }
 
 /// non-Python methods
 impl DaskSQLContext {
+    /// Utility to collect rust futures with GIL released
+    pub fn wait_for_future<F: Future>(py: Python, f: F) -> F::Output
+    where
+        F: Send,
+        F::Output: Send,
+    {
+        let rt = Runtime::new().unwrap();
+        py.allow_threads(|| rt.block_on(f))
+    }
+
     /// Creates a non-optimized Relational Algebra LogicalPlan from an AST Statement
     pub fn _logical_relational_algebra(
         &self,

dask_planner/src/sql/table.rs

@@ -1,7 +1,10 @@
 use std::{any::Any, sync::Arc};
 
 use async_trait::async_trait;
-use datafusion::arrow::datatypes::{DataType, Field, SchemaRef};
+use datafusion::{
+    arrow::datatypes::{DataType, Field, SchemaRef},
+    datasource::empty::EmptyTable,
+};
 use datafusion_common::DFField;
 use datafusion_expr::{Expr, LogicalPlan, TableProviderFilterPushDown, TableSource};
 use datafusion_optimizer::utils::split_conjunction;
@@ -25,12 +28,16 @@ use crate::{
 /// DaskTable wrapper that is compatible with DataFusion logical query plans
 pub struct DaskTableSource {
     schema: SchemaRef,
+    pub provider: Arc<EmptyTable>,
 }
 
 impl DaskTableSource {
     /// Initialize a new `EmptyTable` from a schema.
     pub fn new(schema: SchemaRef) -> Self {
-        Self { schema }
+        Self {
+            schema: schema.clone(),
+            provider: Arc::new(EmptyTable::new(schema)),
+        }
     }
 }
 

dask_sql/context.py

@@ -457,6 +457,7 @@ def sql(
         return_futures: bool = True,
         dataframes: Dict[str, Union[dd.DataFrame, pd.DataFrame]] = None,
         gpu: bool = False,
+        substrait: bool = False,
         config_options: Dict[str, Any] = None,
     ) -> Union[dd.DataFrame, pd.DataFrame]:
         """
@@ -483,6 +484,9 @@ def sql(
                 to register before executing this query
             gpu (:obj:`bool`): Whether or not to load the additional Dask or pandas dataframes (if any) on GPU;
                 requires cuDF / dask-cuDF if enabled. Defaults to False.
+            substrait (:obj:`str`): If True the `sql` argument specifies a path to a Substrait plan file which is loaded
+                and ran as is without any optimizations. Otherwise it is treated as a standard SQL string and parsed by
+                the parsing engine.
             config_options (:obj:`Dict[str,Any]`): Specific configuration options to pass during
                 query execution
         Returns:
@@ -493,14 +497,19 @@ def sql(
                 for df_name, df in dataframes.items():
                     self.create_table(df_name, df, gpu=gpu)
 
-            if isinstance(sql, str):
-                rel, _ = self._get_ral(sql)
-            elif isinstance(sql, LogicalPlan):
-                rel = sql
+            if substrait:
+                logger.debug(f"Executing query using substrait plan: '{sql}'")
+                plan = self.context.plan_from_substrait(sql)
+                print(f"LogicalPlan from substrait: \n{plan}")
             else:
-                raise RuntimeError(
-                    f"Encountered unsupported `LogicalPlan` sql type: {type(sql)}"
-                )
+                if isinstance(sql, str):
+                    rel, _ = self._get_ral(sql)
+                elif isinstance(sql, LogicalPlan):
+                    rel = sql
+                else:
+                    raise RuntimeError(
+                        f"Encountered unsupported `LogicalPlan` sql type: {type(sql)}"
+                    )
 
             return self._compute_table_from_rel(rel, return_futures)
 

tests/integration/test_substrait.py

@@ -0,0 +1,12 @@
+import pandas as pd
+
+from tests.utils import assert_eq
+
+
+def test_usertable_substrait_join(c):
+    return_df = c.sql("./tests/integration/proto/df_simple.proto", substrait=True)
+    expected_df = pd.DataFrame(
+        {"user_id": [1, 1, 2, 2], "b": [3, 3, 1, 3], "c": [1, 2, 3, 3]}
+    )
+
+    assert_eq(return_df, expected_df, check_index=False)