upgrade to datafusion 55

Cargo.toml

@@ -2,8 +2,8 @@
 members = ["benchmarks", "cli", "console"]
 
 [workspace.dependencies]
-datafusion = { version = "54.0.0", default-features = false }
-datafusion-proto = { version = "54.0.0" }
+datafusion = { git = "https://github.com/apache/datafusion.git", rev = "01bf68cdc46ffbf4c226e760b0ba8a63fdb89a6b", default-features = false }
+datafusion-proto = { git = "https://github.com/apache/datafusion.git", rev = "01bf68cdc46ffbf4c226e760b0ba8a63fdb89a6b" }
 
 [package]
 name = "datafusion-distributed"
@@ -23,9 +23,9 @@ datafusion = { workspace = true, features = [
   "datetime_expressions",
 ] }
 datafusion-proto = { workspace = true }
-arrow-flight = "58"
-arrow-select = "58"
-arrow-ipc = { version = "58", features = ["zstd"] }
+arrow-flight = "59"
+arrow-select = "59"
+arrow-ipc = { version = "59", features = ["zstd"] }
 async-trait = "0.1.89"
 tokio = { version = "1.48", features = ["full"] }
 tonic = { version = "0.14.1", features = ["transport"] }
@@ -54,8 +54,8 @@ tonic-prost = "0.14.2"
 
 # integration_tests deps
 insta = { version = "1.46.0", features = ["filters"], optional = true }
-parquet = { version = "58", optional = true }
-arrow = { version = "58", optional = true, features = ["test_utils"] }
+parquet = { version = "59", optional = true }
+arrow = { version = "59", optional = true, features = ["test_utils"] }
 hyper-util = { version = "0.1.16", optional = true }
 
 [features]
@@ -74,8 +74,8 @@ sysinfo = ["dep:sysinfo"]
 datafusion-distributed-benchmarks = { path = "benchmarks" }
 structopt = "0.3"
 insta = { version = "1.46.0", features = ["filters"] }
-parquet = "58"
-arrow = { version = "58", features = ["test_utils"] }
+parquet = "59"
+arrow = { version = "59", features = ["test_utils"] }
 tokio-stream = { version = "0.1.17", features = ["sync"] }
 hyper-util = "0.1.16"
 pretty_assertions = "1.4"

benchmarks/Cargo.toml

@@ -11,7 +11,7 @@ datafusion-distributed = { path = "..", features = [
   "system-metrics",
 ] }
 tokio = { version = "1.48", features = ["full"] }
-parquet = { version = "58" }
+parquet = { version = "59" }
 structopt = { version = "0.3.26" }
 log = "0.4.27"
 serde = "1.0.219"
@@ -20,10 +20,10 @@ env_logger = "0.11.8"
 async-trait = "0.1.89"
 futures = "0.3.31"
 url = "2.5.7"
-arrow = { version = "58", features = ["test_utils"] }
+arrow = { version = "59", features = ["test_utils"] }
 tonic = { version = "0.14.1", features = ["transport"] }
-tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "438e9c2dbc25b2fff82c0efc08b3f13b5707874f" }
-tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "438e9c2dbc25b2fff82c0efc08b3f13b5707874f" }
+tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "f42708c32e7fd816b8ad8d2a880c1802a71b669f" }
+tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "f42708c32e7fd816b8ad8d2a880c1802a71b669f" }
 reqwest = "0.12"
 zip = "6.0"
 axum = "0.7"

cli/Cargo.toml

@@ -6,7 +6,7 @@ edition = "2024"
 [dependencies]
 datafusion = { workspace = true }
 datafusion-distributed = { path = "..", features = ["avro", "integration"] }
-datafusion-cli = { version = "54", default-features = false }
+datafusion-cli = { git = "https://github.com/apache/datafusion.git", rev = "01bf68cdc46ffbf4c226e760b0ba8a63fdb89a6b", default-features = false }
 tokio = { version = "1.48", features = ["full"] }
 clap = { version = "4", features = ["derive"] }
 env_logger = "0.11"

cli/src/main.rs

@@ -18,7 +18,7 @@
 // File mainly copied from https://github.com/apache/datafusion/blob/main/datafusion-cli/src/main.rs
 
 use clap::Parser;
-use datafusion::common::config_err;
+use datafusion::common::{config::ConfigNonZeroUsize, config_err};
 use datafusion::config::ConfigOptions;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::SessionStateBuilder;
@@ -214,7 +214,7 @@ fn get_session_config(args: &Args) -> Result<SessionConfig> {
         if batch_size == 0 {
             return config_err!("batch_size must be greater than 0");
         }
-        config_options.execution.batch_size = batch_size;
+        config_options.execution.batch_size = ConfigNonZeroUsize::try_new(batch_size)?;
     };
 
     // use easier to understand "tree" mode by default

console/Cargo.toml

@@ -17,7 +17,7 @@ url = "2.5.7"
 tokio-stream = "0.1.18"
 
 [dev-dependencies]
-arrow = "58"
+arrow = "59"
 async-trait = "0.1.89"
 datafusion-distributed-benchmarks = { path = "../benchmarks" }
 futures = "0.3.31"

examples/custom_execution_plan.rs

@@ -42,7 +42,7 @@ use datafusion_distributed::{
     DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
     WorkerQueryContext, display_plan_ascii,
 };
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::physical_plan::{PhysicalExtensionCodec, PhysicalProtoConverterExtension};
 use datafusion_proto::protobuf;
 use datafusion_proto::protobuf::proto_error;
 use futures::{TryStreamExt, stream};
@@ -241,6 +241,7 @@ impl PhysicalExtensionCodec for NumbersExecCodec {
         buf: &[u8],
         inputs: &[Arc<dyn ExecutionPlan>],
         _ctx: &TaskContext,
+        _proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         if !inputs.is_empty() {
             return internal_err!("NumbersExec should have no children, got {}", inputs.len());
@@ -261,7 +262,12 @@ impl PhysicalExtensionCodec for NumbersExecCodec {
         )))
     }
 
-    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
+    fn try_encode(
+        &self,
+        node: Arc<dyn ExecutionPlan>,
+        buf: &mut Vec<u8>,
+        _proto_converter: &dyn PhysicalProtoConverterExtension,
+    ) -> Result<()> {
         let Some(exec) = node.downcast_ref::<NumbersExec>() else {
             return internal_err!("Expected plan to be NumbersExec, but was {}", node.name());
         };

examples/custom_worker_url_routing.rs

@@ -44,8 +44,8 @@ use datafusion_distributed::{
     DistributedConfig, DistributedExt, DistributedLeafExec, SessionStateBuilderExt, TaskEstimation,
     TaskEstimator, TaskRoutingContext, WorkerQueryContext, display_plan_ascii,
 };
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-use datafusion_proto::protobuf;
+use datafusion_proto::physical_plan::{PhysicalExtensionCodec, PhysicalProtoConverterExtension};
+use datafusion_proto::{TryFromProto, protobuf};
 use futures::TryStreamExt;
 use prost::Message;
 use std::error::Error;
@@ -151,7 +151,7 @@ impl ExecutionPlan for CacheExec {
 fn hash_key(file_group: &FileGroup) -> usize {
     let mut hasher = DefaultHasher::new();
     for file in file_group.files() {
-        let serialized: protobuf::PartitionedFile = file.try_into().unwrap();
+        let serialized = protobuf::PartitionedFile::try_from_proto(file).unwrap();
         hasher.write(&serialized.encode_to_vec());
     }
     hasher.finish() as usize
@@ -255,14 +255,20 @@ impl PhysicalExtensionCodec for CachedFileScanCodec {
         _buf: &[u8],
         inputs: &[Arc<dyn ExecutionPlan>],
         _ctx: &TaskContext,
+        _proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let [child] = inputs else {
             return internal_err!("CacheExec expects exactly 1 child, got {}", inputs.len());
         };
         Ok(CacheExec::new(Arc::clone(child)))
     }
 
-    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, _buf: &mut Vec<u8>) -> Result<()> {
+    fn try_encode(
+        &self,
+        node: Arc<dyn ExecutionPlan>,
+        _buf: &mut Vec<u8>,
+        _proto_converter: &dyn PhysicalProtoConverterExtension,
+    ) -> Result<()> {
         if node.downcast_ref::<CacheExec>().is_none() {
             return internal_err!("Expected CacheExec, got {}", node.name());
         }

examples/work_unit_feed.rs

@@ -33,7 +33,7 @@ use datafusion_distributed::{
     DistributedExt, DistributedTaskContext, SessionStateBuilderExt, TaskEstimation, TaskEstimator,
     WorkUnitFeed, WorkUnitFeedProto, WorkUnitFeedProvider, WorkerQueryContext, display_plan_ascii,
 };
-use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+use datafusion_proto::physical_plan::{PhysicalExtensionCodec, PhysicalProtoConverterExtension};
 use datafusion_proto::protobuf::proto_error;
 use futures::stream::BoxStream;
 use futures::{StreamExt, TryStreamExt};
@@ -209,6 +209,7 @@ impl PhysicalExtensionCodec for RemoteScanExecCodec {
         buf: &[u8],
         _inputs: &[Arc<dyn ExecutionPlan>],
         _ctx: &TaskContext,
+        _proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let p = RemoteScanProto::decode(buf).map_err(|e| proto_error(format!("{e}")))?;
         let feed = WorkUnitFeed::<ChunkFeedProvider>::from_proto(
@@ -223,7 +224,12 @@ impl PhysicalExtensionCodec for RemoteScanExecCodec {
         )))
     }
 
-    fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> Result<()> {
+    fn try_encode(
+        &self,
+        node: Arc<dyn ExecutionPlan>,
+        buf: &mut Vec<u8>,
+        _proto_converter: &dyn PhysicalProtoConverterExtension,
+    ) -> Result<()> {
         let exec = node
             .downcast_ref::<RemoteScanExec>()
             .ok_or_else(|| proto_error(format!("expected RemoteScanExec, got {}", node.name())))?;

src/distributed_ext.rs

@@ -145,18 +145,18 @@ pub trait DistributedExt: Sized {
     /// # use datafusion::execution::{SessionState, FunctionRegistry, SessionStateBuilder, TaskContext};
     /// # use datafusion::physical_plan::ExecutionPlan;
     /// # use datafusion::prelude::SessionConfig;
-    /// # use datafusion_proto::physical_plan::PhysicalExtensionCodec;
+    /// # use datafusion_proto::physical_plan::{PhysicalExtensionCodec, PhysicalProtoConverterExtension};
     /// # use datafusion_distributed::{DistributedExt, WorkerQueryContext};
     ///
     /// #[derive(Debug)]
     /// struct CustomExecCodec;
     ///
     /// impl PhysicalExtensionCodec for CustomExecCodec {
-    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], ctx: &TaskContext) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
     ///         todo!()
     ///     }
     ///
-    ///     fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> datafusion::common::Result<()> {
+    ///     fn try_encode(&self, node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>, proto_converter: &dyn PhysicalProtoConverterExtension) -> datafusion::common::Result<()> {
     ///         todo!()
     ///     }
     /// }