[HSTACK] Delta table support

Author: adragomir

Cargo.toml

@@ -51,6 +51,8 @@ futures = "0.3"
 object_store = { version = "0.12.3", features = ["aws", "gcp", "azure", "http"] }
 url = "2"
 log = "0.4.27"
+deltalake = { version = "0.28.0", features = ["datafusion", "azure", "s3"] }
+env_logger = "0.11.7"
 
 [build-dependencies]
 prost-types = "0.13.1" # keep in line with `datafusion-substrait`

python/datafusion/context.py

@@ -585,6 +585,9 @@ def register_listing_table(
             self._convert_file_sort_order(file_sort_order),
         )
 
+    def register_delta_table(self, name: str, table_uri: str, storage_opts: dict[str, str] = {}):
+        self.ctx.register_delta_table(name, table_uri, storage_opts)
+
     def sql(self, query: str, options: SQLOptions | None = None) -> DataFrame:
         """Create a :py:class:`~datafusion.DataFrame` from SQL query text.
 

src/context.rs

@@ -34,7 +34,7 @@ use pyo3::prelude::*;
 use crate::catalog::{PyCatalog, PyTable, RustWrappedPyCatalogProvider};
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
-use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionResult};
+use crate::errors::{py_datafusion_err, py_runtime_err, to_datafusion_err, PyDataFusionResult};
 use crate::expr::sort_expr::PySortExpr;
 use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
@@ -72,6 +72,7 @@ use datafusion::prelude::{
 };
 use datafusion_ffi::catalog_provider::{FFI_CatalogProvider, ForeignCatalogProvider};
 use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
+use deltalake::ensure_table_uri;
 use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple, PyType};
 use pyo3::IntoPyObjectExt;
 use tokio::task::JoinHandle;
@@ -320,6 +321,13 @@ impl PySessionContext {
         } else {
             RuntimeEnvBuilder::default()
         };
+        deltalake::azure::register_handlers(None);
+        deltalake::aws::register_handlers(None);
+        let _ = env_logger::try_init();
+
+        let config = config.set_bool("datafusion.sql_parser.enable_ident_normalization", false);
+
+
         let runtime = Arc::new(runtime_env_builder.build()?);
         let session_state = SessionStateBuilder::new()
             .with_config(config)
@@ -875,6 +883,22 @@ impl PySessionContext {
         Ok(())
     }
 
+    pub fn register_delta_table(
+        &self,
+        name: &str,
+        table_uri: &str,
+        storage_opts: HashMap<String, String>,
+        py: Python,
+    ) -> PyResult<()> {
+        deltalake::ensure_initialized();
+        let table_uri = ensure_table_uri(table_uri).map_err(py_runtime_err)?;
+        let table = deltalake::open_table_with_storage_options(table_uri, storage_opts);
+        let table = wait_for_future(py, table)?.map_err(py_datafusion_err)?;
+        self.ctx.register_table(name, Arc::new(table)).map_err(py_datafusion_err)?;
+        Ok(())
+    }
+
+
     #[pyo3(signature = (name="datafusion"))]
     pub fn catalog(&self, name: &str) -> PyResult<PyObject> {
         let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!(

src/physical_plan.rs

@@ -16,10 +16,10 @@
 // under the License.
 
 use datafusion::physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties};
-use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec};
+use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec};
 use prost::Message;
 use std::sync::Arc;
-
+use deltalake::delta_datafusion::DeltaPhysicalCodec;
 use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyBytes};
 
 use crate::{context::PySessionContext, errors::PyDataFusionResult};
@@ -59,10 +59,9 @@ impl PyExecutionPlan {
     }
 
     pub fn to_proto<'py>(&'py self, py: Python<'py>) -> PyDataFusionResult<Bound<'py, PyBytes>> {
-        let codec = DefaultPhysicalExtensionCodec {};
         let proto = datafusion_proto::protobuf::PhysicalPlanNode::try_from_physical_plan(
             self.plan.clone(),
-            &codec,
+            codec(),
         )?;
 
         let bytes = proto.encode_to_vec();
@@ -82,8 +81,7 @@ impl PyExecutionPlan {
                 ))
             })?;
 
-        let codec = DefaultPhysicalExtensionCodec {};
-        let plan = proto_plan.try_into_physical_plan(&ctx.ctx, &ctx.ctx.runtime_env(), &codec)?;
+        let plan = proto_plan.try_into_physical_plan(&ctx.ctx, &ctx.ctx.runtime_env(), codec())?;
         Ok(Self::new(plan))
     }
 
@@ -97,6 +95,11 @@ impl PyExecutionPlan {
     }
 }
 
+pub(crate) fn codec() -> &'static dyn PhysicalExtensionCodec {
+    static CODEC: DeltaPhysicalCodec = DeltaPhysicalCodec {};
+    &CODEC
+}
+
 impl From<PyExecutionPlan> for Arc<dyn ExecutionPlan> {
     fn from(plan: PyExecutionPlan) -> Arc<dyn ExecutionPlan> {
         plan.plan.clone()