This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23825-562f87dff5e2ee381cb37e45ce6802ddefdebbc3 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 16471eeb91635fcb874c91249f9b57128a20a746 Author: Matthew Patton <[email protected]> AuthorDate: Fri Jul 24 18:10:00 2026 -0400 refactor(proto): migrate AsyncFuncExec to self-serializing proto (#23825) Closes #23514. Part of #23494. Migrate `AsyncFuncExec` proto encode/decode into the plan itself via `try_to_proto` / `try_from_proto`, removing its central-arm handling in `physical_plan/mod.rs`. ## Rationale `datafusion-proto` currently downcasts `AsyncFuncExec` in the central encode match and rebuilds it inline on decode. #23495 introduced self-serializing hooks so each plan owns its own wire format. This moves `AsyncFuncExec` onto that pattern. ## Changes - Added `try_to_proto` / `AsyncFuncExec::try_from_proto`, wired into the decode dispatch - Removed the central encode downcast branch - Old helper methods kept as `#[deprecated]` stubs, per existing convention - Wire format unchanged ## Testing Existing `roundtrip_async_func_exec` integration test now exercises the new hooks (old path deleted, so it's the only path left). `cargo fmt` + `cargo clippy --all-targets --features proto -- -D warnings` clean. ## User-facing changes No Co-authored-by: Matthew Patton <[email protected]> --- datafusion/physical-plan/src/async_func.rs | 77 +++++++++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 92 +++++++++--------------------- 2 files changed, 105 insertions(+), 64 deletions(-) diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index 5a65c9aedc..e13a5b986a 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -246,6 +246,83 @@ impl ExecutionPlan for AsyncFuncExec { fn metrics(&self) -> Option<MetricsSet> { Some(self.metrics.clone_inner()) } + + #[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())?; + let async_exprs = + ctx.encode_expressions(self.async_exprs.iter().map(|e| &e.func))?; + let async_expr_names = self + .async_exprs + .iter() + .map(|e| e.name().to_string()) + .collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc(Box::new( + protobuf::AsyncFuncExecNode { + input: Some(Box::new(input)), + async_exprs, + async_expr_names, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AsyncFuncExec { + /// Reconstruct an [`AsyncFuncExec`] 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. Child plans and expressions are 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<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let async_func = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc, + "AsyncFuncExec", + ); + let input = ctx.decode_required_child( + async_func.input.as_deref(), + "AsyncFuncExec", + "input", + )?; + let input_schema = input.schema(); + assert_eq_or_internal_err!( + async_func.async_exprs.len(), + async_func.async_expr_names.len(), + "AsyncFuncExecNode async_exprs length does not match async_expr_names" + ); + let async_exprs = async_func + .async_exprs + .iter() + .zip(async_func.async_expr_names.iter()) + .map(|(expr, name)| { + let physical_expr = ctx.decode_expr(expr, input_schema.as_ref())?; + Ok(Arc::new(AsyncFuncExpr::try_new( + name.clone(), + physical_expr, + input_schema.as_ref(), + )?)) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) + } } struct CoalesceInputStream { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 345ab87bf5..5cb9a922a8 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -57,7 +57,6 @@ use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; -use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_expr::{LexOrdering, LexRequirement}; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; @@ -855,8 +854,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::SortMergeJoin(_) => { SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::AsyncFunc(async_func) => { - self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) + PhysicalPlanType::AsyncFunc(_) => { + AsyncFuncExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::Buffer(_) => { BufferExec::try_from_proto(self.node(), &decode_ctx) @@ -958,14 +957,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(exec) = plan.downcast_ref::<AsyncFuncExec>() { - return protobuf::PhysicalPlanNode::try_from_async_func_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::<ScalarSubqueryExec>() { return protobuf::PhysicalPlanNode::try_from_scalar_subquery_exec( exec, @@ -2244,41 +2235,27 @@ pub trait PhysicalPlanNodeExt: Sized { CooperativeExec::try_from_proto(&node, &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `AsyncFuncExec` deserializes itself via `AsyncFuncExec::try_from_proto`" + )] fn try_into_async_func_physical_plan( &self, async_func: &protobuf::AsyncFuncExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<Arc<dyn ExecutionPlan>> { - let input: Arc<dyn ExecutionPlan> = - into_physical_plan(&async_func.input, ctx, proto_converter)?; - - if async_func.async_exprs.len() != async_func.async_expr_names.len() { - return internal_err!( - "AsyncFuncExecNode async_exprs length does not match async_expr_names" - ); - } - - let async_exprs = async_func - .async_exprs - .iter() - .zip(async_func.async_expr_names.iter()) - .map(|(expr, name)| { - let physical_expr = proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?; - - Ok(Arc::new(AsyncFuncExpr::try_new( - name.clone(), - physical_expr, - input.schema().as_ref(), - )?)) - }) - .collect::<Result<Vec<_>>>()?; - - Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AsyncFunc(Box::new( + async_func.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + AsyncFuncExec::try_from_proto(&node, &decode_ctx) } #[deprecated( @@ -3333,35 +3310,22 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(None) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `AsyncFuncExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_async_func_exec( exec: &AsyncFuncExec, - codec: &dyn PhysicalExtensionCodec, + extension_codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<protobuf::PhysicalPlanNode> { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - codec, + let encoder = ConverterPlanEncoder { + codec: extension_codec, proto_converter, - )?; - - let mut async_exprs = vec![]; - let mut async_expr_names = vec![]; - - for async_expr in exec.async_exprs() { - async_exprs - .push(proto_converter.physical_expr_to_proto(&async_expr.func, codec)?); - async_expr_names.push(async_expr.name.clone()) - } - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::AsyncFunc(Box::new( - protobuf::AsyncFuncExecNode { - input: Some(Box::new(input)), - async_exprs, - async_expr_names, - }, - ))), - }) + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx)? + .ok_or_else(|| internal_datafusion_err!("AsyncFuncExec is not serializable")) } #[deprecated( --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
