feat(connect): sql

src/daft-catalog/python-catalog/src/python.rs

@@ -86,6 +86,7 @@ impl DataCatalogTable for PythonTable {
 }
 
 /// Wrapper around a `daft.catalog.python_catalog.PythonCatalog`
+#[derive(Debug)]
 pub struct PythonCatalog {
     python_catalog_pyobj: PyObject,
 }

src/daft-catalog/src/data_catalog.rs

@@ -4,7 +4,7 @@ use crate::{data_catalog_table::DataCatalogTable, errors::Result};
 ///
 /// It allows registering and retrieving data sources, as well as querying their schemas.
 /// The catalog is used by the query planner to resolve table references in queries.
-pub trait DataCatalog: Sync + Send {
+pub trait DataCatalog: Sync + Send + std::fmt::Debug {
     /// Lists the fully-qualified names of tables in the catalog with the specified prefix
     fn list_tables(&self, prefix: &str) -> Result<Vec<String>>;
 

src/daft-catalog/src/lib.rs

@@ -50,6 +50,7 @@
 ///
 /// Users of Daft can register various [`DataCatalog`] with Daft, enabling
 /// discovery of tables across various [`DataCatalog`] implementations.
+#[derive(Debug, Clone, Default)]
 pub struct DaftMetaCatalog {
     /// Map of catalog names to the DataCatalog impls.
     ///
@@ -95,16 +96,26 @@
     }
 
     /// Registers a LogicalPlan with a name in the DaftMetaCatalog
-    pub fn register_named_table(&mut self, name: &str, view: LogicalPlanBuilder) -> Result<()> {
+    pub fn register_named_table(
+        &mut self,
+        name: &str,
+        view: impl Into<LogicalPlanBuilder>,
+    ) -> Result<()> {
         if !name.chars().all(|c| c.is_alphanumeric() || c == '_') {
             return Err(Error::InvalidTableName {
                 name: name.to_string(),
             });
         }
-        self.named_tables.insert(name.to_string(), view);
+        self.named_tables.insert(name.to_string(), view.into());
         Ok(())
     }
 
+    /// Check if a named table is registered in the DaftMetaCatalog
+    ///
+    pub fn contains_named_table(&self, name: &str) -> bool {
+        self.named_tables.contains_key(name)
+    }
+
     /// Provides high-level functionality for reading a table of data against a [`DaftMetaCatalog`]
     ///
     /// Resolves the provided table_identifier against the catalog:
@@ -146,6 +157,15 @@
             table_id: searched_table_name.to_string(),
         })
     }
+    /// Copy from another catalog, using tables from other in case of conflict
+    pub fn copy_from(&mut self, other: &Self) {
+        for (name, plan) in &other.named_tables {
+            self.named_tables.insert(name.clone(), plan.clone());
+        }
+        for (name, catalog) in &other.data_catalogs {
+            self.data_catalogs.insert(name.clone(), catalog.clone());
+        }
+    }
 }
 
 #[cfg(test)]

src/daft-connect/Cargo.toml

@@ -4,6 +4,7 @@ async-stream = "0.3.6"
 common-daft-config = {workspace = true, optional = true, features = ["python"]}
 common-error = {workspace = true, optional = true, features = ["python"]}
 common-file-formats = {workspace = true, optional = true, features = ["python"]}
+daft-catalog = {path = "../daft-catalog", optional = true, features = ["python"]}
 daft-core = {workspace = true, optional = true, features = ["python"]}
 daft-dsl = {workspace = true, optional = true, features = ["python"]}
 daft-local-execution = {workspace = true, optional = true, features = ["python"]}
@@ -43,7 +44,8 @@ python = [
   "dep:daft-scan",
   "dep:daft-schema",
   "dep:daft-sql",
-  "dep:daft-table"
+  "dep:daft-table",
+  "dep:daft-catalog"
 ]
 
 [lints]

src/daft-connect/src/connect_service.rs

@@ -79,17 +79,24 @@
             }
             OpType::Command(command) => {
                 let command = command.command_type.required("command_type")?;
-
                 match command {
                     CommandType::WriteOperation(op) => {
                         let result = session.execute_write_operation(op, rb).await?;
                         Ok(Response::new(result))
                     }
+                    CommandType::RegisterFunction(_) => todo!(),
+                    CommandType::CreateDataframeView(create_dataframe) => {
+                        let result = session
+                            .execute_create_dataframe_view(create_dataframe, rb)
+                            .await?;
+                        Ok(Response::new(result))
+                    }
+                    CommandType::SqlCommand(sql) => {
+                        let result = session.execute_sql_command(sql, rb).await?;
+                        Ok(Response::new(result))
+                    }
                     other => {
-                        return not_yet_implemented!(
-                            "Command type: {}",
-                            command_type_to_str(&other)
-                        )
+                        not_yet_implemented!("CommandType '{:?}'", command_type_to_str(&other))
                     }
                 }
             }

src/daft-connect/src/execute.rs

@@ -9,7 +9,7 @@
 use daft_micropartition::MicroPartition;
 use daft_ray_execution::RayEngine;
 use daft_table::Table;
-use eyre::bail;
+use eyre::{bail, Context};
 use futures::{
     stream::{self, BoxStream},
     StreamExt, TryFutureExt, TryStreamExt,
@@ -19,7 +19,8 @@
 use spark_connect::{
     relation::RelType,
     write_operation::{SaveMode, SaveType},
-    ExecutePlanResponse, Relation, ShowString, WriteOperation,
+    CreateDataFrameViewCommand, ExecutePlanResponse, Relation, ShowString, SqlCommand,
+    WriteOperation,
 };
 use tonic::{codegen::tokio_stream::wrappers::ReceiverStream, Status};
 use tracing::debug;
@@ -236,6 +237,136 @@
         Ok(Box::pin(stream))
     }
 
+    pub async fn execute_create_dataframe_view(
+        &self,
+        create_dataframe: CreateDataFrameViewCommand,
+        rb: ResponseBuilder<ExecutePlanResponse>,
+    ) -> Result<ExecuteStream, Status> {
+        let CreateDataFrameViewCommand {
+            input,
+            name,
+            is_global,
+            replace,
+        } = create_dataframe;
+
+        if is_global {
+            return not_yet_implemented!("Global dataframe view");
+        }
+
+        let input = input.required("input")?;
+        let input = SparkAnalyzer::new(self)
+            .to_logical_plan(input)
+            .await
+            .map_err(|e| {
+                Status::internal(
+                    textwrap::wrap(&format!("Error in Daft server: {e}"), 120).join("\n"),
+                )
+            })?;
+
+        {
+            let catalog = self.catalog.read().unwrap();
+            if !replace && catalog.contains_named_table(&name) {
+                return Err(Status::internal("Dataframe view already exists"));
+            }
+        }
+
+        let mut catalog = self.catalog.write().unwrap();
+
+        catalog.register_named_table(&name, input).map_err(|e| {
+            Status::internal(textwrap::wrap(&format!("Error in Daft server: {e}"), 120).join("\n"))
+        })?;
+
+        let response = rb.result_complete_response();
+        let stream = stream::once(ready(Ok(response)));
+        Ok(Box::pin(stream))
+    }
+
+    #[allow(deprecated)]
+    pub async fn execute_sql_command(
+        &self,
+        SqlCommand {
+            sql,
+            args,
+            pos_args,
+            named_arguments,
+            pos_arguments,
+            input,
+        }: SqlCommand,
+        res: ResponseBuilder<ExecutePlanResponse>,
+    ) -> Result<ExecuteStream, Status> {
+        if !args.is_empty() {
+            return not_yet_implemented!("Named arguments");
+        }
+        if !pos_args.is_empty() {
+            return not_yet_implemented!("Positional arguments");
+        }
+        if !named_arguments.is_empty() {
+            return not_yet_implemented!("Named arguments");
+        }
+        if !pos_arguments.is_empty() {
+            return not_yet_implemented!("Positional arguments");
+        }
+
+        if input.is_some() {
+            return not_yet_implemented!("Input");
+        }
+
+        let catalog = self.catalog.read().unwrap();
+        let catalog = catalog.clone();
+
+        let mut planner = daft_sql::SQLPlanner::new(catalog);
+
+        let plan = planner
+            .plan_sql(&sql)
+            .wrap_err("Error planning SQL")
+            .map_err(|e| {
+                Status::internal(
+                    textwrap::wrap(&format!("Error in Daft server: {e}"), 120).join("\n"),
+                )
+            })?;
+
+        let plan = LogicalPlanBuilder::from(plan);
+
+        // TODO: code duplication
+        let result_complete = res.result_complete_response();
+
+        let (tx, rx) = tokio::sync::mpsc::channel::<eyre::Result<ExecutePlanResponse>>(1);
+
+        let this = self.clone();
+
+        tokio::spawn(async move {
+            let execution_fut = async {
+                let mut result_stream = this.run_query(plan).await?;
+                while let Some(result) = result_stream.next().await {
+                    let result = result?;
+                    let tables = result.get_tables()?;
+                    for table in tables.as_slice() {
+                        let response = res.arrow_batch_response(table)?;
+                        if tx.send(Ok(response)).await.is_err() {
+                            return Ok(());
+                        }
+                    }
+                }
+                Ok(())
+            };
+            if let Err(e) = execution_fut.await {
+                let _ = tx.send(Err(e)).await;
+            }
+        });
+
+        let stream = ReceiverStream::new(rx);
+
+        let stream = stream
+            .map_err(|e| {
+                Status::internal(
+                    textwrap::wrap(&format!("Error in Daft server: {e}"), 120).join("\n"),
+                )
+            })
+            .chain(stream::once(ready(Ok(result_complete))));
+
+        Ok(Box::pin(stream))
+    }
+
     async fn show_string(
         &self,
         show_string: ShowString,

src/daft-connect/src/session.rs

@@ -1,5 +1,9 @@
-use std::{collections::BTreeMap, sync::Arc};
+use std::{
+    collections::BTreeMap,
+    sync::{Arc, RwLock},
+};
 
+use daft_catalog::DaftMetaCatalog;
 use daft_micropartition::partitioning::InMemoryPartitionSetCache;
 use uuid::Uuid;
 
@@ -15,6 +19,7 @@ pub struct Session {
     /// MicroPartitionSet associated with this session
     /// this will be filled up as the user runs queries
     pub(crate) psets: Arc<InMemoryPartitionSetCache>,
+    pub(crate) catalog: Arc<RwLock<DaftMetaCatalog>>,
 }
 
 impl Session {
@@ -34,6 +39,7 @@ impl Session {
             id,
             server_side_session_id,
             psets: Arc::new(InMemoryPartitionSetCache::empty()),
+            catalog: Arc::new(RwLock::new(DaftMetaCatalog::default())),
         }
     }
 

src/daft-connect/src/spark_analyzer.rs

@@ -18,6 +18,7 @@
 };
 use daft_scan::builder::{CsvScanBuilder, ParquetScanBuilder};
 use daft_schema::schema::{Schema, SchemaRef};
+use daft_sql::SQLPlanner;
 use daft_table::Table;
 use datatype::to_daft_datatype;
 pub use datatype::to_spark_datatype;
@@ -36,7 +37,7 @@
     },
     read::ReadType,
     relation::RelType,
-    Deduplicate, Expression, Limit, Range, Relation, Sort,
+    Deduplicate, Expression, Limit, Range, Relation, Sort, Sql,
 };
 use tracing::debug;
 
@@ -144,6 +145,7 @@
             RelType::ShowString(_) => unreachable!("should already be handled in execute"),
             RelType::Deduplicate(rel) => self.deduplicate(*rel).await,
             RelType::Sort(rel) => self.sort(*rel).await,
+            RelType::Sql(sql) => self.sql(sql).await,
             plan => not_yet_implemented!("relation type: \"{}\"", rel_name(&plan))?,
         }
     }
@@ -644,6 +646,40 @@
         Ok(result)
     }
 
+    #[allow(deprecated)]
+    async fn sql(&self, sql: Sql) -> eyre::Result<LogicalPlanBuilder> {
+        let Sql {
+            query,
+            args,
+            pos_args,
+            named_arguments,
+            pos_arguments,
+        } = sql;
+        if !args.is_empty() {
+            not_yet_implemented!("args")?;
+        }
+        if !pos_args.is_empty() {
+            not_yet_implemented!("pos_args")?;
+        }
+        if !named_arguments.is_empty() {
+            not_yet_implemented!("named_arguments")?;
+        }
+        if !pos_arguments.is_empty() {
+            not_yet_implemented!("pos_arguments")?;
+        }
+
+        let catalog = self
+            .session
+            .catalog
+            .read()
+            .map_err(|e| eyre::eyre!("Failed to read catalog: {e}"))?;
+        let catalog = catalog.clone();
+
+        let mut planner = SQLPlanner::new(catalog);
+        let plan = planner.plan_sql(&query)?;
+        Ok(plan.into())
+    }
+
     pub fn to_daft_expr(&self, expression: &Expression) -> eyre::Result<daft_dsl::ExprRef> {
         if let Some(common) = &expression.common {
             if common.origin.is_some() {