From ddaf3a19b593f1b003d5244fc0e86c2a9d54334f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:05:30 -0500 Subject: [PATCH 1/4] refactor(proto): put Partitioning / PhysicalSortExpr serde on the types `Partitioning`'s protobuf conversion existed in three places: inline in `RepartitionExec::try_to_proto` / `try_from_proto`, and in `datafusion-proto`'s `serialize_partitioning` / `parse_protobuf_partitioning`. Each new plan or data source that needs to encode an output partitioning has so far copied it again. Put the single copy next to the types that own it, taking the expression-level context (`datafusion-physical-expr{,-common}` already carry the `proto` feature): - `PhysicalSortExpr::try_to_proto` / `try_from_proto` (physical-expr-common) - `Partitioning::try_to_proto` / `try_from_proto` (physical-expr) So plan hooks can reach them, `ExecutionPlanEncodeCtx` / `ExecutionPlanDecodeCtx` now back the expression-level contexts and hand one out via `expr_ctx()`; any plan hook can pass its ctx to expression-level conversions from here on. `RepartitionExec` and `datafusion-proto`'s central serializer both route through the type methods, which retires `serialize_range_partitioning`, `serialize_range_split_point`, `parse_protobuf_range_partitioning` and `parse_protobuf_range_split_point`. The wire format is unchanged. Behavior differences: out-of-range partition counts now error instead of wrapping or panicking on an `unwrap`, and a missing sort-expression child reports which field is missing rather than "Unexpected empty physical expression". Co-Authored-By: Claude Opus 5 --- .../physical-expr-common/src/sort_expr.rs | 43 ++++++ datafusion/physical-expr/src/partitioning.rs | 140 ++++++++++++++++++ datafusion/physical-plan/src/proto.rs | 47 ++++++ .../physical-plan/src/repartition/mod.rs | 133 ++--------------- .../proto/src/physical_plan/from_proto.rs | 94 ++---------- .../proto/src/physical_plan/to_proto.rs | 73 +-------- 6 files changed, 262 insertions(+), 268 deletions(-) diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 84ffb92eaa600..82f7a6e5fc33f 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -183,6 +183,49 @@ impl PhysicalSortExpr { } } +/// Protobuf conversions for [`PhysicalSortExpr`]. +/// +/// This is the flat [`PhysicalSortExprNode`] representation used wherever the +/// wire format stores an ordering (scan output orderings, range partitioning, +/// window frames, …). It is *not* the `PhysicalExprNode::Sort` wrapping that +/// `SortExec` uses for its own `expr` field. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +impl PhysicalSortExpr { + /// Serialize this sort expression, encoding its child expression through + /// `ctx`. + pub fn try_to_proto( + &self, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result { + Ok(datafusion_proto_models::protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_child(&self.expr)?)), + asc: !self.options.descending, + nulls_first: self.options.nulls_first, + }) + } + + /// Reconstruct a [`PhysicalSortExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalSortExprNode, + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result { + let expr = ctx.decode_required_expression( + node.expr.as_deref(), + "PhysicalSortExpr", + "expr", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + } +} + impl PartialEq for PhysicalSortExpr { fn eq(&self, other: &Self) -> bool { self.options == other.options && self.expr.eq(&other.expr) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 59d36c4efc1bb..9a24658411300 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -515,6 +515,146 @@ impl Partitioning { } } +/// Protobuf conversions for [`Partitioning`]. +/// +/// Child expressions (hash keys, range orderings) and `ScalarValue` split +/// points are (de)serialized through the expression-level context, so this is +/// the single copy of the partitioning wire format: `RepartitionExec`, +/// `FileScanConfig` and `datafusion-proto`'s central serializer all route +/// through it. +/// +/// [`protobuf::Partitioning`]: datafusion_proto_models::protobuf::Partitioning +#[cfg(feature = "proto")] +impl Partitioning { + /// Serialize this partitioning into its protobuf representation. + pub fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result { + use datafusion_proto_models::protobuf; + + let partition_method = match self { + Partitioning::RoundRobinBatch(n) => { + protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) + } + Partitioning::Hash(exprs, n) => { + protobuf::partitioning::PartitionMethod::Hash( + protobuf::PhysicalHashRepartition { + hash_expr: ctx.encode_children_expressions(exprs)?, + partition_count: *n as u64, + }, + ) + } + Partitioning::Range(range) => { + let sort_expr = range + .ordering() + .iter() + .map(|sort_expr| sort_expr.try_to_proto(ctx)) + .collect::>>()?; + let split_point = range + .split_points() + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>()?; + protobuf::partitioning::PartitionMethod::Range( + protobuf::PhysicalRangePartitioning { + sort_expr, + split_point, + }, + ) + } + Partitioning::UnknownPartitioning(n) => { + protobuf::partitioning::PartitionMethod::Unknown(*n as u64) + } + }; + Ok(protobuf::Partitioning { + partition_method: Some(partition_method), + }) + } + + /// Reconstruct a [`Partitioning`] from its protobuf representation. + /// + /// Returns `Ok(None)` when the message carries no `partition_method`, which + /// the wire format uses to mean "no output partitioning declared"; callers + /// for which it is required should turn that into their own error. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::Partitioning, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::{ScalarValue, internal_datafusion_err, internal_err}; + use datafusion_proto_models::protobuf; + + let Some(partition_method) = node.partition_method.as_ref() else { + return Ok(None); + }; + let partitioning = match partition_method { + protobuf::partitioning::PartitionMethod::RoundRobin(n) => { + Partitioning::RoundRobinBatch(partition_count(*n)?) + } + protobuf::partitioning::PartitionMethod::Hash(hash) => { + let exprs = hash + .hash_expr + .iter() + .map(|expr| ctx.decode(expr)) + .collect::>>()?; + Partitioning::Hash(exprs, partition_count(hash.partition_count)?) + } + protobuf::partitioning::PartitionMethod::Unknown(n) => { + Partitioning::UnknownPartitioning(partition_count(*n)?) + } + protobuf::partitioning::PartitionMethod::Range(range) => { + let sort_exprs = range + .sort_expr + .iter() + .map(|sort_expr| PhysicalSortExpr::try_from_proto(sort_expr, ctx)) + .collect::>>()?; + let sort_expr_count = sort_exprs.len(); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!( + "Range partitioning requires non-empty ordering" + ) + })?; + if ordering.len() != sort_expr_count { + return internal_err!( + "Range partitioning ordering must not contain duplicate expressions" + ); + } + let split_points = range + .split_point + .iter() + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>()?; + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) + } + }; + Ok(Some(partitioning)) + } +} + +/// Narrow a wire partition count to `usize`. +#[cfg(feature = "proto")] +fn partition_count(count: u64) -> Result { + usize::try_from(count).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Partition count {count} exceeds usize::MAX" + ) + }) +} + impl PartialEq for Partitioning { fn eq(&self, other: &Partitioning) -> bool { match (self, other) { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs index f84cd67d46e3b..6085b867dedaa 100644 --- a/datafusion/physical-plan/src/proto.rs +++ b/datafusion/physical-plan/src/proto.rs @@ -66,6 +66,12 @@ use datafusion_execution::TaskContext; use datafusion_expr::physical_planning_context::ScalarSubqueryResults; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, +}; +use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, +}; use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use crate::ExecutionPlan; @@ -210,6 +216,25 @@ impl<'a> ExecutionPlanEncodeCtx<'a> { pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { self.encoder.encode_udwf(udwf) } + + /// An expression-level encode context backed by this plan context. + /// + /// Lets a plan hand `ctx` to expression-level conversions that own their own + /// wire logic — e.g. + /// [`Partitioning::try_to_proto`](datafusion_physical_expr::Partitioning::try_to_proto) + /// and + /// [`PhysicalSortExpr::try_to_proto`](datafusion_physical_expr::PhysicalSortExpr::try_to_proto). + pub fn expr_ctx(&self) -> PhysicalExprEncodeCtx<'_> { + PhysicalExprEncodeCtx::new(self) + } +} + +/// Lets [`ExecutionPlanEncodeCtx`] back an [`PhysicalExprEncodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprEncode for ExecutionPlanEncodeCtx<'_> { + fn encode(&self, expr: &Arc) -> Result { + self.encode_expr(expr) + } } /// Context handed to a plan's `try_from_proto` associated function. @@ -317,6 +342,28 @@ impl<'a> ExecutionPlanDecodeCtx<'a> { ) -> Result> { self.decoder.decode_udwf(name, payload) } + + /// An expression-level decode context backed by this plan context, bound to + /// `input_schema`. + /// + /// The decode counterpart of + /// [`ExecutionPlanEncodeCtx::expr_ctx`], for calling conversions such as + /// [`Partitioning::try_from_proto`](datafusion_physical_expr::Partitioning::try_from_proto). + pub fn expr_ctx<'s>(&'s self, input_schema: &'s Schema) -> PhysicalExprDecodeCtx<'s> { + PhysicalExprDecodeCtx::new(input_schema, self) + } +} + +/// Lets [`ExecutionPlanDecodeCtx`] back a [`PhysicalExprDecodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> { + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.decode_expr(node, schema) + } } /// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 3473aad9b3fc0..873f35fd6aed9 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1713,64 +1713,14 @@ impl ExecutionPlan for RepartitionExec { let input = ctx.encode_child(self.input())?; - // Keep the existing protobuf wire representation unchanged. - let partition_method = match self.partitioning() { - Partitioning::RoundRobinBatch(n) => { - protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) - } - Partitioning::Hash(exprs, n) => { - let hash_expr = ctx.encode_expressions(exprs)?; - protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr, - partition_count: *n as u64, - }, - ) - } - Partitioning::Range(range) => { - let sort_expr = range - .ordering() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; - let split_point = range - .split_points() - .iter() - .map(|split_point| { - let value = split_point - .values() - .iter() - .map(|value| value.try_into().map_err(Into::into)) - .collect::>>()?; - Ok(protobuf::PhysicalRangeSplitPoint { value }) - }) - .collect::>>()?; - protobuf::partitioning::PartitionMethod::Range( - protobuf::PhysicalRangePartitioning { - sort_expr, - split_point, - }, - ) - } - Partitioning::UnknownPartitioning(n) => { - protobuf::partitioning::PartitionMethod::Unknown(*n as u64) - } - }; + let partitioning = self.partitioning().try_to_proto(&ctx.expr_ctx())?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::Repartition(Box::new( protobuf::RepartitionExecNode { input: Some(Box::new(input)), - partitioning: Some(protobuf::Partitioning { - partition_method: Some(partition_method), - }), + partitioning: Some(partitioning), preserve_order: self.preserve_order(), }, )), @@ -1800,84 +1750,23 @@ impl RepartitionExec { )?; let input_schema = input.schema(); - let partition_method = repart + let partitioning = repart .partitioning .as_ref() - .and_then(|p| p.partition_method.as_ref()) + .map(|partitioning| { + Partitioning::try_from_proto( + partitioning, + &ctx.expr_ctx(input_schema.as_ref()), + ) + }) + .transpose()? + .flatten() .ok_or_else(|| { datafusion_common::internal_datafusion_err!( "RepartitionExec is missing required field 'partitioning'" ) })?; - let partitioning = match partition_method { - protobuf::partitioning::PartitionMethod::RoundRobin(n) => { - Partitioning::RoundRobinBatch(*n as usize) - } - protobuf::partitioning::PartitionMethod::Hash(hash) => { - let exprs = hash - .hash_expr - .iter() - .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) - .collect::>>()?; - let partition_count = - usize::try_from(hash.partition_count).map_err(|_| { - datafusion_common::internal_datafusion_err!( - "Hash partition count {} exceeds usize::MAX", - hash.partition_count - ) - })?; - Partitioning::Hash(exprs, partition_count) - } - protobuf::partitioning::PartitionMethod::Unknown(n) => { - Partitioning::UnknownPartitioning(*n as usize) - } - protobuf::partitioning::PartitionMethod::Range(range) => { - let sort_exprs = range - .sort_expr - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_ref().ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "Unexpected empty physical expression" - ) - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, input_schema.as_ref())?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect::>>()?; - let sort_expr_count = sort_exprs.len(); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "Range partitioning requires non-empty ordering" - ) - })?; - if ordering.len() != sort_expr_count { - return datafusion_common::internal_err!( - "Range partitioning ordering must not contain duplicate expressions" - ); - } - let split_points = range - .split_point - .iter() - .map(|split_point| { - let values = split_point - .value - .iter() - .map(|value| ScalarValue::try_from(value).map_err(Into::into)) - .collect::>>()?; - Ok(SplitPoint::new(values)) - }) - .collect::>>()?; - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) - } - }; - let mut repart_exec = RepartitionExec::try_new(input, partitioning)?; if repart.preserve_order { repart_exec = repart_exec.with_preserve_order(); diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 34ad8c7a62fc7..1082571fd8fb1 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -23,9 +23,7 @@ use arrow::array::RecordBatch; use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; -use datafusion_common::{ - DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, -}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -51,9 +49,7 @@ use datafusion_physical_plan::expressions::{ }; use datafusion_physical_plan::joins::HashExpr; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; -use datafusion_physical_plan::{ - Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use datafusion_proto_common::common::proto_error; use super::{ @@ -423,83 +419,20 @@ pub fn parse_protobuf_partitioning( input_schema: &Schema, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - match partitioning { - Some(protobuf::Partitioning { partition_method }) => match partition_method { - Some(protobuf::partitioning::PartitionMethod::RoundRobin( - partition_count, - )) => Ok(Some(Partitioning::RoundRobinBatch( - *partition_count as usize, - ))), - Some(protobuf::partitioning::PartitionMethod::Hash(hash_repartition)) => { - parse_protobuf_hash_partitioning( - Some(hash_repartition), - ctx, - input_schema, - proto_converter, - ) - } - Some(protobuf::partitioning::PartitionMethod::Range(range_partitioning)) => { - Ok(Some(parse_protobuf_range_partitioning( - range_partitioning, - ctx, - input_schema, - proto_converter, - )?)) - } - Some(protobuf::partitioning::PartitionMethod::Unknown(partition_count)) => { - Ok(Some(Partitioning::UnknownPartitioning( - *partition_count as usize, - ))) - } - None => Ok(None), - }, - None => Ok(None), - } -} - -fn parse_protobuf_range_partitioning( - range_partitioning: &protobuf::PhysicalRangePartitioning, - ctx: &PhysicalPlanDecodeContext<'_>, - input_schema: &Schema, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - let sort_exprs = parse_physical_sort_exprs( - &range_partitioning.sort_expr, + let decoder = ConverterDecoder { ctx, - input_schema, proto_converter, - )?; - let sort_expr_count = sort_exprs.len(); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - internal_datafusion_err!("Range partitioning requires non-empty ordering") - })?; - if ordering.len() != sort_expr_count { - return Err(internal_datafusion_err!( - "Range partitioning ordering must not contain duplicate expressions" - )); - } - let split_points = range_partitioning - .split_point - .iter() - .map(parse_protobuf_range_split_point) - .collect::>()?; - Ok(Partitioning::Range(RangePartitioning::try_new( - ordering, - split_points, - )?)) -} - -fn parse_protobuf_range_split_point( - split_point: &protobuf::PhysicalRangeSplitPoint, -) -> Result { - let values = split_point - .value - .iter() - .map(|value| ScalarValue::try_from(value).map_err(Into::into)) - .collect::>()?; - Ok(SplitPoint::new(values)) + }; + let decode_ctx = + datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::new( + input_schema, + &decoder, + ); + partitioning + .map(|partitioning| Partitioning::try_from_proto(partitioning, &decode_ctx)) + .transpose() + .map(Option::flatten) } - pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result> { @@ -776,6 +709,7 @@ mod tests { use super::*; use arrow::datatypes::{DataType, Field, Schema}; use chrono::{TimeZone, Utc}; + use datafusion_common::ScalarValue; use object_store::ObjectMeta; use object_store::path::Path; diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 5189972f0e200..bab7af2ab48f5 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -36,9 +36,7 @@ use datafusion_physical_expr::{HigherOrderFunctionExpr, ScalarFunctionExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; -use datafusion_physical_plan::{ - Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use super::{ DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -358,70 +356,13 @@ pub fn serialize_partitioning( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let serialized_partitioning = match partitioning { - Partitioning::RoundRobinBatch(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::RoundRobin( - *partition_count as u64, - )), - }, - Partitioning::Hash(exprs, partition_count) => { - let serialized_exprs = - serialize_physical_exprs(exprs, codec, proto_converter)?; - protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr: serialized_exprs, - partition_count: *partition_count as u64, - }, - )), - } - } - Partitioning::Range(range) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Range( - serialize_range_partitioning(range, codec, proto_converter)?, - )), - }, - Partitioning::UnknownPartitioning(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Unknown( - *partition_count as u64, - )), - }, + let encoder = ConverterEncoder { + codec, + proto_converter, }; - Ok(serialized_partitioning) -} - -fn serialize_range_partitioning( - range: &RangePartitioning, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - Ok(protobuf::PhysicalRangePartitioning { - sort_expr: serialize_physical_sort_exprs( - range.ordering().iter().cloned(), - codec, - proto_converter, - )?, - split_point: range - .split_points() - .iter() - .map(serialize_range_split_point) - .collect::>()?, - }) -} - -fn serialize_range_split_point( - split_point: &SplitPoint, -) -> Result { - Ok(protobuf::PhysicalRangeSplitPoint { - value: split_point - .values() - .iter() - .map(|value| { - TryInto::::try_into(value) - .map_err(Into::into) - }) - .collect::>()?, - }) + partitioning.try_to_proto( + &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx::new(&encoder), + ) } /// Thin shim over `TryFrom<&PartitionedFile>`, which owns the wire logic. From 38b5b2a81ba21f88c848ccb5b7939bc9c6f2ac29 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:02:33 -0500 Subject: [PATCH 2/4] refactor(proto): put ordering serde on the sort-expression types `PhysicalSortExpr::try_to_proto` / `try_from_proto` convert one sort expression, but every caller holds a sequence of them, so each one hand-rolls the same map/collect over the flat `PhysicalSortExprNode` list: `AggregateExec`'s ordering requirement, `SymmetricHashJoinExec`'s left/right sort expressions, the window expressions, range partitioning, and (in flight) `FileScanConfig`'s output orderings and the file sinks' required ordering. Add `sort_exprs_try_to_proto` / `sort_exprs_try_from_proto` next to `PhysicalSortExpr` and route those callers through them. The encode side takes anything that borrows a `PhysicalSortExpr`, so a `LexOrdering`, a `&[PhysicalSortExpr]` and a `LexRequirement` mapped through `PhysicalSortExpr::from` all work. The decode side returns the expressions rather than a `LexOrdering`, because callers disagree on what an empty list means: "no ordering declared" for a scan, an error for an operator that requires one. Decode now reports the missing inner expression uniformly as "PhysicalSortExpr is missing required field 'expr'", replacing three bespoke messages. No wire-format change. --- .../physical-expr-common/src/sort_expr.rs | 53 ++++++++ datafusion/physical-expr/src/partitioning.rs | 126 ++++++++++++++++-- .../physical-plan/src/aggregates/mod.rs | 46 ++----- .../src/joins/symmetric_hash_join.rs | 66 +++------ datafusion/physical-plan/src/windows/proto.rs | 37 +---- 5 files changed, 207 insertions(+), 121 deletions(-) diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 82f7a6e5fc33f..72e877234752f 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -226,6 +226,59 @@ impl PhysicalSortExpr { } } +/// Serialize a sequence of sort expressions into the flat +/// [`PhysicalSortExprNode`] list the wire format uses for an ordering. +/// +/// Accepts anything that yields [`PhysicalSortExpr`]s by value or by reference, +/// so a [`LexOrdering`], a `&[PhysicalSortExpr]`, or a [`LexRequirement`] +/// mapped through [`PhysicalSortExpr::from`] all work: +/// +/// ```ignore +/// let nodes = sort_exprs_try_to_proto(ordering.iter(), ctx)?; +/// let nodes = sort_exprs_try_to_proto( +/// requirement.iter().map(|req| PhysicalSortExpr::from(req.clone())), +/// ctx, +/// )?; +/// ``` +/// +/// The `PhysicalSortExprNodeCollection` message some plans use is just this +/// list in a wrapper, so those callers wrap the result themselves rather than +/// this function guessing which shape they mean. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +pub fn sort_exprs_try_to_proto>( + exprs: impl IntoIterator, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, +) -> Result> { + exprs + .into_iter() + .map(|expr| expr.borrow().try_to_proto(ctx)) + .collect() +} + +/// Reconstruct a sequence of sort expressions from the flat +/// [`PhysicalSortExprNode`] list, the counterpart of +/// [`sort_exprs_try_to_proto`]. +/// +/// Returns the expressions rather than a [`LexOrdering`] or a +/// [`LexRequirement`], because callers differ in what an empty list means: +/// `LexOrdering::new` / `LexRequirement::new` return `None` for it, which is +/// "no ordering declared" for a scan and an error for an operator that requires +/// one. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +pub fn sort_exprs_try_from_proto( + nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode], + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, +) -> Result> { + nodes + .iter() + .map(|node| PhysicalSortExpr::try_from_proto(node, ctx)) + .collect() +} + impl PartialEq for PhysicalSortExpr { fn eq(&self, other: &Self) -> bool { self.options == other.options && self.expr.eq(&other.expr) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 9a24658411300..ab63f52d836a3 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -25,6 +25,10 @@ pub use datafusion_common::SplitPoint; use datafusion_common::{Result, validate_range_split_points}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; use std::fmt; use std::fmt::Display; use std::sync::Arc; @@ -546,11 +550,7 @@ impl Partitioning { ) } Partitioning::Range(range) => { - let sort_expr = range - .ordering() - .iter() - .map(|sort_expr| sort_expr.try_to_proto(ctx)) - .collect::>>()?; + let sort_expr = sort_exprs_try_to_proto(range.ordering().iter(), ctx)?; let split_point = range .split_points() .iter() @@ -610,11 +610,7 @@ impl Partitioning { Partitioning::UnknownPartitioning(partition_count(*n)?) } protobuf::partitioning::PartitionMethod::Range(range) => { - let sort_exprs = range - .sort_expr - .iter() - .map(|sort_expr| PhysicalSortExpr::try_from_proto(sort_expr, ctx)) - .collect::>>()?; + let sort_exprs = sort_exprs_try_from_proto(&range.sort_expr, ctx)?; let sort_expr_count = sort_exprs.len(); let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { internal_datafusion_err!( @@ -1278,3 +1274,113 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod ordering_proto_tests { + use std::sync::Arc; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_physical_expr_common::sort_expr::{ + LexRequirement, PhysicalSortExpr, PhysicalSortRequirement, + sort_exprs_try_from_proto, sort_exprs_try_to_proto, + }; + + use crate::expressions::Column; + use crate::proto_test_util::{StubDecoder, StubEncoder}; + + fn schema() -> Schema { + Schema::new(vec![Field::new("a", DataType::Int32, false)]) + } + + fn sort_expr(descending: bool, nulls_first: bool) -> PhysicalSortExpr { + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions { + descending, + nulls_first, + }, + ) + } + + #[test] + fn sort_exprs_round_trip_preserves_options_and_order() { + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let exprs = vec![sort_expr(true, false), sort_expr(false, true)]; + + let nodes = sort_exprs_try_to_proto(&exprs, &encode_ctx).unwrap(); + // `asc` is the inverse of `descending` on the wire. + assert_eq!( + nodes + .iter() + .map(|node| (node.asc, node.nulls_first)) + .collect::>(), + vec![(false, false), (true, true)] + ); + + let schema = schema(); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let decoded = sort_exprs_try_from_proto(&nodes, &decode_ctx).unwrap(); + assert_eq!( + decoded.iter().map(|expr| expr.options).collect::>(), + exprs.iter().map(|expr| expr.options).collect::>() + ); + } + + #[test] + fn sort_exprs_accepts_owned_requirements() { + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let requirement = LexRequirement::from([PhysicalSortRequirement::new( + Arc::new(Column::new("a", 0)), + Some(SortOptions { + descending: true, + nulls_first: true, + }), + )]); + + let nodes = sort_exprs_try_to_proto( + requirement + .iter() + .map(|req| PhysicalSortExpr::from(req.clone())), + &encode_ctx, + ) + .unwrap(); + + assert_eq!(nodes.len(), 1); + assert!(!nodes[0].asc); + assert!(nodes[0].nulls_first); + } + + #[test] + fn sort_exprs_propagate_encode_errors() { + let encoder = StubEncoder::failing_on(2); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let exprs = vec![sort_expr(false, false), sort_expr(true, true)]; + + let err = sort_exprs_try_to_proto(&exprs, &encode_ctx).unwrap_err(); + assert!(err.to_string().contains("stub encode failure on call 2")); + } + + #[test] + fn sort_exprs_reject_missing_inner_expr() { + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let mut nodes = + sort_exprs_try_to_proto(&[sort_expr(false, false)], &encode_ctx).unwrap(); + nodes[0].expr = None; + + let schema = schema(); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = sort_exprs_try_from_proto(&nodes, &decode_ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalSortExpr is missing required field 'expr'") + ); + } +} diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fbc..0523628ad7e6a 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -2264,17 +2264,11 @@ fn encode_aggregate_expr( let expressions = aggr_expr.expressions(); let expr = ctx.encode_expressions(expressions.iter())?; - let ordering_req = aggr_expr - .order_bys() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; + let ordering_req = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( + aggr_expr.order_bys(), + &ctx.expr_ctx(), + )?; let name = aggr_expr.fun().name().to_string(); // The context already applies `(!buf.is_empty()).then_some(buf)`. let fun_definition = ctx.encode_udaf(aggr_expr.fun())?; @@ -2314,7 +2308,6 @@ impl AggregateExec { node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_proto_models::protobuf; use protobuf::physical_aggregate_expr_node::AggregateFunction; @@ -2421,24 +2414,11 @@ impl AggregateExec { .iter() .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) .collect::>>()?; - let order_by = aggregate - .ordering_req - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_deref().ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "AggregateExec ordering expression is missing its inner expr" - ) - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, input_schema.as_ref())?, - options: arrow::compute::SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect::>>()?; + let order_by = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( + &aggregate.ordering_req, + &ctx.expr_ctx(input_schema.as_ref()), + )?; let Some(AggregateFunction::UserDefinedAggrFunction(udaf_name)) = aggregate.aggregate_function.as_ref() else { @@ -2448,10 +2428,8 @@ impl AggregateExec { }; // The context owns the payload-to-codec and // registry-to-codec fallback order. - let udaf = ctx.decode_udaf( - udaf_name, - aggregate.fun_definition.as_deref(), - )?; + let udaf = + ctx.decode_udaf(udaf_name, aggregate.fun_definition.as_deref())?; let (human_display, human_display_alias) = split_human_display_alias(&aggregate.human_display, name); let builder = AggregateExprBuilder::new(udaf, args) diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 95f4c35871431..eb358b10b4bfd 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -698,23 +698,18 @@ impl ExecutionPlan for SymmetricHashJoinExec { }) }) .transpose()?; + let expr_ctx = ctx.expr_ctx(); let encode_sort_exprs = |exprs: Option<&LexOrdering>| -> Result> { - exprs - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }) - }) - .collect::>>() - }) - .transpose() - .map(Option::unwrap_or_default) + exprs.map_or_else( + || Ok(vec![]), + |exprs| { + datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( + exprs.iter(), + &expr_ctx, + ) + }, + ) }; let left_sort_exprs = encode_sort_exprs(self.left_sort_exprs())?; let right_sort_exprs = encode_sort_exprs(self.right_sort_exprs())?; @@ -751,7 +746,6 @@ impl SymmetricHashJoinExec { ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, ) -> Result> { use datafusion_common::internal_datafusion_err; - use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_proto_models::protobuf; let sym_join = crate::expect_plan_variant!( @@ -883,39 +877,19 @@ impl SymmetricHashJoinExec { }) .transpose()?; let decode_sort_exprs = |sort_exprs: &[protobuf::PhysicalSortExprNode], - schema: &Schema, - field: &str| + schema: &Schema| -> Result> { - let sort_exprs = sort_exprs - .iter() - .map(|sort_expr| { - let expr = ctx.decode_required_expr( - sort_expr.expr.as_deref(), - schema, - "SymmetricHashJoinExec", - field, - )?; - Ok(PhysicalSortExpr { - expr, - options: arrow::compute::SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect::>>()?; + let sort_exprs = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( + sort_exprs, + &ctx.expr_ctx(schema), + )?; Ok(LexOrdering::new(sort_exprs)) }; - let left_sort_exprs = decode_sort_exprs( - &sym_join.left_sort_exprs, - left_schema.as_ref(), - "left_sort_exprs", - )?; - let right_sort_exprs = decode_sort_exprs( - &sym_join.right_sort_exprs, - right_schema.as_ref(), - "right_sort_exprs", - )?; + let left_sort_exprs = + decode_sort_exprs(&sym_join.left_sort_exprs, left_schema.as_ref())?; + let right_sort_exprs = + decode_sort_exprs(&sym_join.right_sort_exprs, right_schema.as_ref())?; Self::try_new( left, diff --git a/datafusion/physical-plan/src/windows/proto.rs b/datafusion/physical-plan/src/windows/proto.rs index aa62158d18fa0..e96b0a9fb1087 100644 --- a/datafusion/physical-plan/src/windows/proto.rs +++ b/datafusion/physical-plan/src/windows/proto.rs @@ -19,7 +19,6 @@ use std::sync::Arc; -use arrow::compute::SortOptions; use arrow::datatypes::Schema; use datafusion_common::{ Result, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, @@ -28,7 +27,9 @@ use datafusion_expr::{ WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, }; use datafusion_physical_expr::window::SlidingAggregateWindowExpr; -use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; use datafusion_proto_common::protobuf_common; use datafusion_proto_models::protobuf::{self, physical_window_expr_node}; @@ -93,17 +94,7 @@ pub(super) fn encode_physical_window_expr( let args = ctx.encode_expressions(&args)?; let partition_by = ctx.encode_expressions(window_expr.partition_by())?; - let order_by = window_expr - .order_by() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; + let order_by = sort_exprs_try_to_proto(window_expr.order_by(), &ctx.expr_ctx())?; Ok(protobuf::PhysicalWindowExprNode { args, @@ -133,24 +124,8 @@ pub(super) fn decode_physical_window_expr( .iter() .map(|expr| ctx.decode_expr(expr, input_schema)) .collect::>>()?; - let order_by = proto - .order_by - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!( - "Missing expr in window order_by sort expression" - ) - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, input_schema)?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect::>>()?; + let order_by = + sort_exprs_try_from_proto(&proto.order_by, &ctx.expr_ctx(input_schema))?; let window_frame = proto .window_frame .as_ref() From 4a85a4d1a4800bfa8c7bff3000b8eb454d08d66b Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 13:11:17 -0500 Subject: [PATCH 3/4] refactor(proto): address review feedback on partitioning serde - `parse_protobuf_hash_partitioning` delegates to the shared `Partitioning::try_from_proto` instead of keeping its own copy of the hash wire format, so its `try_into().unwrap()` on the partition count can no longer panic for direct callers on a 32-bit target. Covered by a new unit test, since the function has no in-tree callers left. - Encode validates the partition count with `u64::try_from` rather than `as u64`, mirroring the decode side instead of truncating silently. - Reword the `Partitioning` proto docs: `FileScanConfig` is an intended consumer, not a current one. - Grammar: "a `PhysicalExprEncodeCtx`". Co-Authored-By: Claude Opus 5 --- datafusion/physical-expr/src/partitioning.rs | 30 +++++++++--- datafusion/physical-plan/src/proto.rs | 2 +- .../proto/src/physical_plan/from_proto.rs | 26 ++++------- .../tests/cases/roundtrip_physical_plan.rs | 46 +++++++++++++++++++ 4 files changed, 81 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index ab63f52d836a3..124a4e4439f94 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -523,9 +523,10 @@ impl Partitioning { /// /// Child expressions (hash keys, range orderings) and `ScalarValue` split /// points are (de)serialized through the expression-level context, so this is -/// the single copy of the partitioning wire format: `RepartitionExec`, -/// `FileScanConfig` and `datafusion-proto`'s central serializer all route -/// through it. +/// the single copy of the partitioning wire format: `RepartitionExec` and +/// `datafusion-proto`'s central serializer route through it, and the remaining +/// per-plan migrations (`FileScanConfig` and friends) are meant to do the same +/// rather than grow another copy. /// /// [`protobuf::Partitioning`]: datafusion_proto_models::protobuf::Partitioning #[cfg(feature = "proto")] @@ -539,13 +540,15 @@ impl Partitioning { let partition_method = match self { Partitioning::RoundRobinBatch(n) => { - protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) + protobuf::partitioning::PartitionMethod::RoundRobin(wire_partition_count( + *n, + )?) } Partitioning::Hash(exprs, n) => { protobuf::partitioning::PartitionMethod::Hash( protobuf::PhysicalHashRepartition { hash_expr: ctx.encode_children_expressions(exprs)?, - partition_count: *n as u64, + partition_count: wire_partition_count(*n)?, }, ) } @@ -571,7 +574,9 @@ impl Partitioning { ) } Partitioning::UnknownPartitioning(n) => { - protobuf::partitioning::PartitionMethod::Unknown(*n as u64) + protobuf::partitioning::PartitionMethod::Unknown(wire_partition_count( + *n, + )?) } }; Ok(protobuf::Partitioning { @@ -651,6 +656,19 @@ fn partition_count(count: u64) -> Result { }) } +/// Widen a partition count to its `u64` wire representation. +/// +/// The mirror of [`partition_count`]: an out-of-range count is an error on both +/// sides rather than a silent truncation on the way out. +#[cfg(feature = "proto")] +fn wire_partition_count(count: usize) -> Result { + u64::try_from(count).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Partition count {count} exceeds u64::MAX" + ) + }) +} + impl PartialEq for Partitioning { fn eq(&self, other: &Partitioning) -> bool { match (self, other) { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs index 6085b867dedaa..7640d76c3e010 100644 --- a/datafusion/physical-plan/src/proto.rs +++ b/datafusion/physical-plan/src/proto.rs @@ -229,7 +229,7 @@ impl<'a> ExecutionPlanEncodeCtx<'a> { } } -/// Lets [`ExecutionPlanEncodeCtx`] back an [`PhysicalExprEncodeCtx`], so +/// Lets [`ExecutionPlanEncodeCtx`] back a [`PhysicalExprEncodeCtx`], so /// expression-level conversions can be reused from plan hooks. impl PhysicalExprEncode for ExecutionPlanEncodeCtx<'_> { fn encode(&self, expr: &Arc) -> Result { diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 1082571fd8fb1..60647bd7aa840 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -395,22 +395,16 @@ pub fn parse_protobuf_hash_partitioning( input_schema: &Schema, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - match partitioning { - Some(hash_part) => { - let expr = parse_physical_exprs( - &hash_part.hash_expr, - ctx, - input_schema, - proto_converter, - )?; - - Ok(Some(Partitioning::Hash( - expr, - hash_part.partition_count.try_into().unwrap(), - ))) - } - None => Ok(None), - } + // Delegate to the shared decoder rather than keep a second copy of the hash + // wire format: a partition count that does not fit in `usize` (a 32-bit + // target reading a plan written on a 64-bit one) is then an error here too + // instead of a panic. + let hash = partitioning.map(|hash_part| protobuf::Partitioning { + partition_method: Some(protobuf::partitioning::PartitionMethod::Hash( + hash_part.clone(), + )), + }); + parse_protobuf_partitioning(hash.as_ref(), ctx, input_schema, proto_converter) } pub fn parse_protobuf_partitioning( diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 864e6d68676ee..e0a5a9e2b6889 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2374,6 +2374,52 @@ fn roundtrip_range_partitioning() -> Result<()> { roundtrip_test(Arc::new(repartition)) } +/// `parse_protobuf_hash_partitioning` has no in-tree callers left; it delegates +/// to the shared `Partitioning::try_from_proto`, so pin that it still decodes +/// the hash message it is handed. +#[test] +fn parse_hash_partitioning_delegates_to_shared_decoder() -> Result<()> { + use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning; + + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_ctx = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let proto_converter = DefaultPhysicalProtoConverter {}; + + let hash_expr = serialize_physical_expr_with_converter( + &col("a", &schema)?, + &codec, + &proto_converter, + )?; + let hash = protobuf::PhysicalHashRepartition { + hash_expr: vec![hash_expr], + partition_count: 4, + }; + + let partitioning = parse_protobuf_hash_partitioning( + Some(&hash), + &decode_ctx, + &schema, + &proto_converter, + )?; + let Some(Partitioning::Hash(exprs, count)) = partitioning else { + panic!("expected hash partitioning, got {partitioning:?}"); + }; + assert_eq!(count, 4); + assert_eq!(exprs.len(), 1); + assert_eq!(exprs[0].to_string(), col("a", &schema)?.to_string()); + + // No message means no partitioning, as before. + assert!( + parse_protobuf_hash_partitioning(None, &decode_ctx, &schema, &proto_converter)? + .is_none() + ); + + Ok(()) +} + #[test] fn roundtrip_interleave() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); From eeb9ed302d95c2f9b2345c20eae49502ef460a9e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 04:06:45 +0000 Subject: [PATCH 4/4] test(proto): cover out-of-range partition counts The partition count is a `u64` on the wire and a `usize` in memory, and routing both directions through the shared helpers was part of what motivated moving the serde onto `Partitioning`. Nothing exercised the narrowing, though: the decoder test only used a count of 4. Cover both crossings. `partition_count` / `wire_partition_count` are pinned at the `usize` ceiling, and each counted partition method (`RoundRobin`, `Unknown`, `Hash`) is driven through `try_to_proto` and `try_from_proto` with the widest count the wire can carry. The `parse_protobuf_hash_partitioning` test grows the same case. The reject-instead-of-panic arm is only reachable on a target narrower than 64 bits, so it is asserted under `cfg(target_pointer_width)`; on a 64-bit target every `u64` fits and the assertion is that the largest possible count decodes whole rather than being truncated. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01QTJQMADVSt9yvDf2LPLk6S --- datafusion/physical-expr/src/partitioning.rs | 126 ++++++++++++++++++ .../tests/cases/roundtrip_physical_plan.rs | 37 +++++ 2 files changed, 163 insertions(+) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 124a4e4439f94..98f082f7256db 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -1402,3 +1402,129 @@ mod ordering_proto_tests { ); } } + +/// Partition counts are `usize` in memory and `u64` on the wire, so every +/// counted [`Partitioning`] variant crosses a width boundary in both +/// directions. These pin that neither crossing wraps or panics. +#[cfg(all(test, feature = "proto"))] +mod partition_count_proto_tests { + use std::sync::Arc; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf; + + use super::{Partitioning, partition_count, wire_partition_count}; + use crate::expressions::Column; + use crate::proto_test_util::{StubDecoder, StubEncoder, column_node}; + + fn partitioning_node( + method: protobuf::partitioning::PartitionMethod, + ) -> protobuf::Partitioning { + protobuf::Partitioning { + partition_method: Some(method), + } + } + + /// The counted variants, each carrying `count`. `Range` is excluded: it + /// derives its partition count from its split points rather than reading + /// one off the wire. + fn counted_methods(count: u64) -> Vec { + use protobuf::partitioning::PartitionMethod; + + vec![ + PartitionMethod::RoundRobin(count), + PartitionMethod::Unknown(count), + PartitionMethod::Hash(protobuf::PhysicalHashRepartition { + hash_expr: vec![column_node("a")], + partition_count: count, + }), + ] + } + + #[test] + fn partition_count_round_trips_at_the_usize_ceiling() { + // `usize::MAX` is the largest count that can exist in memory, so it has + // to widen onto the wire and narrow back unchanged. + let wire = wire_partition_count(usize::MAX).unwrap(); + assert_eq!(wire, u64::try_from(usize::MAX).unwrap()); + assert_eq!(partition_count(wire).unwrap(), usize::MAX); + } + + #[test] + fn out_of_range_partition_count_is_reported_not_wrapped() { + // A count wider than the target's `usize` can only be reached by + // decoding on a narrower host than the one that encoded. That used to + // wrap (`as usize`) or panic (`unwrap`); it is an error now. On a + // 64-bit target every `u64` fits, so the same input has to decode + // losslessly instead of being rejected. + let narrowed = partition_count(u64::MAX); + + #[cfg(target_pointer_width = "64")] + assert_eq!(narrowed.unwrap(), usize::MAX); + + #[cfg(not(target_pointer_width = "64"))] + assert!( + narrowed + .unwrap_err() + .to_string() + .contains("Partition count 18446744073709551615 exceeds usize::MAX") + ); + } + + #[test] + fn try_from_proto_narrows_every_counted_variant() { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + for method in counted_methods(u64::MAX) { + let decoded = + Partitioning::try_from_proto(&partitioning_node(method), &decode_ctx); + + #[cfg(target_pointer_width = "64")] + assert_eq!(decoded.unwrap().unwrap().partition_count(), usize::MAX); + + #[cfg(not(target_pointer_width = "64"))] + assert!( + decoded + .unwrap_err() + .to_string() + .contains("exceeds usize::MAX") + ); + } + } + + #[test] + fn try_to_proto_widens_every_counted_variant() { + use protobuf::partitioning::PartitionMethod; + + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let hash_key: Arc = Arc::new(Column::new("a", 0)); + + let encoded = [ + Partitioning::RoundRobinBatch(usize::MAX), + Partitioning::UnknownPartitioning(usize::MAX), + Partitioning::Hash(vec![hash_key], usize::MAX), + ] + .iter() + .map(|partitioning| { + match partitioning + .try_to_proto(&encode_ctx) + .unwrap() + .partition_method + { + Some(PartitionMethod::RoundRobin(n) | PartitionMethod::Unknown(n)) => n, + Some(PartitionMethod::Hash(hash)) => hash.partition_count, + other => panic!("expected a counted partition method, got {other:?}"), + } + }) + .collect::>(); + + // Every variant widens to the same wire value, with no truncation. + assert_eq!(encoded, vec![u64::try_from(usize::MAX).unwrap(); 3]); + } +} diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index e0a5a9e2b6889..508ba5c020d4d 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2417,6 +2417,43 @@ fn parse_hash_partitioning_delegates_to_shared_decoder() -> Result<()> { .is_none() ); + // The count is a `u64` on the wire and a `usize` in memory, so decoding + // narrows it. A count that does not fit is the case that motivated routing + // this through the shared decoder: it used to `unwrap()` and panic, and now + // reports an error. Only a target narrower than 64 bits can reach that arm + // -- on a 64-bit target every `u64` fits, and the assertion there is that + // the largest possible count survives whole rather than being truncated. + let oversized = protobuf::PhysicalHashRepartition { + hash_expr: vec![serialize_physical_expr_with_converter( + &col("a", &schema)?, + &codec, + &proto_converter, + )?], + partition_count: u64::MAX, + }; + let decoded = parse_protobuf_hash_partitioning( + Some(&oversized), + &decode_ctx, + &schema, + &proto_converter, + ); + + #[cfg(target_pointer_width = "64")] + { + let Some(Partitioning::Hash(_, count)) = decoded? else { + panic!("expected hash partitioning"); + }; + assert_eq!(count, usize::MAX); + } + + #[cfg(not(target_pointer_width = "64"))] + assert!( + decoded + .unwrap_err() + .to_string() + .contains("Partition count 18446744073709551615 exceeds usize::MAX") + ); + Ok(()) }