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/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/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/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 e5f8aa072db71..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,14 +778,15 @@ 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, - ), - PhysicalPlanType::Merge(merge) => { - self.try_into_merge_physical_plan(merge, ctx, proto_converter) + #[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" + )] + PhysicalPlanType::CoalesceBatches(_) => { + CoalesceBatchesExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::Merge(_) => { + CoalescePartitionsExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::Repartition(repart) => { self.try_into_repartition_physical_plan(repart, ctx, proto_converter) @@ -850,8 +854,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) @@ -862,8 +866,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) @@ -985,15 +989,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, @@ -1004,14 +999,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, @@ -1094,14 +1081,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)? @@ -1117,14 +1096,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, @@ -1511,33 +1482,52 @@ 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, 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( + 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, 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( @@ -2777,14 +2767,27 @@ 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, 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( @@ -2824,16 +2827,25 @@ 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, 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( @@ -3500,25 +3512,26 @@ pub trait PhysicalPlanNodeExt: Sized { }) } - #[expect(deprecated)] + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CoalesceBatchesExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] + #[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") }) } @@ -3694,23 +3707,22 @@ 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, 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") }) } @@ -4130,23 +4142,22 @@ 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, 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") }) } @@ -4304,25 +4315,22 @@ 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, 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( diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index cbd6fd912abef..96b382b6d19eb 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -60,9 +60,11 @@ 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; +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 +1125,31 @@ 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_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 =