Upgrade to DF v51

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,17 @@ edition = "2024"
 
 [dependencies]
 chrono = { version = "0.4.42" }
-datafusion = { workspace = true }
+datafusion = { workspace = true, features = [
+    "parquet",
+    "sql",
+    "unicode_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 +31,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 +56,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

@@ -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

@@ -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

@@ -160,8 +160,8 @@ pub fn df_metrics_set_to_proto(
             Err(err) => {
                 // Check if this is the specific custom metrics error we want to filter out
                 if let DataFusionError::Internal(msg) = &err {
-                    if msg == CUSTOM_METRICS_NOT_SUPPORTED {
-                        // Filter out custom metrics error - continue processing other metrics
+                    if msg == CUSTOM_METRICS_NOT_SUPPORTED || msg == UNSUPPORTED_METRICS {
+                        // Filter out custom/unsupported metrics error - continue processing other metrics
                         continue;
                     }
                 }
@@ -191,6 +191,9 @@ pub fn metrics_set_proto_to_df(
 const CUSTOM_METRICS_NOT_SUPPORTED: &str =
     "custom metrics are not supported in metrics proto conversion";
 
+/// New DataFusion metrics that are not yet supported in proto conversion.
+const UNSUPPORTED_METRICS: &str = "metric type not supported in proto conversion";
+
 /// df_metric_to_proto converts a `datafusion::physical_plan::metrics::Metric` to a `MetricProto`. It does not consume the Arc<Metric>.
 pub fn df_metric_to_proto(metric: Arc<Metric>) -> Result<MetricProto, DataFusionError> {
     let partition = metric.partition().map(|p| p as u64);
@@ -285,6 +288,10 @@ pub fn df_metric_to_proto(metric: Arc<Metric>) -> Result<MetricProto, DataFusion
             labels,
         }),
         MetricValue::Custom { .. } => internal_err!("{}", CUSTOM_METRICS_NOT_SUPPORTED),
+        MetricValue::OutputBytes(_) | MetricValue::PruningMetrics { .. } | MetricValue::Ratio { .. } => {
+            // TODO: Support these metrics
+            internal_err!("{}", UNSUPPORTED_METRICS)
+        }
     }
 }
 

src/metrics/proto.rs

@@ -8,11 +8,11 @@ use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::internal_datafusion_err;
 use datafusion::error::DataFusionError;
-use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
+use datafusion::execution::TaskContext;
 use datafusion::physical_expr::EquivalenceProperties;
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::{ExecutionPlan, Partitioning, PlanProperties};
-use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion::prelude::SessionConfig;
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning;
 use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
 use datafusion_proto::physical_plan::{ComposedPhysicalExtensionCodec, PhysicalExtensionCodec};
@@ -40,7 +40,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
         &self,
         buf: &[u8],
         inputs: &[Arc<dyn ExecutionPlan>],
-        registry: &dyn FunctionRegistry,
+        ctx: &TaskContext,
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         let DistributedExecProto {
             node: Some(distributed_exec_node),
@@ -51,20 +51,6 @@ impl PhysicalExtensionCodec for DistributedCodec {
             ));
         };
 
-        // TODO: The PhysicalExtensionCodec trait doesn't provide access to session state,
-        // so we create a new SessionContext which loses any custom UDFs, UDAFs, and other
-        // user configurations. This is a limitation of the current trait design.
-        let state = SessionStateBuilder::new()
-            .with_scalar_functions(
-                registry
-                    .udfs()
-                    .iter()
-                    .map(|f| registry.udf(f))
-                    .collect::<Result<Vec<_>, _>>()?,
-            )
-            .build();
-        let ctx = SessionContext::from(state);
-
         fn parse_stage_proto(
             proto: Option<StageProto>,
             inputs: &[Arc<dyn ExecutionPlan>],
@@ -114,7 +100,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
 
                 let partitioning = parse_protobuf_partitioning(
                     partitioning.as_ref(),
-                    &ctx,
+                    ctx,
                     &schema,
                     &DistributedCodec {},
                 )?
@@ -138,7 +124,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
 
                 let partitioning = parse_protobuf_partitioning(
                     partitioning.as_ref(),
-                    &ctx,
+                    ctx,
                     &schema,
                     &DistributedCodec {},
                 )?
@@ -403,11 +389,12 @@ mod tests {
     use datafusion::physical_expr::LexOrdering;
     use datafusion::physical_plan::empty::EmptyExec;
     use datafusion::{
-        execution::registry::MemoryFunctionRegistry,
         physical_expr::{Partitioning, PhysicalSortExpr, expressions::Column, expressions::col},
         physical_plan::{ExecutionPlan, displayable, sorts::sort::SortExec, union::UnionExec},
     };
 
+    use datafusion::prelude::SessionContext;
+
     fn empty_exec() -> Arc<dyn ExecutionPlan> {
         Arc::new(EmptyExec::new(SchemaRef::new(Schema::empty())))
     }
@@ -429,10 +416,14 @@ mod tests {
         displayable(plan.as_ref()).indent(true).to_string()
     }
 
+    fn create_context() -> Arc<TaskContext> {
+        SessionContext::new().task_ctx()
+    }
+
     #[test]
     fn test_roundtrip_single_flight() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("a");
         let part = Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4);
@@ -442,7 +433,7 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[empty_exec()], &registry)?;
+        let decoded = codec.try_decode(&buf, &[empty_exec()], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -451,7 +442,7 @@ mod tests {
     #[test]
     fn test_roundtrip_isolator_flight() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("b");
         let flight = Arc::new(new_network_hash_shuffle_exec(
@@ -466,7 +457,7 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[flight], &registry)?;
+        let decoded = codec.try_decode(&buf, &[flight], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -475,7 +466,7 @@ mod tests {
     #[test]
     fn test_roundtrip_isolator_union() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("c");
         let left = Arc::new(new_network_hash_shuffle_exec(
@@ -489,14 +480,14 @@ mod tests {
             dummy_stage(),
         ));
 
-        let union = Arc::new(UnionExec::new(vec![left.clone(), right.clone()]));
+        let union = UnionExec::try_new(vec![left.clone(), right.clone()])?;
         let plan: Arc<dyn ExecutionPlan> =
             Arc::new(PartitionIsolatorExec::new_ready(union.clone(), 1)?);
 
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[union], &registry)?;
+        let decoded = codec.try_decode(&buf, &[union], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -505,7 +496,7 @@ mod tests {
     #[test]
     fn test_roundtrip_isolator_sort_flight() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("d");
         let flight = Arc::new(new_network_hash_shuffle_exec(
@@ -529,7 +520,7 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[sort], &registry)?;
+        let decoded = codec.try_decode(&buf, &[sort], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -538,7 +529,7 @@ mod tests {
     #[test]
     fn test_roundtrip_single_flight_coalesce() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("e");
         let plan: Arc<dyn ExecutionPlan> = Arc::new(new_network_coalesce_tasks_exec(
@@ -550,7 +541,7 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[empty_exec()], &registry)?;
+        let decoded = codec.try_decode(&buf, &[empty_exec()], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -559,7 +550,7 @@ mod tests {
     #[test]
     fn test_roundtrip_isolator_flight_coalesce() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("f");
         let flight = Arc::new(new_network_coalesce_tasks_exec(
@@ -574,7 +565,7 @@ mod tests {
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[flight], &registry)?;
+        let decoded = codec.try_decode(&buf, &[flight], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())
@@ -583,7 +574,7 @@ mod tests {
     #[test]
     fn test_roundtrip_isolator_union_coalesce() -> datafusion::common::Result<()> {
         let codec = DistributedCodec;
-        let registry = MemoryFunctionRegistry::new();
+        let ctx = create_context();
 
         let schema = schema_i32("g");
         let left = Arc::new(new_network_coalesce_tasks_exec(
@@ -597,14 +588,14 @@ mod tests {
             dummy_stage(),
         ));
 
-        let union = Arc::new(UnionExec::new(vec![left.clone(), right.clone()]));
+        let union = UnionExec::try_new(vec![left.clone(), right.clone()])?;
         let plan: Arc<dyn ExecutionPlan> =
             Arc::new(PartitionIsolatorExec::new_ready(union.clone(), 3)?);
 
         let mut buf = Vec::new();
         codec.try_encode(plan.clone(), &mut buf)?;
 
-        let decoded = codec.try_decode(&buf, &[union], &registry)?;
+        let decoded = codec.try_decode(&buf, &[union], &ctx)?;
         assert_eq!(repr(&plan), repr(&decoded));
 
         Ok(())