[HSTACK] Datafusion compat - added delta table, ATS support, disable LTO

Author: adragomir

Cargo.lock

@@ -17,7 +17,7 @@
 
 [package]
 name = "datafusion-python"
-version = "45.2.0"
+version = "46.0.0"
 homepage = "https://datafusion.apache.org/python"
 repository = "https://github.com/apache/datafusion-python"
 authors = ["Apache DataFusion <dev@datafusion.apache.org>"]
@@ -38,17 +38,19 @@ tokio = { version = "1.42", features = ["macros", "rt", "rt-multi-thread", "sync
 pyo3 = { version = "0.23", features = ["extension-module", "abi3", "abi3-py38"] }
 pyo3-async-runtimes = { version = "0.23", features = ["tokio-runtime"]}
 arrow = { version = "54", features = ["pyarrow"] }
-datafusion = { version = "45.0.0", features = ["avro", "unicode_expressions"] }
-datafusion-substrait = { version = "45.0.0", optional = true }
-datafusion-proto = { version = "45.0.0" }
-datafusion-ffi = { version = "45.0.0" }
+datafusion = { version = "46.0.0", features = ["avro", "unicode_expressions"] }
+datafusion-substrait = { version = "46.0.0", optional = true }
+datafusion-proto = { version = "46.0.0" }
+datafusion-ffi = { version = "46.0.0" }
 prost = "0.13" # keep in line with `datafusion-substrait`
 uuid = { version = "1.12", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] }
 async-trait = "0.1"
 futures = "0.3"
 object_store = { version = "0.11.0", features = ["aws", "gcp", "azure", "http"] }
 url = "2"
+deltalake = { version = "0.25.0", features = ["datafusion", "azure", "s3"] }
+
 
 [build-dependencies]
 prost-types = "0.13" # keep in line with `datafusion-substrait`
@@ -58,6 +60,66 @@ pyo3-build-config = "0.23"
 name = "datafusion_python"
 crate-type = ["cdylib", "rlib"]
 
-[profile.release]
-lto = true
-codegen-units = 1
+[patch.crates-io]
+datafusion = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-catalog = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-common = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-common-runtime = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-execution = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-expr = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-expr-common = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions-aggregate = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions-aggregate-common = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions-nested = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions-table = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions-window = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-functions-window-common = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-optimizer = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-physical-expr = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-physical-expr-common = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-physical-optimizer = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-physical-plan = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-proto = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-proto-common = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-sql = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+datafusion-substrait = { git = 'https://github.com/hstack/arrow-datafusion.git', branch = 'main' }
+deltalake = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+deltalake-catalog-glue = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+deltalake-core = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+deltalake-aws = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+deltalake-azure = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+deltalake-mount = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+deltalake-sql = { git = 'https://github.com/hstack/delta-rs.git', branch = 'main' }
+reqwest = { git = 'https://github.com/hstack/reqwest.git', branch = 'disable-proxy-tunnel' }
+#datafusion-table-providers = { git = "https://github.com/hstack/datafusion-table-providers", branch = "main" }
+#datafusion = { path = "../arrow-datafusion/datafusion/core" }
+#datafusion-catalog = { path = "../arrow-datafusion/datafusion/catalog" }
+#datafusion-common = { path = "../arrow-datafusion/datafusion/common" }
+#datafusion-common-runtime = { path = "../arrow-datafusion/datafusion/common-runtime" }
+#datafusion-execution = { path = "../arrow-datafusion/datafusion/execution" }
+#datafusion-expr = { path = "../arrow-datafusion/datafusion/expr" }
+#datafusion-expr-common = { path = "../arrow-datafusion/datafusion/expr-common" }
+#datafusion-functions = { path = "../arrow-datafusion/datafusion/functions" }
+#datafusion-functions-aggregate = { path = "../arrow-datafusion/datafusion/functions-aggregate" }
+#datafusion-functions-aggregate-common = { path = "../arrow-datafusion/datafusion/functions-aggregate-common" }
+#datafusion-functions-nested = { path = "../arrow-datafusion/datafusion/functions-nested" }
+#datafusion-functions-table = { path = "../arrow-datafusion/datafusion/functions-table" }
+#datafusion-functions-window = { path = "../arrow-datafusion/datafusion/functions-window" }
+#datafusion-functions-window-common = { path = "../arrow-datafusion/datafusion/functions-window-common" }
+#datafusion-optimizer = { path = "../arrow-datafusion/datafusion/optimizer" }
+#datafusion-physical-expr = { path = "../arrow-datafusion/datafusion/physical-expr" }
+#datafusion-physical-expr-common = { path = "../arrow-datafusion/datafusion/physical-expr-common" }
+#datafusion-physical-optimizer = { path = "../arrow-datafusion/datafusion/physical-optimizer" }
+#datafusion-physical-plan = { path = "../arrow-datafusion/datafusion/physical-plan" }
+#datafusion-proto = { path = "../arrow-datafusion/datafusion/proto" }
+#datafusion-proto-common = { path = "../arrow-datafusion/datafusion/proto-common" }
+#datafusion-sql = { path = "../arrow-datafusion/datafusion/sql" }
+#datafusion-substrait = { path = "../arrow-datafusion/datafusion/substrait" }
+#deltalake = { path = "../delta-rs/crates/deltalake" }
+#deltalake-catalog-glue = { path = "../delta-rs/crates/catalog-glue" }
+#deltalake-core = { path = "../delta-rs/crates/core" }
+#deltalake-aws = { path = "../delta-rs/crates/aws" }
+#deltalake-azure = { path = "../delta-rs/crates/azure" }
+#deltalake-mount = { path = "../delta-rs/crates/mount" }
+#deltalake-sql = { path = "../delta-rs/crates/sql" }

Cargo.toml

@@ -560,6 +560,9 @@ def register_listing_table(
             file_sort_order_raw,
         )
 
+    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.
 

python/datafusion/context.py

@@ -289,6 +289,10 @@ impl PySessionContext {
         } else {
             RuntimeEnvBuilder::default()
         };
+        deltalake::azure::register_handlers(None);
+        deltalake::aws::register_handlers(None);
+        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)
@@ -390,6 +394,20 @@ 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 = 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(())
+    }
+
     /// Returns a PyDataFrame whose plan corresponds to the SQL statement.
     pub fn sql(&mut self, query: &str, py: Python) -> PyDataFusionResult<PyDataFrame> {
         let result = self.ctx.sql(query);

src/context.rs

@@ -30,10 +30,10 @@ use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::functions::core::expr_ext::FieldAccessor;
 use datafusion::logical_expr::{
     col,
-    expr::{AggregateFunction, InList, InSubquery, ScalarFunction, WindowFunction},
+    expr::{AggregateFunction, AggregateFunctionParams, InList, InSubquery, ScalarFunction, WindowFunction},
     lit, Between, BinaryExpr, Case, Cast, Expr, Like, Operator, TryCast,
 };
-
+use datafusion::logical_expr::expr::WindowFunctionParams;
 use crate::common::data_type::{DataTypeMap, NullTreatment, PyScalarValue, RexType};
 use crate::errors::{
     py_runtime_err, py_type_err, py_unsupported_variant_err, PyDataFusionError, PyDataFusionResult,
@@ -394,9 +394,9 @@ impl PyExpr {
             | Expr::InSubquery(InSubquery { expr, .. }) => Ok(vec![PyExpr::from(*expr.clone())]),
 
             // Expr variants containing a collection of Expr(s) for operands
-            Expr::AggregateFunction(AggregateFunction { args, .. })
+            Expr::AggregateFunction(AggregateFunction { params: AggregateFunctionParams { args, .. }, .. })
             | Expr::ScalarFunction(ScalarFunction { args, .. })
-            | Expr::WindowFunction(WindowFunction { args, .. }) => {
+            | Expr::WindowFunction(WindowFunction { params: WindowFunctionParams { args, .. }, .. }) => {
                 Ok(args.iter().map(|arg| PyExpr::from(arg.clone())).collect())
             }
 
@@ -575,7 +575,7 @@ impl PyExpr {
             Expr::AggregateFunction(agg_fn) => {
                 let window_fn = Expr::WindowFunction(WindowFunction::new(
                     WindowFunctionDefinition::AggregateUDF(agg_fn.func.clone()),
-                    agg_fn.args.clone(),
+                    agg_fn.params.args.clone(),
                 ));
 
                 add_builder_fns_to_window(

src/expr.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use datafusion::common::DataFusionError;
-use datafusion::logical_expr::expr::{AggregateFunction, Alias};
+use datafusion::logical_expr::expr::{AggregateFunction, AggregateFunctionParams, Alias};
 use datafusion::logical_expr::logical_plan::Aggregate;
 use datafusion::logical_expr::Expr;
 use pyo3::{prelude::*, IntoPyObjectExt};
@@ -126,7 +126,7 @@ impl PyAggregate {
         match expr {
             // TODO: This Alias logic seems to be returning some strange results that we should investigate
             Expr::Alias(Alias { expr, .. }) => self._aggregation_arguments(expr.as_ref()),
-            Expr::AggregateFunction(AggregateFunction { func: _, args, .. }) => {
+            Expr::AggregateFunction(AggregateFunction { func: _, params: AggregateFunctionParams { args, .. }}) => {
                 Ok(args.iter().map(|e| PyExpr::from(e.clone())).collect())
             }
             _ => Err(py_type_err(

src/expr/aggregate.rs

@@ -40,7 +40,7 @@ impl From<AggregateFunction> for PyAggregateFunction {
 
 impl Display for PyAggregateFunction {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        let args: Vec<String> = self.aggr.args.iter().map(|expr| expr.to_string()).collect();
+        let args: Vec<String> = self.aggr.params.args.iter().map(|expr| expr.to_string()).collect();
         write!(f, "{}({})", self.aggr.func.name(), args.join(", "))
     }
 }
@@ -54,12 +54,13 @@ impl PyAggregateFunction {
 
     /// is this a distinct aggregate such as `COUNT(DISTINCT expr)`
     fn is_distinct(&self) -> bool {
-        self.aggr.distinct
+        self.aggr.params.distinct
     }
 
     /// Get the arguments to the aggregate function
     fn args(&self) -> Vec<PyExpr> {
         self.aggr
+            .params
             .args
             .iter()
             .map(|expr| PyExpr::from(expr.clone()))

src/expr/aggregate_expr.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use datafusion::common::{DataFusionError, ScalarValue};
-use datafusion::logical_expr::expr::WindowFunction;
+use datafusion::logical_expr::expr::{WindowFunction, WindowFunctionParams};
 use datafusion::logical_expr::{Expr, Window, WindowFrame, WindowFrameBound, WindowFrameUnits};
 use pyo3::{prelude::*, IntoPyObjectExt};
 use std::fmt::{self, Display, Formatter};
@@ -118,15 +118,15 @@ impl PyWindowExpr {
     /// Returns order by columns in a window function expression
     pub fn get_sort_exprs(&self, expr: PyExpr) -> PyResult<Vec<PySortExpr>> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction { order_by, .. }) => py_sort_expr_list(&order_by),
+            Expr::WindowFunction(WindowFunction { params: WindowFunctionParams { order_by, .. }, .. }) => py_sort_expr_list(&order_by),
             other => Err(not_window_function_err(other)),
         }
     }
 
     /// Return partition by columns in a window function expression
     pub fn get_partition_exprs(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction { partition_by, .. }) => {
+            Expr::WindowFunction(WindowFunction { params: WindowFunctionParams { partition_by, .. }, .. }) => {
                 py_expr_list(&partition_by)
             }
             other => Err(not_window_function_err(other)),
@@ -136,7 +136,7 @@ impl PyWindowExpr {
     /// Return input args for window function
     pub fn get_args(&self, expr: PyExpr) -> PyResult<Vec<PyExpr>> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction { args, .. }) => py_expr_list(&args),
+            Expr::WindowFunction(WindowFunction { params: WindowFunctionParams { args, .. }, .. }) => py_expr_list(&args),
             other => Err(not_window_function_err(other)),
         }
     }
@@ -152,7 +152,7 @@ impl PyWindowExpr {
     /// Returns a Pywindow frame for a given window function expression
     pub fn get_frame(&self, expr: PyExpr) -> Option<PyWindowFrame> {
         match expr.expr.unalias() {
-            Expr::WindowFunction(WindowFunction { window_frame, .. }) => Some(window_frame.into()),
+            Expr::WindowFunction(WindowFunction { params: WindowFunctionParams {window_frame, .. }, .. }) => Some(window_frame.into()),
             _ => None,
         }
     }

src/expr/window.rs

@@ -36,7 +36,7 @@ use datafusion::execution::FunctionRegistry;
 use datafusion::functions;
 use datafusion::functions_aggregate;
 use datafusion::functions_window;
-use datafusion::logical_expr::expr::Alias;
+use datafusion::logical_expr::expr::{Alias, WindowFunctionParams};
 use datafusion::logical_expr::sqlparser::ast::NullTreatment as DFNullTreatment;
 use datafusion::logical_expr::{expr::WindowFunction, lit, Expr, WindowFunctionDefinition};
 
@@ -196,10 +196,7 @@ fn alias(expr: PyExpr, name: &str) -> PyResult<PyExpr> {
 #[pyfunction]
 fn col(name: &str) -> PyResult<PyExpr> {
     Ok(PyExpr {
-        expr: datafusion::logical_expr::Expr::Column(Column {
-            relation: None,
-            name: name.to_string(),
-        }),
+        expr: Expr::Column(Column::new_unqualified(name)),
     })
 }
 
@@ -314,19 +311,21 @@ fn window(
     Ok(PyExpr {
         expr: datafusion::logical_expr::Expr::WindowFunction(WindowFunction {
             fun,
-            args: args.into_iter().map(|x| x.expr).collect::<Vec<_>>(),
-            partition_by: partition_by
-                .unwrap_or_default()
-                .into_iter()
-                .map(|x| x.expr)
-                .collect::<Vec<_>>(),
-            order_by: order_by
-                .unwrap_or_default()
-                .into_iter()
-                .map(|x| x.into())
-                .collect::<Vec<_>>(),
-            window_frame,
-            null_treatment: None,
+            params: WindowFunctionParams {
+                args: args.into_iter().map(|x| x.expr).collect::<Vec<_>>(),
+                partition_by: partition_by
+                    .unwrap_or_default()
+                    .into_iter()
+                    .map(|x| x.expr)
+                    .collect::<Vec<_>>(),
+                order_by: order_by
+                    .unwrap_or_default()
+                    .into_iter()
+                    .map(|x| x.into())
+                    .collect::<Vec<_>>(),
+                window_frame,
+                null_treatment: None,
+            },
         }),
     })
 }