You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/23 19:55:24 UTC
(datafusion) branch main updated: refactor(proto): migrate RepartitionExec serde (#23792)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 13dadad9ca refactor(proto): migrate RepartitionExec serde (#23792)
13dadad9ca is described below
commit 13dadad9cab1f9a259e3833c91fc413ef4b75885
Author: Burak Şen <bu...@gmail.com>
AuthorDate: Thu Jul 23 22:40:16 2026 +0300
refactor(proto): migrate RepartitionExec serde (#23792)
## Which issue does this PR close?
- Closes #23504.
## Rationale for this change
Part of epic #23494. Moves `RepartitionExec` protobuf serialization from
central dispatch.
## What changes are included in this PR?
Add new proto functions in physical plan and deprecate proto ones.
## Are these changes tested?
Yes, existing round trip tests pass.
## Are there any user-facing changes?
Existing methods are deprecated with no immediate API change.
---
datafusion/physical-plan/src/repartition/mod.rs | 181 ++++++++++++++++++++++++
datafusion/proto/src/physical_plan/mod.rs | 62 ++++----
2 files changed, 205 insertions(+), 38 deletions(-)
diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
index 12229c26b7..3473aad9b3 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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ 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::<Result<Vec<_>>>()?;
+ 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::<Result<Vec<_>>>()?;
+ Ok(protobuf::PhysicalRangeSplitPoint { value })
+ })
+ .collect::<Result<Vec<_>>>()?;
+ 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<Arc<dyn ExecutionPlan>> {
+ 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::<Result<Vec<_>>>()?;
+ 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::<Result<Vec<_>>>()?;
+ 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::<Result<Vec<_>>>()?;
+ Ok(SplitPoint::new(values))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ 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 f3f1d88e08..9c4fd8e347 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(_) => {
GlobalLimitExec::try_from_proto(self.node(), &decode_ctx)
@@ -966,14 +964,6 @@ pub trait PhysicalPlanNodeExt: Sized {
return Ok(node);
}
- if let Some(exec) = plan.downcast_ref::<RepartitionExec>() {
- return protobuf::PhysicalPlanNode::try_from_repartition_exec(
- exec,
- codec,
- proto_converter,
- );
- }
-
if let Some(union) = plan.downcast_ref::<UnionExec>() {
return protobuf::PhysicalPlanNode::try_from_union_exec(
union,
@@ -1432,25 +1422,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<Arc<dyn ExecutionPlan>> {
- let input: Arc<dyn ExecutionPlan> =
- 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)
}
#[deprecated(
@@ -3362,28 +3354,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<protobuf::PhysicalPlanNode> {
- 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")
})
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org