Upgrade to DF v51 (#236)

Author: ahmed-mez

Cargo.lock

@@ -2,8 +2,8 @@
 members = ["benchmarks"]
 
 [workspace.dependencies]
-datafusion = { version = "50.0.0", default-features = false }
-datafusion-proto = { version = "50.0.0" }
+datafusion = { version = "51.0.0", default-features = false }
+datafusion-proto = { version = "51.0.0" }
 
 [package]
 name = "datafusion-distributed"
@@ -12,14 +12,18 @@ edition = "2024"
 
 [dependencies]
 chrono = { version = "0.4.42" }
-datafusion = { workspace = true }
+datafusion = { workspace = true, features = [
+    "parquet",
+    "sql",
+    "unicode_expressions",
+    "datetime_expressions",
+] }
 datafusion-proto = { workspace = true }
-arrow-flight = "56.1.0"
-arrow-select = "56.1.0"
+arrow-flight = "57.0.0"
+arrow-select = "57.0.0"
 async-trait = "0.1.88"
 tokio = { version = "1.46.1", features = ["full"] }
-# Updated to 0.13.1 to match arrow-flight 56.1.0
-tonic = { version = "0.13.1", features = ["transport"] }
+tonic = { version = "0.14.2", features = ["transport"] }
 tower = "0.5.2"
 http = "1.3.1"
 itertools = "0.14.0"
@@ -28,18 +32,18 @@ url = "2.5.4"
 uuid = "1.17.0"
 delegate = "0.13.4"
 dashmap = "6.1.0"
-prost = "0.13.5"
+prost = "0.14.0"
 rand = "0.8.5"
 object_store = "0.12.3"
 bytes = "1.10.1"
 pin-project = "1.1.10"
 
 # integration_tests deps
 insta = { version = "1.43.1", features = ["filters"], optional = true }
-tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "482ee68", optional = true }
-tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "482ee68", optional = true }
-parquet = { version = "56.1.0", optional = true }
-arrow = { version = "56.1.0", optional = true }
+tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "e83365a5a9101906eb9f78c5607b83bc59849acf", optional = true }
+tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "e83365a5a9101906eb9f78c5607b83bc59849acf", optional = true }
+parquet = { version = "57.0.0", optional = true }
+arrow = { version = "57.0.0", optional = true }
 tokio-stream = { version = "0.1.17", optional = true }
 hyper-util = { version = "0.1.16", optional = true }
 pretty_assertions = { version = "1.4", optional = true }
@@ -53,18 +57,18 @@ integration = [
     "arrow",
     "tokio-stream",
     "hyper-util",
-    "pretty_assertions"
+    "pretty_assertions",
 ]
 
 tpch = ["integration"]
 
 [dev-dependencies]
 structopt = "0.3"
 insta = { version = "1.43.1", features = ["filters"] }
-tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "482ee68" }
-tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "482ee68" }
-parquet = "56.1.0"
-arrow = "56.1.0"
+tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "e83365a5a9101906eb9f78c5607b83bc59849acf" }
+tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "e83365a5a9101906eb9f78c5607b83bc59849acf" }
+parquet = "57.0.0"
+arrow = "57.0.0"
 tokio-stream = "0.1.17"
 hyper-util = "0.1.16"
 pretty_assertions = "1.4"

Cargo.toml

@@ -1,3 +1,3 @@
 [toolchain]
-channel = "1.86.0"
+channel = "1.88.0"
 profile = "default"

rust-toolchain.toml

@@ -409,7 +409,7 @@ mod tests {
             }
             tokio::time::sleep(Duration::from_millis(10)).await;
         }
-        panic!("Assertion failed within {:?}", timeout);
+        panic!("Assertion failed within {timeout:?}");
     }
 
     #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
@@ -543,9 +543,9 @@ mod tests {
         let elapsed = start_time.elapsed();
 
         println!("\n=== TTLMap Lock Contention Benchmark ===");
-        println!("Tasks: {}", task_count);
-        println!("Total time: {:.2?}", elapsed);
-        println!("Average latency: {:.2} ns per operation", avg_time);
+        println!("Tasks: {task_count}");
+        println!("Total time: {elapsed:.2?}");
+        println!("Average latency: {avg_time:.2} ns per operation");
         println!("Entries remaining: {}", ttl_map.data.len());
         println!(
             "DashMap Lock contention time: {}ms",

src/common/ttl_map.rs

@@ -137,7 +137,7 @@ pub trait DistributedExt: Sized {
     /// ```
     /// # use std::sync::Arc;
     /// # use datafusion::common::DataFusionError;
-    /// # use datafusion::execution::{SessionState, FunctionRegistry, SessionStateBuilder};
+    /// # use datafusion::execution::{SessionState, FunctionRegistry, SessionStateBuilder, TaskContext};
     /// # use datafusion::physical_plan::ExecutionPlan;
     /// # use datafusion::prelude::SessionConfig;
     /// # use datafusion_proto::physical_plan::PhysicalExtensionCodec;
@@ -147,7 +147,7 @@ pub trait DistributedExt: Sized {
     /// struct CustomExecCodec;
     ///
     /// impl PhysicalExtensionCodec for CustomExecCodec {
-    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], registry: &dyn FunctionRegistry) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], ctx: &TaskContext) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
     ///         todo!()
     ///     }
     ///

src/distributed_ext.rs

@@ -112,6 +112,7 @@ use tonic::metadata::MetadataMap;
 /// 2. Ready: runs within a distributed stage and queries the next input stage over the network
 ///    using Arrow Flight.
 #[derive(Debug, Clone)]
+#[allow(clippy::large_enum_variant)]
 pub enum NetworkShuffleExec {
     Pending(NetworkShufflePendingExec),
     Ready(NetworkShuffleReadyExec),

src/execution_plans/network_shuffle.rs

@@ -50,6 +50,7 @@ use std::{fmt::Formatter, sync::Arc};
 /// └───────────────────────────┘  └───────────────────────────┘ └───────────────────────────┘     ■
 /// ```
 #[derive(Debug)]
+#[allow(clippy::large_enum_variant)]
 pub enum PartitionIsolatorExec {
     Pending(PartitionIsolatorPendingExec),
     Ready(PartitionIsolatorReadyExec),
@@ -223,7 +224,7 @@ impl ExecutionPlan for PartitionIsolatorExec {
         // then look up that index in our group and execute that partition, in this
         // example partition 8
 
-        let output_stream = match partition_group.get(partition) {
+        match partition_group.get(partition) {
             Some(actual_partition_number) => {
                 if *actual_partition_number >= input_partitions {
                     //trace!("{} returning empty stream", ctx_name);
@@ -239,8 +240,7 @@ impl ExecutionPlan for PartitionIsolatorExec {
                 Box::pin(EmptyRecordBatchStream::new(self_ready.input.schema()))
                     as SendableRecordBatchStream,
             ),
-        };
-        output_stream
+        }
     }
 }
 

src/execution_plans/partition_isolator.rs

@@ -96,7 +96,7 @@ impl ArrowFlightEndpoint {
         let stage_data = once
             .get_or_try_init(|| async {
                 let proto_node = PhysicalPlanNode::try_decode(doget.plan_proto.as_ref())?;
-                let mut plan = proto_node.try_into_physical_plan(&ctx, &self.runtime, &codec)?;
+                let mut plan = proto_node.try_into_physical_plan(&ctx.task_ctx(), &codec)?;
                 for hook in self.hooks.on_plan.iter() {
                     plan = hook(plan)
                 }

src/flight_service/do_get.rs

@@ -25,7 +25,7 @@ pub trait DistributedSessionBuilder {
     /// # use std::sync::Arc;
     /// # use async_trait::async_trait;
     /// # use datafusion::error::DataFusionError;
-    /// # use datafusion::execution::{FunctionRegistry, SessionState, SessionStateBuilder};
+    /// # use datafusion::execution::{FunctionRegistry, SessionState, SessionStateBuilder, TaskContext};
     /// # use datafusion::physical_plan::ExecutionPlan;
     /// # use datafusion_proto::physical_plan::PhysicalExtensionCodec;
     /// # use datafusion_distributed::{DistributedExt, DistributedSessionBuilder, DistributedSessionBuilderContext};
@@ -34,7 +34,7 @@ pub trait DistributedSessionBuilder {
     /// struct CustomExecCodec;
     ///
     /// impl PhysicalExtensionCodec for CustomExecCodec {
-    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], registry: &dyn FunctionRegistry) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+    ///     fn try_decode(&self, buf: &[u8], inputs: &[Arc<dyn ExecutionPlan>], ctx: &TaskContext) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
     ///         todo!()
     ///     }
     ///

src/flight_service/session_builder.rs

@@ -47,7 +47,7 @@ where
 
         let metadata =
             FlightAppMetadata::decode(flight_data.app_metadata.as_ref()).map_err(|e| {
-                FlightError::ProtocolError(format!("failed to decode app_metadata: {}", e))
+                FlightError::ProtocolError(format!("failed to decode app_metadata: {e}"))
             })?;
 
         let Some(content) = metadata.content else {