diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 12229c26b7d98..3473aad9b3fc0 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1703,6 +1703,187 @@ impl ExecutionPlan for RepartitionExec { cache: new_properties.into(), }))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + 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) + } + }; + + 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), + }), + preserve_order: self.preserve_order(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl RepartitionExec { + /// Reconstruct a [`RepartitionExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let repart = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Repartition, + "RepartitionExec", + ); + let input = ctx.decode_required_child( + repart.input.as_deref(), + "RepartitionExec", + "input", + )?; + let input_schema = input.schema(); + + let partition_method = repart + .partitioning + .as_ref() + .and_then(|p| p.partition_method.as_ref()) + .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(); + } + Ok(Arc::new(repart_exec)) + } } impl RepartitionExec { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 7a822324789f8..3e8e0738cd3ad 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -105,8 +105,6 @@ use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, PhysicalExpr, Wind use prost::Message; use prost::bytes::BufMut; -use self::from_proto::parse_protobuf_partitioning; -use self::to_proto::serialize_partitioning; use crate::common::{byte_to_string, str_to_byte}; use crate::convert::{FromProto, TryFromProto}; use crate::convert_required; @@ -786,8 +784,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::Merge(_) => { CoalescePartitionsExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Repartition(repart) => { - self.try_into_repartition_physical_plan(repart, ctx, proto_converter) + PhysicalPlanType::Repartition(_) => { + RepartitionExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::GlobalLimit(limit) => { self.try_into_global_limit_physical_plan(limit, ctx, proto_converter) @@ -982,14 +980,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_repartition_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(union) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_union_exec( union, @@ -1448,25 +1438,27 @@ pub trait PhysicalPlanNodeExt: Sized { CoalescePartitionsExec::try_from_proto(&node, &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `RepartitionExec` deserializes itself via `RepartitionExec::try_from_proto`" + )] fn try_into_repartition_physical_plan( &self, repart: &protobuf::RepartitionExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input: Arc = - into_physical_plan(&repart.input, ctx, proto_converter)?; - let partitioning = parse_protobuf_partitioning( - repart.partitioning.as_ref(), + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new( + repart.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { ctx, - input.schema().as_ref(), proto_converter, - )?; - let mut repart_exec = RepartitionExec::try_new(input, partitioning.unwrap())?; - if repart.preserve_order { - repart_exec = repart_exec.with_preserve_order(); - } - Ok(Arc::new(repart_exec)) + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + RepartitionExec::try_from_proto(&node, &decode_ctx) } fn try_into_global_limit_physical_plan( @@ -3370,28 +3362,22 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `RepartitionExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_repartition_exec( exec: &RepartitionExec, codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), + let encoder = ConverterPlanEncoder { codec, proto_converter, - )?; - - let pb_partitioning = - serialize_partitioning(exec.partitioning(), codec, proto_converter)?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new( - protobuf::RepartitionExecNode { - input: Some(Box::new(input)), - partitioning: Some(pb_partitioning), - preserve_order: exec.preserve_order(), - }, - ))), + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("RepartitionExec is not serializable") }) }