Upgrade DataFusion to 27.0.0 (#834)

* Upgrade DataFusion to 27.0.0

* Fix optimizer rule 'push_down_projection' failed

* Fix tests error of bencmarks

* Fix physical plan encode/decode error

Author: r4ntix

Cargo.toml

@@ -19,21 +19,21 @@
 members = ["ballista-cli", "ballista/client", "ballista/core", "ballista/executor", "ballista/scheduler", "benchmarks", "examples"]
 
 [workspace.dependencies]
-arrow = { version = "40.0.0" }
-arrow-flight = { version = "40.0.0", features = ["flight-sql-experimental"] }
-arrow-schema = { version = "40.0.0", default-features = false }
-configure_me = { version = "0.4.0" }
-configure_me_codegen = { version = "0.4.4" }
-datafusion = "26.0.0"
-datafusion-cli = "26.0.0"
-datafusion-proto = "26.0.0"
-object_store = "0.5.6"
-sqlparser = "0.34.0"
-tonic = { version = "0.9" }
-tonic-build = { version = "0.9", default-features = false, features = ["transport", "prost"] }
+arrow = {version = "43.0.0"}
+arrow-flight = {version = "43.0.0", features = ["flight-sql-experimental"]}
+arrow-schema = {version = "43.0.0", default-features = false}
+configure_me = {version = "0.4.0"}
+configure_me_codegen = {version = "0.4.4"}
+datafusion = "27.0.0"
+datafusion-cli = "27.0.0"
+datafusion-proto = "27.0.0"
+object_store = "0.6.1"
+sqlparser = "0.35.0"
+tonic = {version = "0.9"}
+tonic-build = {version = "0.9", default-features = false, features = ["transport", "prost"]}
 tracing = "0.1.36"
 tracing-appender = "0.2.2"
-tracing-subscriber = { version = "0.3.15", features = ["env-filter"] }
+tracing-subscriber = {version = "0.3.15", features = ["env-filter"]}
 
 # cargo build --profile release-lto
 [profile.release-lto]
@@ -54,3 +54,10 @@ opt-level = 3
 overflow-checks = false
 panic = 'unwind'
 rpath = false
+
+[patch.crates-io]
+# TODO remove on upgrade to DataFusion 28.0.0
+# fix for https://github.com/apache/arrow-datafusion/issues/6819 and https://github.com/apache/arrow-datafusion/issues/6898
+datafusion = {git = "https://github.com/apache/arrow-datafusion.git", rev = "4e2a72f6c7109d40a4986e3d05360524be078dd4"}
+datafusion-cli = {git = "https://github.com/apache/arrow-datafusion.git", rev = "4e2a72f6c7109d40a4986e3d05360524be078dd4"}
+datafusion-proto = {git = "https://github.com/apache/arrow-datafusion.git", rev = "4e2a72f6c7109d40a4986e3d05360524be078dd4"}

ballista/client/README.md

@@ -84,8 +84,8 @@ To build a simple ballista example, add the following dependencies to your `Carg
 
 ```toml
 [dependencies]
-ballista = "0.10"
-datafusion = "21.0.0"
+ballista = "0.11"
+datafusion = "27.0.0"
 tokio = "1.0"
 ```
 

ballista/core/src/execution_plans/distributed_query.rs

@@ -33,7 +33,8 @@ use datafusion::logical_expr::LogicalPlan;
 use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
+    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
+    Statistics,
 };
 use datafusion_proto::logical_plan::{
     AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec,
@@ -119,6 +120,24 @@ impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
     }
 }
 
+impl<T: 'static + AsLogicalPlan> DisplayAs for DistributedQueryExec<T> {
+    fn fmt_as(
+        &self,
+        t: DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(
+                    f,
+                    "DistributedQueryExec: scheduler_url={}",
+                    self.scheduler_url
+                )
+            }
+        }
+    }
+}
+
 impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
     fn as_any(&self) -> &dyn Any {
         self
@@ -191,22 +210,6 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
         Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
     }
 
-    fn fmt_as(
-        &self,
-        t: DisplayFormatType,
-        f: &mut std::fmt::Formatter,
-    ) -> std::fmt::Result {
-        match t {
-            DisplayFormatType::Default => {
-                write!(
-                    f,
-                    "DistributedQueryExec: scheduler_url={}",
-                    self.scheduler_url
-                )
-            }
-        }
-    }
-
     fn statistics(&self) -> Statistics {
         // This execution plan sends the logical plan to the scheduler without
         // performing the node by node conversion to a full physical plan.

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -37,7 +37,7 @@ use datafusion::error::Result;
 use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use datafusion::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
+    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
     SendableRecordBatchStream, Statistics,
 };
 use futures::{Stream, StreamExt, TryStreamExt};
@@ -83,6 +83,20 @@ impl ShuffleReaderExec {
     }
 }
 
+impl DisplayAs for ShuffleReaderExec {
+    fn fmt_as(
+        &self,
+        t: DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "ShuffleReaderExec: partitions={}", self.partition.len())
+            }
+        }
+    }
+}
+
 impl ExecutionPlan for ShuffleReaderExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -154,18 +168,6 @@ impl ExecutionPlan for ShuffleReaderExec {
         Ok(Box::pin(result))
     }
 
-    fn fmt_as(
-        &self,
-        t: DisplayFormatType,
-        f: &mut std::fmt::Formatter,
-    ) -> std::fmt::Result {
-        match t {
-            DisplayFormatType::Default => {
-                write!(f, "ShuffleReaderExec: partitions={}", self.partition.len())
-            }
-        }
-    }
-
     fn metrics(&self) -> Option<MetricsSet> {
         Some(self.metrics.clone_inner())
     }

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -47,7 +47,8 @@ use datafusion::physical_plan::metrics::{
 };
 
 use datafusion::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
+    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
+    Statistics,
 };
 use futures::{StreamExt, TryFutureExt, TryStreamExt};
 
@@ -294,6 +295,24 @@ impl ShuffleWriterExec {
     }
 }
 
+impl DisplayAs for ShuffleWriterExec {
+    fn fmt_as(
+        &self,
+        t: DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(
+                    f,
+                    "ShuffleWriterExec: {:?}",
+                    self.shuffle_output_partitioning
+                )
+            }
+        }
+    }
+}
+
 impl ExecutionPlan for ShuffleWriterExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -399,22 +418,6 @@ impl ExecutionPlan for ShuffleWriterExec {
         Some(self.metrics.clone_inner())
     }
 
-    fn fmt_as(
-        &self,
-        t: DisplayFormatType,
-        f: &mut std::fmt::Formatter,
-    ) -> std::fmt::Result {
-        match t {
-            DisplayFormatType::Default => {
-                write!(
-                    f,
-                    "ShuffleWriterExec: {:?}",
-                    self.shuffle_output_partitioning
-                )
-            }
-        }
-    }
-
     fn statistics(&self) -> Statistics {
         self.plan.statistics()
     }

ballista/core/src/execution_plans/unresolved_shuffle.rs

@@ -23,7 +23,8 @@ use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::TaskContext;
 use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::{
-    DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
+    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
+    Statistics,
 };
 
 /// UnresolvedShuffleExec represents a dependency on the results of a ShuffleWriterExec node which hasn't computed yet.
@@ -57,6 +58,20 @@ impl UnresolvedShuffleExec {
     }
 }
 
+impl DisplayAs for UnresolvedShuffleExec {
+    fn fmt_as(
+        &self,
+        t: DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "UnresolvedShuffleExec")
+            }
+        }
+    }
+}
+
 impl ExecutionPlan for UnresolvedShuffleExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -100,18 +115,6 @@ impl ExecutionPlan for UnresolvedShuffleExec {
         ))
     }
 
-    fn fmt_as(
-        &self,
-        t: DisplayFormatType,
-        f: &mut std::fmt::Formatter,
-    ) -> std::fmt::Result {
-        match t {
-            DisplayFormatType::Default => {
-                write!(f, "UnresolvedShuffleExec")
-            }
-        }
-    }
-
     fn statistics(&self) -> Statistics {
         // The full statistics are computed in the `ShuffleReaderExec` node
         // that replaces this one once the previous stage is completed.

ballista/core/src/serde/scheduler/from_proto.rs

@@ -18,7 +18,7 @@
 use chrono::{TimeZone, Utc};
 use datafusion::common::tree_node::{Transformed, TreeNode};
 use datafusion::execution::runtime_env::RuntimeEnv;
-use datafusion::logical_expr::{AggregateUDF, ScalarUDF};
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use datafusion::physical_plan::metrics::{
     Count, Gauge, MetricValue, MetricsSet, Time, Timestamp,
 };
@@ -279,6 +279,7 @@ pub fn get_task_definition<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
     runtime: Arc<RuntimeEnv>,
     scalar_functions: HashMap<String, Arc<ScalarUDF>>,
     aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    window_functions: HashMap<String, Arc<WindowUDF>>,
     codec: BallistaCodec<T, U>,
 ) -> Result<TaskDefinition, BallistaError> {
     let mut props = HashMap::new();
@@ -289,16 +290,21 @@ pub fn get_task_definition<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
 
     let mut task_scalar_functions = HashMap::new();
     let mut task_aggregate_functions = HashMap::new();
+    let mut task_window_functions = HashMap::new();
     // TODO combine the functions from Executor's functions and TaskDefinition's function resources
     for scalar_func in scalar_functions {
         task_scalar_functions.insert(scalar_func.0, scalar_func.1);
     }
     for agg_func in aggregate_functions {
         task_aggregate_functions.insert(agg_func.0, agg_func.1);
     }
+    for agg_func in window_functions {
+        task_window_functions.insert(agg_func.0, agg_func.1);
+    }
     let function_registry = Arc::new(SimpleFunctionRegistry {
         scalar_functions: task_scalar_functions,
         aggregate_functions: task_aggregate_functions,
+        window_functions: task_window_functions,
     });
 
     let encoded_plan = task.plan.as_slice();
@@ -342,6 +348,7 @@ pub fn get_task_definition_vec<
     runtime: Arc<RuntimeEnv>,
     scalar_functions: HashMap<String, Arc<ScalarUDF>>,
     aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    window_functions: HashMap<String, Arc<WindowUDF>>,
     codec: BallistaCodec<T, U>,
 ) -> Result<Vec<TaskDefinition>, BallistaError> {
     let mut props = HashMap::new();
@@ -352,16 +359,21 @@ pub fn get_task_definition_vec<
 
     let mut task_scalar_functions = HashMap::new();
     let mut task_aggregate_functions = HashMap::new();
+    let mut task_window_functions = HashMap::new();
     // TODO combine the functions from Executor's functions and TaskDefinition's function resources
     for scalar_func in scalar_functions {
         task_scalar_functions.insert(scalar_func.0, scalar_func.1);
     }
     for agg_func in aggregate_functions {
         task_aggregate_functions.insert(agg_func.0, agg_func.1);
     }
+    for agg_func in window_functions {
+        task_window_functions.insert(agg_func.0, agg_func.1);
+    }
     let function_registry = Arc::new(SimpleFunctionRegistry {
         scalar_functions: task_scalar_functions,
         aggregate_functions: task_aggregate_functions,
+        window_functions: task_window_functions,
     });
 
     let encoded_plan = multi_task.plan.as_slice();

ballista/core/src/serde/scheduler/mod.rs

@@ -25,7 +25,7 @@ use datafusion::arrow::array::{
 use datafusion::arrow::datatypes::{DataType, Field};
 use datafusion::common::DataFusionError;
 use datafusion::execution::FunctionRegistry;
-use datafusion::logical_expr::{AggregateUDF, ScalarUDF};
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::physical_plan::Partitioning;
 use serde::Serialize;
@@ -295,6 +295,7 @@ pub struct TaskDefinition {
 pub struct SimpleFunctionRegistry {
     pub scalar_functions: HashMap<String, Arc<ScalarUDF>>,
     pub aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
+    pub window_functions: HashMap<String, Arc<WindowUDF>>,
 }
 
 impl FunctionRegistry for SimpleFunctionRegistry {
@@ -321,4 +322,14 @@ impl FunctionRegistry for SimpleFunctionRegistry {
             ))
         })
     }
+
+    fn udwf(&self, name: &str) -> datafusion::common::Result<Arc<WindowUDF>> {
+        let result = self.window_functions.get(name);
+
+        result.cloned().ok_or_else(|| {
+            DataFusionError::Internal(format!(
+                "There is no UDWF named \"{name}\" in the TaskContext"
+            ))
+        })
+    }
 }

ballista/core/src/utils.rs

@@ -27,6 +27,7 @@ use datafusion::arrow::{ipc::writer::FileWriter, record_batch::RecordBatch};
 use datafusion::datasource::object_store::{
     DefaultObjectStoreRegistry, ObjectStoreRegistry,
 };
+use datafusion::datasource::physical_plan::{CsvExec, ParquetExec};
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::{
     QueryPlanner, SessionConfig, SessionContext, SessionState,
@@ -38,7 +39,6 @@ use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::common::batch_byte_size;
 use datafusion::physical_plan::empty::EmptyExec;
-use datafusion::physical_plan::file_format::{CsvExec, ParquetExec};
 use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::joins::HashJoinExec;
 use datafusion::physical_plan::metrics::MetricsSet;