From d250bf572c2685b247ce10250879d501b4a1fb77 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 20 Jul 2026 13:02:21 +0800 Subject: [PATCH 1/5] refactor(proto): migrate CoalesceBatches serde CoalesceBatches serialization still depends on centralized plan downcasts, keeping protobuf ownership outside the execution plan. Move encoding and decoding behind its ExecutionPlan hooks while preserving the existing wire fields. Keep the legacy helpers as deprecated compatibility APIs. Refs #23500 Signed-off-by: Jiawei Zhao --- .../physical-plan/src/coalesce_batches.rs | 55 +++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 27 ++++----- 2 files changed, 67 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 667a8a3ce0697..c5b91767777f2 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -290,6 +290,61 @@ impl ExecutionPlan for CoalesceBatchesExec { ) as Arc) }) } + + #[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())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches( + Box::new(protobuf::CoalesceBatchesExecNode { + input: Some(Box::new(input)), + target_batch_size: self.target_batch_size() as u32, + fetch: self.fetch().map(|n| n as u32), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +#[expect(deprecated)] +impl CoalesceBatchesExec { + /// Reconstruct a [`CoalesceBatchesExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. The child plan is decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let coalesce_batches = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches, + "CoalesceBatchesExec", + ); + let input = ctx.decode_required_child( + coalesce_batches.input.as_deref(), + "CoalesceBatchesExec", + "input", + )?; + Ok(Arc::new( + CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) + .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), + )) + } } /// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details. diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index e5f8aa072db71..764d67a887038 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -775,12 +775,10 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::ArrowScan(scan) => { self.try_into_arrow_scan_physical_plan(scan, ctx, proto_converter) } - PhysicalPlanType::CoalesceBatches(coalesce_batches) => self - .try_into_coalesce_batches_physical_plan( - coalesce_batches, - ctx, - proto_converter, - ), + #[expect(deprecated)] + PhysicalPlanType::CoalesceBatches(_) => { + CoalesceBatchesExec::try_from_proto(self.node(), &decode_ctx) + } PhysicalPlanType::Merge(merge) => { self.try_into_merge_physical_plan(merge, ctx, proto_converter) } @@ -985,15 +983,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - #[expect(deprecated)] - if let Some(coalesce_batches) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_coalesce_batches_exec( - coalesce_batches, - codec, - proto_converter, - ); - } - if let Some(data_source_exec) = plan.downcast_ref::() && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_source_exec( data_source_exec, @@ -1511,6 +1500,10 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(DataSourceExec::from_data_source(source)) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CoalesceBatchesExec` deserializes itself via `CoalesceBatchesExec::try_from_proto`" + )] fn try_into_coalesce_batches_physical_plan( &self, coalesce_batches: &protobuf::CoalesceBatchesExecNode, @@ -3500,6 +3493,10 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CoalesceBatchesExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] #[expect(deprecated)] fn try_from_coalesce_batches_exec( coalesce_batches: &CoalesceBatchesExec, From 8d79da8edbad87db18525efaf15946f31e224376 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 20 Jul 2026 13:07:20 +0800 Subject: [PATCH 2/5] refactor(proto): migrate CoalescePartitions CoalescePartitions serialization still depends on the central Merge dispatch, keeping protobuf ownership outside the execution plan. Move encoding and decoding behind its ExecutionPlan hooks while preserving the historical Merge wire variant and optional fetch field. Keep the legacy helpers as deprecated compatibility APIs. Refs #23500 Signed-off-by: Jiawei Zhao --- .../physical-plan/src/coalesce_partitions.rs | 50 +++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 20 ++++---- 2 files changed, 60 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 1c6f95c53b018..f9694e0d16817 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -346,6 +346,56 @@ impl ExecutionPlan for CoalescePartitionsExec { } }) } + + #[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())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Merge(Box::new( + protobuf::CoalescePartitionsExecNode { + input: Some(Box::new(input)), + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CoalescePartitionsExec { + /// Reconstruct a [`CoalescePartitionsExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. Note the protobuf + /// variant is named `Merge` (node [`CoalescePartitionsExecNode`]). + /// + /// [`CoalescePartitionsExecNode`]: datafusion_proto_models::protobuf::CoalescePartitionsExecNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let merge = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Merge, + "CoalescePartitionsExec", + ); + let input = ctx.decode_required_child( + merge.input.as_deref(), + "CoalescePartitionsExec", + "input", + )?; + Ok(Arc::new( + CoalescePartitionsExec::new(input) + .with_fetch(merge.fetch.map(|f| f as usize)), + )) + } } #[cfg(test)] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 764d67a887038..9667b43c7afe8 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -779,8 +779,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::CoalesceBatches(_) => { CoalesceBatchesExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Merge(merge) => { - self.try_into_merge_physical_plan(merge, ctx, proto_converter) + PhysicalPlanType::Merge(_) => { + CoalescePartitionsExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::Repartition(repart) => { self.try_into_repartition_physical_plan(repart, ctx, proto_converter) @@ -993,14 +993,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_coalesce_partitions_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_repartition_exec( exec, @@ -1519,6 +1511,10 @@ pub trait PhysicalPlanNodeExt: Sized { )) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CoalescePartitionsExec` deserializes itself via `CoalescePartitionsExec::try_from_proto`" + )] fn try_into_merge_physical_plan( &self, merge: &protobuf::CoalescePartitionsExecNode, @@ -3691,6 +3687,10 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(None) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CoalescePartitionsExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_coalesce_partitions_exec( exec: &CoalescePartitionsExec, codec: &dyn PhysicalExtensionCodec, From b293a7f41eeddc3221fbeb55d62efc40f13b92a3 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 20 Jul 2026 13:12:45 +0800 Subject: [PATCH 3/5] refactor(proto): migrate CooperativeExec serde CooperativeExec serialization still depends on centralized plan dispatch despite carrying only one child. Move encoding and decoding behind its ExecutionPlan hooks, retain the legacy helpers as deprecated compatibility APIs, and add a direct round-trip test that proves the hooks are reached. Refs #23500 Signed-off-by: Jiawei Zhao --- datafusion/physical-plan/src/coop.rs | 44 +++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 20 ++++----- .../tests/cases/roundtrip_physical_plan.rs | 9 ++++ 3 files changed, 63 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index 94e4fdca2b53e..a5b57f546bbfa 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -369,6 +369,50 @@ impl ExecutionPlan for CooperativeExec { } } } + + #[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())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Cooperative(Box::new( + protobuf::CooperativeExecNode { + input: Some(Box::new(input)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CooperativeExec { + /// Reconstruct a [`CooperativeExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let cooperative = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Cooperative, + "CooperativeExec", + ); + let input = ctx.decode_required_child( + cooperative.input.as_deref(), + "CooperativeExec", + "input", + )?; + Ok(Arc::new(CooperativeExec::new(input))) + } } /// Creates a [`CooperativeStream`] wrapper around the given [`RecordBatchStream`]. diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 9667b43c7afe8..73a6651492098 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -848,8 +848,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::Unnest(unnest) => { self.try_into_unnest_physical_plan(unnest, ctx, proto_converter) } - PhysicalPlanType::Cooperative(cooperative) => { - self.try_into_cooperative_physical_plan(cooperative, ctx, proto_converter) + PhysicalPlanType::Cooperative(_) => { + CooperativeExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::GenerateSeries(generate_series) => { self.try_into_generate_series_physical_plan(generate_series) @@ -1075,14 +1075,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_cooperative_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() && let Some(node) = protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)? @@ -2766,6 +2758,10 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CooperativeExec` deserializes itself via `CooperativeExec::try_from_proto`" + )] fn try_into_cooperative_physical_plan( &self, field_stream: &protobuf::CooperativeExecNode, @@ -4127,6 +4123,10 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CooperativeExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_cooperative_exec( exec: &CooperativeExec, codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index cbd6fd912abef..d060175572944 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -63,6 +63,7 @@ use datafusion::physical_plan::analyze::AnalyzeExec; #[expect(deprecated)] use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::coop::CooperativeExec; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::expressions::{ BinaryExpr, Column, DynamicFilterPhysicalExpr, NotExpr, PhysicalSortExpr, binary, @@ -1123,6 +1124,14 @@ fn roundtrip_coalesce_partitions_with_fetch() -> Result<()> { )) } +#[test] +fn roundtrip_cooperative() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Boolean, false)])); + roundtrip_test(Arc::new(CooperativeExec::new(Arc::new(EmptyExec::new( + schema, + ))))) +} + #[test] fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { let file_schema = From 48ab1735ba23dbbf6a8fd862a90577ada5313a0c Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 20 Jul 2026 13:18:00 +0800 Subject: [PATCH 4/5] refactor(proto): migrate BufferExec serde BufferExec serialization still depends on centralized plan dispatch, keeping its protobuf capacity representation outside the execution plan. Move encoding and decoding behind its ExecutionPlan hooks, retain the legacy helpers as deprecated compatibility APIs, and verify capacity preservation with a direct round-trip test. CLOSES #23500 Signed-off-by: Jiawei Zhao --- datafusion/physical-plan/src/buffer.rs | 42 +++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 20 ++++----- .../tests/cases/roundtrip_physical_plan.rs | 18 ++++++++ 3 files changed, 70 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 6f83e3719690a..3be331a1ee1ba 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -298,6 +298,48 @@ impl ExecutionPlan for BufferExec { Ok(Arc::new(Self::new(new_input, self.capacity)) as Arc) }) } + + #[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())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Buffer(Box::new( + protobuf::BufferExecNode { + input: Some(Box::new(input)), + capacity: self.capacity() as u64, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl BufferExec { + /// Reconstruct a [`BufferExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let buffer = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Buffer, + "BufferExec", + ); + let input = + ctx.decode_required_child(buffer.input.as_deref(), "BufferExec", "input")?; + Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) + } } /// Represents anything that occupies a capacity in a [MemoryBufferedStream]. diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 73a6651492098..dadb6e74e5192 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -860,8 +860,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::AsyncFunc(async_func) => { self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) } - PhysicalPlanType::Buffer(buffer) => { - self.try_into_buffer_physical_plan(buffer, ctx, proto_converter) + PhysicalPlanType::Buffer(_) => { + BufferExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::ScalarSubquery(sq) => { self.try_into_scalar_subquery_physical_plan(sq, ctx, proto_converter) @@ -1090,14 +1090,6 @@ pub trait PhysicalPlanNodeExt: Sized { ); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_buffer_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_scalar_subquery_exec( exec, @@ -2809,6 +2801,10 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `BufferExec` deserializes itself via `BufferExec::try_from_proto`" + )] fn try_into_buffer_physical_plan( &self, buffer: &protobuf::BufferExecNode, @@ -4301,6 +4297,10 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `BufferExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_buffer_exec( exec: &BufferExec, extension_codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index d060175572944..96b382b6d19eb 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -60,6 +60,7 @@ use datafusion::physical_plan::aggregates::{ AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, }; use datafusion::physical_plan::analyze::AnalyzeExec; +use datafusion::physical_plan::buffer::BufferExec; #[expect(deprecated)] use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; @@ -1132,6 +1133,23 @@ fn roundtrip_cooperative() -> Result<()> { ))))) } +#[test] +fn roundtrip_buffer() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Boolean, false)])); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result = roundtrip_test_and_return( + Arc::new(BufferExec::new(Arc::new(EmptyExec::new(schema)), 4096)), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.capacity(), 4096); + Ok(()) +} + #[test] fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { let file_schema = From 1c33ac2f9985408a5379717800fa8f3ea87b1407 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 11:56:24 +0800 Subject: [PATCH 5/5] refactor(proto): delegate deprecated plan serde helpers The migrated PhysicalPlanNodeExt helpers retained independent encode and decode implementations after per-plan protobuf hooks became canonical. This duplicated wire-format logic and allowed the deprecated compatibility paths to drift from normal dispatch. Make the CoalesceBatches, CoalescePartitions, Cooperative, and Buffer helpers delegate through the execution-plan encode and decode contexts. Add coverage that compares legacy helper output with the canonical hooks for both serialization and deserialization. Signed-off-by: Jiawei Zhao --- datafusion/proto/src/physical_plan/mod.rs | 143 ++++++++++++---------- 1 file changed, 77 insertions(+), 66 deletions(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index dadb6e74e5192..62c3b955e31e3 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -68,7 +68,10 @@ use datafusion_physical_plan::aggregates::{ use datafusion_physical_plan::analyze::AnalyzeExec; use datafusion_physical_plan::async_func::AsyncFuncExec; use datafusion_physical_plan::buffer::BufferExec; -#[expect(deprecated)] +#[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" +)] use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::coop::CooperativeExec; @@ -775,7 +778,10 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::ArrowScan(scan) => { self.try_into_arrow_scan_physical_plan(scan, ctx, proto_converter) } - #[expect(deprecated)] + #[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" + )] PhysicalPlanType::CoalesceBatches(_) => { CoalesceBatchesExec::try_from_proto(self.node(), &decode_ctx) } @@ -1486,13 +1492,21 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input: Arc = - into_physical_plan(&coalesce_batches.input, ctx, proto_converter)?; - Ok(Arc::new( - #[expect(deprecated)] - CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) - .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), - )) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new( + coalesce_batches.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + #[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" + )] + CoalesceBatchesExec::try_from_proto(&node, &decode_ctx) } #[deprecated( @@ -1505,12 +1519,15 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input: Arc = - into_physical_plan(&merge.input, ctx, proto_converter)?; - Ok(Arc::new( - CoalescePartitionsExec::new(input) - .with_fetch(merge.fetch.map(|f| f as usize)), - )) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Merge(Box::new(merge.clone()))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + CoalescePartitionsExec::try_from_proto(&node, &decode_ctx) } fn try_into_repartition_physical_plan( @@ -2760,8 +2777,17 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input = into_physical_plan(&field_stream.input, ctx, proto_converter)?; - Ok(Arc::new(CooperativeExec::new(input))) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Cooperative(Box::new( + field_stream.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + CooperativeExec::try_from_proto(&node, &decode_ctx) } fn try_into_async_func_physical_plan( @@ -2811,10 +2837,15 @@ pub trait PhysicalPlanNodeExt: Sized { ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input: Arc = - into_physical_plan(&buffer.input, ctx, proto_converter)?; - - Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Buffer(Box::new(buffer.clone()))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + BufferExec::try_from_proto(&node, &decode_ctx) } fn try_into_scalar_subquery_physical_plan( @@ -3485,25 +3516,22 @@ pub trait PhysicalPlanNodeExt: Sized { since = "55.0.0", note = "unused by DataFusion; `CoalesceBatchesExec` serializes itself via `ExecutionPlan::try_to_proto`" )] - #[expect(deprecated)] + #[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" + )] fn try_from_coalesce_batches_exec( coalesce_batches: &CoalesceBatchesExec, codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - coalesce_batches.input().to_owned(), + let encoder = ConverterPlanEncoder { codec, proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new( - protobuf::CoalesceBatchesExecNode { - input: Some(Box::new(input)), - target_batch_size: coalesce_batches.target_batch_size() as u32, - fetch: coalesce_batches.fetch().map(|n| n as u32), - }, - ))), + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + coalesce_batches.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("CoalesceBatchesExec is not serializable") }) } @@ -3688,18 +3716,13 @@ pub trait PhysicalPlanNodeExt: Sized { 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, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Merge(Box::new( - protobuf::CoalescePartitionsExecNode { - input: Some(Box::new(input)), - fetch: exec.fetch().map(|f| f as u32), - }, - ))), + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("CoalescePartitionsExec is not serializable") }) } @@ -4128,18 +4151,13 @@ pub trait PhysicalPlanNodeExt: Sized { 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, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Cooperative(Box::new( - protobuf::CooperativeExecNode { - input: Some(Box::new(input)), - }, - ))), + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)?.ok_or_else(|| { + internal_datafusion_err!("CooperativeExec is not serializable") }) } @@ -4306,20 +4324,13 @@ pub trait PhysicalPlanNodeExt: Sized { extension_codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - extension_codec, + let encoder = ConverterPlanEncoder { + codec: extension_codec, proto_converter, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Buffer(Box::new( - protobuf::BufferExecNode { - input: Some(Box::new(input)), - capacity: exec.capacity() as u64, - }, - ))), - }) + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)? + .ok_or_else(|| internal_datafusion_err!("BufferExec is not serializable")) } fn try_from_scalar_subquery_exec(