diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 84ffb92eaa600..72e877234752f 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -183,6 +183,102 @@ 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, + }, + }) + } +} + +/// 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 59d36c4efc1bb..98f082f7256db 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; @@ -515,6 +519,156 @@ 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` 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")] +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(wire_partition_count( + *n, + )?) + } + Partitioning::Hash(exprs, n) => { + protobuf::partitioning::PartitionMethod::Hash( + protobuf::PhysicalHashRepartition { + hash_expr: ctx.encode_children_expressions(exprs)?, + partition_count: wire_partition_count(*n)?, + }, + ) + } + Partitioning::Range(range) => { + let sort_expr = sort_exprs_try_to_proto(range.ordering().iter(), ctx)?; + 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(wire_partition_count( + *n, + )?) + } + }; + 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 = 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!( + "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" + ) + }) +} + +/// 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) { @@ -1138,3 +1292,239 @@ 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'") + ); + } +} + +/// 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/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/proto.rs b/datafusion/physical-plan/src/proto.rs index f84cd67d46e3b..7640d76c3e010 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 a [`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/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() diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 34ad8c7a62fc7..60647bd7aa840 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::{ @@ -399,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( @@ -423,83 +413,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 +703,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. diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 864e6d68676ee..508ba5c020d4d 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2374,6 +2374,89 @@ 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() + ); + + // 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(()) +} + #[test] fn roundtrip_interleave() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false);