From 63b18508dde30d93823c209d439386ebf31fb109 Mon Sep 17 00:00:00 2001 From: Matthew Patton Date: Wed, 22 Jul 2026 19:52:55 -0400 Subject: [PATCH] refactor(proto): migrate AsyncFuncExec to self-serializing proto --- 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 5a65c9aedc2f1..e13a5b986aa2c 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 { Some(self.metrics.clone_inner()) } + + #[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())?; + 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> { + 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::>>()?; + 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 cea334e42aace..555f78a0c6dc6 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -59,7 +59,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, PhysicalExprRef}; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; @@ -865,8 +864,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) @@ -1074,14 +1073,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_async_func_exec( - exec, - codec, - proto_converter, - ); - } - if let Some(exec) = plan.downcast_ref::() { return protobuf::PhysicalPlanNode::try_from_scalar_subquery_exec( exec, @@ -2657,41 +2648,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> { - let input: Arc = - 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::>>()?; - - 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( @@ -4069,35 +4046,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 { - 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(