feat: e2e flow for unity requests via SQL (#14)

* feat: setup simple configuration

* feat: scaffold query planning

Author: roeap

Cargo.toml

@@ -18,3 +18,9 @@ delta_kernel = { version = "0.14.0", features = [
   "arrow",
   "internal-api",
 ] }
+# unitycatalog-common = { path = "../unitycatalog-rs/crates/common", default-features = false, features = [
+#   "rest-client",
+# ] }
+unitycatalog-common = { git = "https://github.com/unitycatalog-incubator/unitycatalog-rs", rev = "b840bd195c2c1c3ec5c692e245acef794f677190", default-features = false, features = [
+  "rest-client",
+] }

crates/datafusion/Cargo.toml

@@ -11,18 +11,25 @@ version.workspace = true
 
 [dependencies]
 delta_kernel = { workspace = true }
+unitycatalog-common = { workspace = true }
 
 datafusion = { version = "49", features = ["avro"] }
-datafusion-functions = { version = "49" }
 datafusion-session = { version = "49" }
 
+arrow = "55"
 async-trait = "0.1"
 bytes = "1.6.0"
 chrono = "0.4.40"
 dashmap = "6.0.1"
 futures = "0.3"
 itertools = "0.14"
 parking_lot = "0.12"
+serde = "1.0"
+serde_json = "1.0"
+sqlparser = { version = "0.55.0", default-features = false, features = [
+  "std",
+  "visitor",
+] }
 tokio = { version = "1", features = ["full"] }
 tracing = "0.1"
 url = "2.5.4"

crates/datafusion/examples/simple_read.rs

@@ -4,6 +4,7 @@ use datafusion::arrow::util::pretty::print_batches;
 use datafusion::catalog::CatalogProvider;
 use datafusion::catalog::MemoryCatalogProvider;
 use datafusion::execution::context::SessionContext;
+use deltalake_datafusion::config::OpenLakehouseConfig;
 use deltalake_datafusion::{DeltaLakeSchemaProvider, KernelContextExt as _};
 use url::Url;
 
@@ -14,7 +15,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let base = Url::from_directory_path(std::fs::canonicalize(PATH)?).unwrap();
     let table_url = base.join("column_mapping/delta")?;
 
-    let ctx = SessionContext::new().enable_delta_kernel(None);
+    let config = OpenLakehouseConfig::session_config().set_str(
+        "lakehouse.unity.uri",
+        "http://localhost:8080/api/2.1/unity-catalog/",
+    );
+
+    let ctx = SessionContext::new_with_config(config).enable_delta_kernel(None);
 
     // register and read a single table
     // TODO: right now the table provider does not internally refresh the snapshot
@@ -35,5 +41,19 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let df = df.collect().await?;
     print_batches(&df)?;
 
+    create_catalog(&ctx).await?;
+
+    Ok(())
+}
+
+async fn create_catalog(ctx: &SessionContext) -> Result<(), Box<dyn std::error::Error>> {
+    let df = ctx.sql_delta("CREATE CATALOG hello_from_df").await?;
+    let df = df.collect().await?;
+    print_batches(&df)?;
+
+    let df = ctx.sql_delta("DROP CATALOG hello_from_df").await?;
+    let df = df.collect().await?;
+    print_batches(&df)?;
+
     Ok(())
 }

crates/datafusion/src/config.rs

@@ -0,0 +1,73 @@
+use datafusion::common::extensions_options;
+use datafusion::config::ConfigExtension;
+use datafusion::prelude::SessionConfig;
+
+extensions_options! {
+    /// Configuration for integration with Unity Catalog
+    pub struct UnityCatalogConfig {
+        /// Server Url for Unity Catalog
+        pub uri: Option<String>, default = None
+        /// Access Token (PAT) for Unity Catalog server
+        pub token: Option<String>, default = None
+    }
+}
+
+impl ConfigExtension for UnityCatalogConfig {
+    const PREFIX: &'static str = "unity";
+}
+
+extensions_options! {
+    pub struct DeltaLakeConfig {
+        pub enable_caching: bool, default = false
+    }
+}
+
+impl ConfigExtension for DeltaLakeConfig {
+    const PREFIX: &'static str = "delta";
+}
+
+extensions_options! {
+    pub struct OpenLakehouseConfig {
+        /// Configuration to connect to unity catalog server
+        pub unity: UnityCatalogConfig, default = UnityCatalogConfig::default()
+        /// Configuration for interacting with delta tables
+        pub delta: DeltaLakeConfig, default = DeltaLakeConfig::default()
+    }
+}
+
+impl ConfigExtension for OpenLakehouseConfig {
+    const PREFIX: &'static str = "lakehouse";
+}
+
+impl OpenLakehouseConfig {
+    pub fn session_config() -> SessionConfig {
+        SessionConfig::new().with_option_extension(OpenLakehouseConfig::default())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use datafusion::config::{ConfigOptions, Extensions};
+
+    use super::*;
+
+    #[test]
+    fn test_delta_lake_config() {
+        let mut extensions = Extensions::new();
+        extensions.insert(OpenLakehouseConfig::default());
+
+        let mut config = ConfigOptions::new().with_extensions(extensions);
+        config
+            .set("lakehouse.unity.uri", "http://example.com")
+            .unwrap();
+        config.set("lakehouse.unity.token", "token").unwrap();
+
+        let lakehouse_config = config.extensions.get::<OpenLakehouseConfig>().unwrap();
+
+        assert_eq!(
+            lakehouse_config.unity.uri,
+            Some("http://example.com".to_string())
+        );
+        assert_eq!(lakehouse_config.unity.token, Some("token".to_string()));
+    }
+}

crates/datafusion/src/engine/expressions/to_datafusion.rs

@@ -1,8 +1,8 @@
 use datafusion::common::scalar::ScalarStructBuilder;
 use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue, not_impl_err};
+use datafusion::functions::core::expr_ext::FieldAccessor;
+use datafusion::functions::expr_fn::named_struct;
 use datafusion::logical_expr::{Expr, col, lit};
-use datafusion_functions::core::expr_ext::FieldAccessor;
-use datafusion_functions::expr_fn::named_struct;
 use delta_kernel::Predicate;
 use delta_kernel::arrow::datatypes::{DataType as ArrowDataType, Field as ArrowField};
 use delta_kernel::engine::arrow_conversion::TryIntoArrow;

crates/datafusion/src/exec.rs

@@ -0,0 +1,121 @@
+use std::{any::Any, fmt, sync::Arc};
+
+use arrow::array::RecordBatch;
+use datafusion::{
+    common::{DFSchemaRef, Statistics, internal_err},
+    error::Result,
+    execution::{SendableRecordBatchStream, TaskContext},
+    physical_expr::EquivalenceProperties,
+    physical_plan::{
+        DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties,
+        execution_plan::{Boundedness, EmissionType},
+        stream::RecordBatchStreamAdapter,
+    },
+};
+use unitycatalog_common::client::UnityCatalogClient;
+
+use crate::KernelTaskContextExt;
+
+#[async_trait::async_trait]
+pub trait ExecutableUnityCatalogStement: std::fmt::Debug + Send + Sync + 'static {
+    fn name(&self) -> &str;
+    async fn execute(&self, client: UnityCatalogClient) -> Result<RecordBatch>;
+    fn return_schema(&self) -> &DFSchemaRef;
+}
+
+pub struct UnityCatalogRequestExec {
+    request: Arc<dyn ExecutableUnityCatalogStement>,
+    cache: PlanProperties,
+}
+
+impl UnityCatalogRequestExec {
+    pub fn new(request: Arc<dyn ExecutableUnityCatalogStement>) -> Self {
+        Self {
+            cache: PlanProperties::new(
+                EquivalenceProperties::new(Arc::new(request.return_schema().as_arrow().clone())),
+                Partitioning::UnknownPartitioning(1),
+                EmissionType::Incremental,
+                Boundedness::Bounded,
+            ),
+            request,
+        }
+    }
+}
+
+impl std::fmt::Debug for UnityCatalogRequestExec {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("UnityCatalogRequestExec")
+            .field("request", &self.request)
+            .finish()
+    }
+}
+
+impl DisplayAs for UnityCatalogRequestExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "ExecuteUCStatement: statement={:?}", self.request)
+            }
+            DisplayFormatType::TreeRender => {
+                // TODO: collect info
+                write!(f, "")
+            }
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl ExecutionPlan for UnityCatalogRequestExec {
+    fn name(&self) -> &'static str {
+        Self::static_name()
+    }
+
+    /// Return a reference to Any that can be used for downcasting
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn properties(&self) -> &PlanProperties {
+        &self.cache
+    }
+
+    fn required_input_distribution(&self) -> Vec<Distribution> {
+        vec![Distribution::SinglePartition]
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        _: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(self)
+    }
+
+    /// Execute one partition and return an iterator over RecordBatch
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        if 0 != partition {
+            return internal_err!("CreateCatalogExec invalid partition {partition}");
+        }
+
+        let uc_client = context.kernel_ext()?.unity_catalog_client()?;
+        let request = self.request.clone();
+
+        Ok(Box::pin(RecordBatchStreamAdapter::new(
+            self.schema(),
+            Box::pin(futures::stream::once(async move {
+                request.execute(uc_client).await
+            })),
+        )))
+    }
+
+    fn statistics(&self) -> Result<Statistics> {
+        Ok(Statistics::new_unknown(&self.schema()))
+    }
+}

crates/datafusion/src/lib.rs

@@ -1,11 +1,15 @@
 // Datafusion error will trigger this warning.
 #![allow(clippy::result_large_err)]
 
+pub mod config;
 mod engine;
 mod error;
+mod exec;
 mod log_table_provider;
+mod planner;
 mod schema_provider;
 mod session;
+pub mod sql;
 mod table_provider;
 mod utils;