DataFusion 46.0.0

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 <[email protected]>"]
@@ -38,18 +38,18 @@ 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.23", features = ["datafusion", "azure", "s3"] }
+deltalake = { version = "0.25.0", features = ["datafusion", "azure", "s3"] }
 
 
 [build-dependencies]

Cargo.toml

@@ -6,7 +6,19 @@ WORKDIR /root
 
 RUN <<EOF
   dnf install -y epel-release
-  dnf install -y curl pkg-config openssl ca-certificates openssl-devel patchelf autoconf automake make libtool unzip clang libatomic protobuf-c-compiler openssh-clients
+  dnf install -y curl pkg-config openssl ca-certificates openssl-devel patchelf autoconf automake make libtool unzip clang libatomic openssh-clients wget
+EOF
+
+RUN <<EOF
+  # echo 1
+  wget https://github.com/protocolbuffers/protobuf/archive/refs/tags/v3.21.12.zip
+  unzip v3.21.12.zip
+  cd protobuf-3.21.12
+  autoupdate
+  ./autogen.sh
+  ./configure --prefix=/usr --disable-dependency-tracking
+  make -j8
+  make install
 EOF
 
 RUN <<EOF

Dockerfile-build-wheel.arm64

@@ -6,7 +6,19 @@ WORKDIR /root
 
 RUN <<EOF
   dnf install -y epel-release
-  dnf install -y curl pkg-config openssl ca-certificates openssl-devel patchelf autoconf automake make libtool unzip clang libatomic protobuf-c-compiler openssh-clients
+  dnf install -y curl pkg-config openssl ca-certificates openssl-devel patchelf autoconf automake make libtool unzip clang libatomic openssh-clients wget
+EOF
+
+RUN <<EOF
+  # echo 1
+  wget https://github.com/protocolbuffers/protobuf/archive/refs/tags/v3.21.12.zip
+  unzip v3.21.12.zip
+  cd protobuf-3.21.12
+  autoupdate
+  ./autogen.sh
+  ./configure --prefix=/usr --disable-dependency-tracking
+  make -j8
+  make install
 EOF
 
 RUN <<EOF

Dockerfile-build-wheel.x86_64

@@ -46,7 +46,7 @@ classifiers = [
 ]
 dependencies = ["pyarrow>=11.0.0", "typing-extensions;python_version<'3.13'"]
 #dynamic = ["version"]
-version = "44.0.0+adobe.1"
+version = "46.0.0+adobe.1"
 
 [project.urls]
 homepage = "https://datafusion.apache.org/python"

pyproject.toml

@@ -30,11 +30,12 @@ use datafusion::arrow::util::pretty;
 use datafusion::common::stats::Precision;
 use datafusion::common::{DFSchema, DataFusionError, Statistics, UnnestOptions};
 use datafusion::common::tree_node::{Transformed, TreeNode};
-use datafusion::config::{ConfigOptions, CsvOptions, TableParquetOptions};
+use datafusion::config::{CsvOptions, TableParquetOptions};
 use datafusion::dataframe::{DataFrame, DataFrameWriteOptions};
+use datafusion::datasource::memory::DataSourceExec;
 use datafusion::datasource::TableProvider;
-use datafusion::execution::runtime_env::RuntimeEnvBuilder;
-use datafusion::datasource::physical_plan::ParquetExec;
+use datafusion::datasource::physical_plan::FileScanConfig;
+use datafusion::datasource::source::DataSource;
 use datafusion::execution::SendableRecordBatchStream;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
@@ -753,25 +754,28 @@ impl DistributedPlan {
             return Ok(())
         }
         let updated_plan = self.plan().clone().transform_up(|node| {
-            if let Some(parquet) = node.as_any().downcast_ref::<ParquetExec>() {
+            if let Some(exec) = node.as_any().downcast_ref::<DataSourceExec>() {
                 // Remove redundant ranges from partition files because ParquetExec refuses to repartition
                 // if any file has a range defined (even when the range actually covers the entire file).
                 // The EnforceDistribution optimizer rule adds ranges for both full and partial files,
                 // so this tries to rever that to trigger a repartition when no files are actually split.
-                let mut file_groups = parquet.base_config().file_groups.clone();
-                for group in file_groups.iter_mut() {
-                    for file in group.iter_mut() {
-                        if let Some(range) = &file.range {
-                            if range.start == 0 && range.end == file.object_meta.size as i64 {
-                                file.range = None;  // remove redundant range
+                if let Some(file_scan) = exec.data_source().as_any().downcast_ref::<FileScanConfig>() {
+                    let mut file_groups = file_scan.file_groups.clone();
+                    for group in file_groups.iter_mut() {
+                        for file in group.iter_mut() {
+                            if let Some(range) = &file.range {
+                                if range.start == 0 && range.end == file.object_meta.size as i64 {
+                                    file.range = None;  // remove redundant range
+                                }
                             }
                         }
                     }
-                }
-                if let Some(repartitioned) = parquet.clone().into_builder().with_file_groups(file_groups)
-                    .build_arc()
-                    .repartitioned(desired_parallelism, &ConfigOptions::default())? {
-                    Ok(Transformed::yes(repartitioned))
+                    if let Some(repartitioned) = file_scan.clone().with_file_groups(file_groups)
+                        .repartitioned(desired_parallelism, 10 * 1024 * 1024, None)? {
+                        Ok(Transformed::yes(Arc::new(DataSourceExec::new(repartitioned))))
+                    } else {
+                        Ok(Transformed::no(node))
+                    }
                 } else {
                     Ok(Transformed::no(node))
                 }

src/dataframe.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,
         }
     }