update to df 50

adriangb · adriangb · commit a52686e2c00b · 2025-09-21T16:08:36.000-04:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -2,22 +2,26 @@
 members = ["benchmarks"]
 
 [workspace.dependencies]
-datafusion = { version = "49.0.0" }
-datafusion-proto = { version = "49.0.0" }
+datafusion = { version = "50.0.0" }
+datafusion-proto = { version = "50.0.0" }
 
 [package]
 name = "datafusion-distributed"
-version = "0.1.0"
-edition = "2021"
+version = "0.1.1"
+edition = "2024"
+
+[workspace.package]
+edition = "2024"
+rust-version = "1.86.0"
 
 [dependencies]
 datafusion = { workspace = true }
 datafusion-proto = { workspace = true }
-arrow-flight = "55.2.0"
+arrow-flight = "56.1.0"
 async-trait = "0.1.88"
 tokio = { version = "1.46.1", features = ["full"] }
-# Fixed to 0.12.3 because of arrow-flight
-tonic = { version = "0.12.3", features = ["transport"] }
+# Updated to 0.13.1 to match arrow-flight 56.1.0
+tonic = { version = "0.13.1", features = ["transport"] }
 tower = "0.5.2"
 http = "1.3.1"
 itertools = "0.14.0"
@@ -29,16 +33,16 @@ dashmap = "6.1.0"
 prost = "0.13.5"
 rand = "0.8.5"
 object_store = "0.12.3"
+chrono = { version = "0.4.42" }
 
 # integration_tests deps
 insta = { version = "1.43.1", features = ["filters"], optional = true }
-tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8", optional = true }
-tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8", optional = true }
-parquet = { version = "55.2.0", optional = true }
-arrow = { version = "55.2.0", 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 }
 tokio-stream = { version = "0.1.17", optional = true }
 hyper-util = { version = "0.1.16", optional = true }
-chrono = { version = "0.4.42", optional = true }
 
 [features]
 integration = [
@@ -49,17 +53,16 @@ integration = [
     "arrow",
     "tokio-stream",
     "hyper-util",
-    "chrono"
 ]
 
 tpch = ["integration"]
 
 [dev-dependencies]
 structopt = "0.3"
 insta = { version = "1.43.1", features = ["filters"] }
-tpchgen = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8" }
-tpchgen-arrow = { git = "https://github.com/clflushopt/tpchgen-rs", rev = "c8d823432528eed4f70fca5a1296a66c68a389a8" }
-parquet = "55.2.0"
-arrow = "55.2.0"
+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"
 tokio-stream = "0.1.17"
 hyper-util = "0.1.16"
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
@@ -9,7 +9,7 @@ datafusion = { workspace = true }
 datafusion-proto = { workspace = true }
 datafusion-distributed = { path = "..", features = ["integration"] }
 tokio = { version = "1.46.1", features = ["full"] }
-parquet = { version = "55.2.0" }
+parquet = { version = "56.1.0" }
 structopt = { version = "0.3.26" }
 log = "0.4.27"
 serde = "1.0.219"
diff --git a/rust-toolchain.toml b/rust-toolchain.toml
diff --git a/src/channel_resolver_ext.rs b/src/channel_resolver_ext.rs
@@ -3,7 +3,7 @@ use datafusion::common::exec_datafusion_err;
 use datafusion::error::DataFusionError;
 use datafusion::prelude::SessionConfig;
 use std::sync::Arc;
-use tonic::body::BoxBody;
+use tonic::body::Body;
 use url::Url;
 
 pub(crate) fn set_distributed_channel_resolver(
@@ -27,8 +27,8 @@ pub(crate) fn get_distributed_channel_resolver(
 struct ChannelResolverExtension(Arc<dyn ChannelResolver + Send + Sync>);
 
 pub type BoxCloneSyncChannel = tower::util::BoxCloneSyncService<
-    http::Request<BoxBody>,
-    http::Response<BoxBody>,
+    http::Request<Body>,
+    http::Response<Body>,
     tonic::transport::Error,
 >;
 
diff --git a/src/errors/arrow_error.rs b/src/errors/arrow_error.rs
@@ -129,6 +129,10 @@ impl ArrowErrorProto {
                 inner: Some(ArrowErrorInnerProto::RunEndIndexOverflowError(true)),
                 ctx: ctx.cloned(),
             },
+            ArrowError::OffsetOverflowError(size) => ArrowErrorProto {
+                inner: Some(ArrowErrorInnerProto::ExternalError(format!("Offset overflow error: {}", size))),
+                ctx: ctx.cloned(),
+            },
         }
     }
 
diff --git a/src/errors/schema_error.rs b/src/errors/schema_error.rs
@@ -168,7 +168,7 @@ impl TableReferenceProto {
 impl SchemaErrorProto {
     pub fn from_schema_error(err: &SchemaError, backtrace: Option<&String>) -> Self {
         match err {
-            SchemaError::AmbiguousReference { ref field } => SchemaErrorProto {
+            SchemaError::AmbiguousReference { field } => SchemaErrorProto {
                 inner: Some(SchemaErrorInnerProto::AmbiguousReference(
                     AmbiguousReferenceProto {
                         field: Some(ColumnProto::from_column(field)),
diff --git a/src/execution_plans/stage.rs b/src/execution_plans/stage.rs
@@ -530,7 +530,7 @@ pub fn display_plan(
     plan: &Arc<dyn ExecutionPlan>,
     partition_group: &[usize],
     stage_num: usize,
-    distributed: bool,
+    _distributed: bool,
 ) -> Result<String> {
     // draw all plans
     // we need to label the nodes including depth to uniquely identify them within this task
diff --git a/src/metrics/proto.rs b/src/metrics/proto.rs
@@ -381,10 +381,10 @@ pub fn metric_proto_to_df(metric: MetricProto) -> Result<Arc<Metric>, DataFusion
             )))
         }
         Some(MetricValueProto::StartTimestamp(start_ts)) => match start_ts.value {
-            Some(value) => {
+            Some(_value) => {
                 let timestamp = Timestamp::new();
                 #[cfg(feature = "integration")]
-                timestamp.set(DateTime::from_timestamp_nanos(value));
+                timestamp.set(DateTime::from_timestamp_nanos(_value));
                 #[cfg(not(feature = "integration"))]
                 { /* DateTime not available without integration feature */ }
                 Ok(Arc::new(Metric::new_with_labels(
@@ -400,10 +400,10 @@ pub fn metric_proto_to_df(metric: MetricProto) -> Result<Arc<Metric>, DataFusion
             ))),
         },
         Some(MetricValueProto::EndTimestamp(end_ts)) => match end_ts.value {
-            Some(value) => {
+            Some(_value) => {
                 let timestamp = Timestamp::new();
                 #[cfg(feature = "integration")]
-                timestamp.set(DateTime::from_timestamp_nanos(value));
+                timestamp.set(DateTime::from_timestamp_nanos(_value));
                 #[cfg(not(feature = "integration"))]
                 { /* DateTime not available without integration feature */ }
                 Ok(Arc::new(Metric::new_with_labels(
diff --git a/src/protobuf/distributed_codec.rs b/src/protobuf/distributed_codec.rs
@@ -2,7 +2,7 @@ use super::get_distributed_user_codec;
 use crate::common::ComposedPhysicalExtensionCodec;
 use crate::{ArrowFlightReadExec, PartitionIsolatorExec};
 use datafusion::arrow::datatypes::Schema;
-use datafusion::execution::FunctionRegistry;
+use datafusion::execution::{FunctionRegistry, context::SessionContext};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionConfig;
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning;
@@ -19,7 +19,7 @@ use std::sync::Arc;
 pub struct DistributedCodec;
 
 impl DistributedCodec {
-    pub fn new_combined_with_user(cfg: &SessionConfig) -> impl PhysicalExtensionCodec {
+    pub fn new_combined_with_user(cfg: &SessionConfig) -> impl PhysicalExtensionCodec + use<> {
         let mut combined_codec = ComposedPhysicalExtensionCodec::default();
         combined_codec.push(DistributedCodec {});
         if let Some(ref user_codec) = get_distributed_user_codec(cfg) {
@@ -34,7 +34,7 @@ impl PhysicalExtensionCodec for DistributedCodec {
         &self,
         buf: &[u8],
         inputs: &[Arc<dyn ExecutionPlan>],
-        registry: &dyn FunctionRegistry,
+        _registry: &dyn FunctionRegistry,
     ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         let DistributedExecProto {
             node: Some(distributed_exec_node),
@@ -56,9 +56,12 @@ impl PhysicalExtensionCodec for DistributedCodec {
                     .map(|s| s.try_into())
                     .ok_or(proto_error("ArrowFlightReadExec is missing schema"))??;
 
+                // Create a default SessionContext for the protobuf parsing
+                // TODO: This loses the original function registry, but DataFusion 50.0.0 requires SessionContext
+                let session_ctx = SessionContext::new();
                 let partioning = parse_protobuf_partitioning(
                     partitioning.as_ref(),
-                    registry,
+                    &session_ctx,
                     &schema,
                     &DistributedCodec {},
                 )?
diff --git a/src/protobuf/stage_proto.rs b/src/protobuf/stage_proto.rs
@@ -2,7 +2,7 @@ use crate::execution_plans::{ExecutionTask, StageExec};
 use datafusion::{
     common::internal_datafusion_err,
     error::{DataFusionError, Result},
-    execution::{runtime_env::RuntimeEnv, FunctionRegistry},
+    execution::{runtime_env::RuntimeEnv, SessionState, context::SessionContext},
     physical_plan::ExecutionPlan,
 };
 use datafusion_proto::{
@@ -75,21 +75,22 @@ pub fn proto_from_stage(
 
 pub fn stage_from_proto(
     msg: StageExecProto,
-    registry: &dyn FunctionRegistry,
+    session_state: &SessionState,
     runtime: &RuntimeEnv,
     codec: &dyn PhysicalExtensionCodec,
 ) -> Result<StageExec> {
     let plan_node = msg.plan.ok_or(internal_datafusion_err!(
         "ExecutionStageMsg is missing the plan"
     ))?;
 
-    let plan = plan_node.try_into_physical_plan(registry, runtime, codec)?;
+    let session_ctx = SessionContext::new_with_state(session_state.clone());
+    let plan = plan_node.try_into_physical_plan(&session_ctx, runtime, codec)?;
 
     let inputs = msg
         .inputs
         .into_iter()
         .map(|s| {
-            stage_from_proto(s, registry, runtime, codec)
+            stage_from_proto(s, session_state, runtime, codec)
                 .map(|s| Arc::new(s) as Arc<dyn ExecutionPlan>)
         })
         .collect::<Result<Vec<_>>>()?;
@@ -204,7 +205,7 @@ mod tests {
         // Convert back to ExecutionStage
         let round_trip_stage = stage_from_proto(
             decoded_msg,
-            &ctx,
+            &ctx.state(),
             ctx.runtime_env().as_ref(),
             &DefaultPhysicalExtensionCodec {},
         )?;
diff --git a/src/test_utils/insta.rs b/src/test_utils/insta.rs
@@ -12,7 +12,9 @@ macro_rules! assert_snapshot {
 }
 
 pub fn settings() -> insta::Settings {
-    env::set_var("INSTA_WORKSPACE_ROOT", env!("CARGO_MANIFEST_DIR"));
+    unsafe {
+        env::set_var("INSTA_WORKSPACE_ROOT", env!("CARGO_MANIFEST_DIR"));
+    }
     let mut settings = insta::Settings::clone_current();
     let cwd = env::current_dir().unwrap();
     let cwd = cwd.to_str().unwrap();

Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,10 @@ impl ArrowErrorProto {`
`129`	`129`	`inner: Some(ArrowErrorInnerProto::RunEndIndexOverflowError(true)),`
`130`	`130`	`ctx: ctx.cloned(),`
`131`	`131`	`},`
	`132`	`+ ArrowError::OffsetOverflowError(size) => ArrowErrorProto {`
	`133`	`+ inner: Some(ArrowErrorInnerProto::ExternalError(format!("Offset overflow error: {}", size))),`
	`134`	`+ ctx: ctx.cloned(),`
	`135`	`+ },`
`132`	`136`	`}`
`133`	`137`	`}`
`134`	`138`