From 1388fb1e1a675167a088a8ac875ab851b5ad18d5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 10 Jul 2026 21:55:18 -0500 Subject: [PATCH 1/5] prototype: ExecutionPlan try_to_proto/try_from_proto hook + ProjectionExec reference (#22419) General part: ExecutionPlanEncodeCtx/DecodeCtx + internal dispatch traits in datafusion-physical-plan (mirrors the PhysicalExpr pattern), the try_to_proto trait hook, expect_plan_variant! macro, and the ConverterPlanEncoder/Decoder adapters wiring the hook into the central dispatch in datafusion-proto. ProjectionExec migrated as the reference: encode downcast arm deleted, decode arm reduced to a one-liner delegating to ProjectionExec::try_from_proto. All 189 proto_integration roundtrip tests pass (incl. TPC-H). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01SfBPC2yYCwHUn4V8pQCVj3 --- .../physical-plan/src/execution_plan.rs | 21 ++ datafusion/physical-plan/src/lib.rs | 2 + datafusion/physical-plan/src/projection.rs | 65 +++++ datafusion/physical-plan/src/proto.rs | 230 ++++++++++++++++++ datafusion/proto/src/physical_plan/mod.rs | 158 ++++++------ 5 files changed, 402 insertions(+), 74 deletions(-) create mode 100644 datafusion/physical-plan/src/proto.rs diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index e58acee4ce6bd..d03738b44e056 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -823,6 +823,27 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { ) -> Option> { None } + + /// Serialize this plan to its protobuf representation, if it knows how. + /// + /// This is the `ExecutionPlan` analog of + /// [`PhysicalExpr::try_to_proto`](datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto). + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller + /// (`datafusion-proto`) falls back to the central downcast chain. Every + /// un-migrated plan keeps its existing behavior. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// * `Err(_)` — a real failure (e.g. a child failed to serialize). + /// + /// Only *self-contained* plans should override this — see [`crate::proto`] + /// for the session-dependency boundary. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn ExecutionPlan { diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 0ab232bc102bc..8cba650b79770 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -88,6 +88,8 @@ pub mod metrics; pub mod operator_statistics; pub mod placeholder_row; pub mod projection; +#[cfg(feature = "proto")] +pub mod proto; pub mod recursive_query; pub mod repartition; pub mod scalar_subquery; diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index d55363297bd52..42501f22395b4 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -500,6 +500,71 @@ impl ExecutionPlan for ProjectionExec { .ok() }) } + + #[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 expr = ctx.encode_expressions(self.expr().iter().map(|p| &p.expr))?; + let expr_name = self.expr().iter().map(|p| p.alias.clone()).collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Projection(Box::new( + protobuf::ProjectionExecNode { + input: Some(Box::new(input)), + expr, + expr_name, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ProjectionExec { + /// Reconstruct a [`ProjectionExec`] 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 projection = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Projection, + "ProjectionExec", + ); + let input = ctx.decode_required_child( + projection.input.as_deref(), + "ProjectionExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = projection + .expr + .iter() + .zip(projection.expr_name.iter()) + .map(|(expr, name)| { + Ok(ProjectionExpr { + expr: ctx.decode_expr(expr, input_schema.as_ref())?, + alias: name.to_string(), + }) + }) + .collect::>>()?; + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) + } } impl ProjectionStream { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs new file mode 100644 index 0000000000000..10aa72170a04f --- /dev/null +++ b/datafusion/physical-plan/src/proto.rs @@ -0,0 +1,230 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Serialization hooks for [`ExecutionPlan`], mirroring the +//! `try_to_proto`/`try_from_proto` pattern used for `PhysicalExpr`. +//! +//! # Why the indirection +//! +//! An `ExecutionPlan` must be able to (de)serialize its child plans and its +//! child physical expressions recursively. The concrete recursion lives in +//! `datafusion-proto` (it owns the extension codec, the session context and the +//! central converter), but `datafusion-proto` sits *above* `datafusion-physical-plan` +//! in the crate graph. To let a plan drive that recursion without a dependency +//! cycle, this module defines: +//! +//! * [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] — the stable, +//! concrete context types a plan author interacts with. New capabilities can +//! be added here without changing every plan's hook signature. +//! * [`ExecutionPlanEncode`] / [`ExecutionPlanDecode`] — internal dispatch +//! traits, *defined* here but *implemented* in `datafusion-proto`, that the +//! context types delegate to. This is the dependency inversion that keeps the +//! proto types flowing in one direction only. +//! +//! `datafusion-physical-plan` depends on the pure prost types in +//! `datafusion-proto-models` (feature `proto`), never on `datafusion-proto`. +//! +//! # Scope: self-contained plans only +//! +//! These hooks are for plans that are *self-contained*: everything they need to +//! round-trip is reachable from their child plans, their child expressions and +//! their input schema. Plans that reference session-scoped machinery only +//! `datafusion-proto` owns — the [`PhysicalExtensionCodec`] (for UDAF/UDWF and +//! extension nodes) — stay as typed dispatch arms inside `datafusion-proto`, +//! exactly as `ScalarFunctionExpr` does on the expression side. The decode +//! context intentionally exposes the [`TaskContext`] (for function-registry and +//! session-config lookups reachable from `datafusion-execution`) but *not* the +//! extension codec. +//! +//! [`PhysicalExtensionCodec`]: (owned by `datafusion-proto`) +//! [`ExecutionPlan`]: crate::ExecutionPlan + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; + +use crate::ExecutionPlan; + +/// Internal dispatch trait backing [`ExecutionPlanEncodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanEncodeCtx`] instead. +pub trait ExecutionPlanEncode { + /// Serialize a child execution plan (recursing through the central + /// serializer, so the child's own `try_to_proto` hook is honored). + fn encode_plan(&self, plan: &Arc) -> Result; + + /// Serialize a physical expression owned by the plan. + fn encode_expr(&self, expr: &Arc) -> Result; +} + +/// Internal dispatch trait backing [`ExecutionPlanDecodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanDecodeCtx`] instead. +pub trait ExecutionPlanDecode { + /// Deserialize a child execution plan (recursing through the central + /// deserializer, so the child's own `try_from_proto` is honored). + fn decode_plan(&self, node: &PhysicalPlanNode) -> Result>; + + /// Deserialize a physical expression against `input_schema`. + fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result>; + + /// The session task context, used by plans that need the function registry + /// or session configuration. Never exposes the proto extension codec. + fn task_ctx(&self) -> &TaskContext; +} + +/// Context handed to [`ExecutionPlan::try_to_proto`](crate::ExecutionPlan::try_to_proto). +/// +/// Provides the primitives a plan needs to serialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanEncodeCtx<'a> { + encoder: &'a dyn ExecutionPlanEncode, +} + +impl<'a> ExecutionPlanEncodeCtx<'a> { + /// Create a new encode context wrapping an [`ExecutionPlanEncode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(encoder: &'a dyn ExecutionPlanEncode) -> Self { + Self { encoder } + } + + /// Serialize a single child plan. + pub fn encode_child( + &self, + plan: &Arc, + ) -> Result { + self.encoder.encode_plan(plan) + } + + /// Serialize an iterator of child plans. + pub fn encode_children<'b, I>(&self, plans: I) -> Result> + where + I: IntoIterator>, + { + plans.into_iter().map(|p| self.encode_child(p)).collect() + } + + /// Serialize a single physical expression. + pub fn encode_expr(&self, expr: &Arc) -> Result { + self.encoder.encode_expr(expr) + } + + /// Serialize an iterator of physical expressions. + pub fn encode_expressions<'b, I>(&self, exprs: I) -> Result> + where + I: IntoIterator>, + { + exprs.into_iter().map(|e| self.encode_expr(e)).collect() + } +} + +/// Context handed to a plan's `try_from_proto` associated function. +/// +/// Provides the primitives a plan needs to deserialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanDecodeCtx<'a> { + decoder: &'a dyn ExecutionPlanDecode, +} + +impl<'a> ExecutionPlanDecodeCtx<'a> { + /// Create a new decode context wrapping an [`ExecutionPlanDecode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(decoder: &'a dyn ExecutionPlanDecode) -> Self { + Self { decoder } + } + + /// Deserialize a single child plan. + pub fn decode_child( + &self, + node: &PhysicalPlanNode, + ) -> Result> { + self.decoder.decode_plan(node) + } + + /// Deserialize a required child plan, producing a uniform "missing required + /// field" error when the optional wire field is absent. + pub fn decode_required_child( + &self, + node: Option<&PhysicalPlanNode>, + plan_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_child(node) + } + + /// Deserialize a physical expression against `input_schema`. + pub fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result> { + self.decoder.decode_expr(node, input_schema) + } + + /// Deserialize a required physical expression against `input_schema`. + pub fn decode_required_expr( + &self, + node: Option<&PhysicalExprNode>, + input_schema: &Schema, + plan_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_expr(node, input_schema) + } + + /// The session task context (function registry + session config). Never + /// exposes the proto extension codec — session-codec-dependent plans stay as + /// typed dispatch arms in `datafusion-proto`. + pub fn task_ctx(&self) -> &TaskContext { + self.decoder.task_ctx() + } +} + +/// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` +/// variant, returning a reference to the inner payload, else an `internal_err!`. +/// Mirrors `expect_expr_variant!` on the expression side. Field access on the +/// result auto-derefs through the `Box` that boxed variants use. +#[macro_export] +macro_rules! expect_plan_variant { + ($node:expr, $variant:path, $plan_name:literal $(,)?) => {{ + match &$node.physical_plan_type { + Some($variant(inner)) => inner, + _ => { + return ::datafusion_common::internal_err!(concat!( + "PhysicalPlanNode is not a ", + $plan_name + )); + } + } + }}; +} diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 0744a94dcebd1..1879332b45454 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -85,7 +85,11 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; use datafusion_physical_plan::metrics::MetricCategory; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecode, ExecutionPlanDecodeCtx, ExecutionPlanEncode, + ExecutionPlanEncodeCtx, +}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; use datafusion_physical_plan::sorts::sort::SortExec; @@ -312,12 +316,20 @@ pub trait PhysicalPlanNodeExt: Sized { self.node(), )) })?; + // Decode context for plans migrated to the `try_from_proto` pattern + // (#22419). Arms for migrated plans are one-liners delegating to the + // plan's own crate; un-migrated arms keep their inline bodies. + let plan_decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&plan_decoder); match plan { PhysicalPlanType::Explain(explain) => { self.try_into_explain_physical_plan(explain, ctx, proto_converter) } - PhysicalPlanType::Projection(projection) => { - self.try_into_projection_physical_plan(projection, ctx, proto_converter) + PhysicalPlanType::Projection(_) => { + ProjectionExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::Filter(filter) => { self.try_into_filter_physical_plan(filter, ctx, proto_converter) @@ -444,16 +456,20 @@ pub trait PhysicalPlanNodeExt: Sized { let plan_clone = Arc::clone(&plan); let plan = plan.as_ref(); - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec); + // Self-serializing plans handle themselves via the `try_to_proto` hook + // (#22419). `Ok(None)` means "not migrated" and falls through to the + // central downcast chain below. + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + if let Some(node) = plan.try_to_proto(&encode_ctx)? { + return Ok(node); } - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_projection_exec( - exec, - codec, - proto_converter, - ); + if let Some(exec) = plan.downcast_ref::() { + return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec); } if let Some(exec) = plan.downcast_ref::() { @@ -731,36 +747,6 @@ pub trait PhysicalPlanNodeExt: Sized { ))) } - fn try_into_projection_physical_plan( - &self, - projection: &protobuf::ProjectionExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&projection.input, ctx, proto_converter)?; - let exprs = projection - .expr - .iter() - .zip(projection.expr_name.iter()) - .map(|(expr, name)| { - Ok(( - proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - name.to_string(), - )) - }) - .collect::, String)>>>()?; - let proj_exprs: Vec = exprs - .into_iter() - .map(|(expr, alias)| ProjectionExpr { expr, alias }) - .collect(); - Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) - } - fn try_into_filter_physical_plan( &self, filter: &protobuf::FilterExecNode, @@ -2430,39 +2416,6 @@ pub trait PhysicalPlanNodeExt: Sized { }) } - fn try_from_projection_exec( - exec: &ProjectionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let expr = exec - .expr() - .iter() - .map(|proj_expr| { - proto_converter.physical_expr_to_proto(&proj_expr.expr, codec) - }) - .collect::>>()?; - let expr_name = exec - .expr() - .iter() - .map(|proj_expr| proj_expr.alias.clone()) - .collect(); - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Projection(Box::new( - protobuf::ProjectionExecNode { - input: Some(Box::new(input)), - expr, - expr_name, - }, - ))), - }) - } - fn try_from_analyze_exec( exec: &AnalyzeExec, codec: &dyn PhysicalExtensionCodec, @@ -4387,3 +4340,60 @@ fn into_physical_plan( Err(proto_error("Missing required field in protobuf")) } } + +/// Adapter backing [`ExecutionPlanEncodeCtx`] for plans migrated to the +/// `try_to_proto` hook (#22419). Routes child-plan and child-expr encoding back +/// through the central converter so nested plans honor their own hooks. +struct ConverterPlanEncoder<'a> { + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl ExecutionPlanEncode for ConverterPlanEncoder<'_> { + fn encode_plan( + &self, + plan: &Arc, + ) -> Result { + self.proto_converter + .execution_plan_to_proto(plan, self.codec) + } + + fn encode_expr( + &self, + expr: &Arc, + ) -> Result { + self.proto_converter + .physical_expr_to_proto(expr, self.codec) + } +} + +/// Adapter backing [`ExecutionPlanDecodeCtx`] for plans migrated to the +/// `try_from_proto` pattern (#22419). Routes child-plan and child-expr decoding +/// back through the central converter, and exposes the session task context +/// (never the extension codec). +struct ConverterPlanDecoder<'a, 'ctx> { + ctx: &'a PhysicalPlanDecodeContext<'ctx>, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl ExecutionPlanDecode for ConverterPlanDecoder<'_, '_> { + fn decode_plan( + &self, + node: &protobuf::PhysicalPlanNode, + ) -> Result> { + self.proto_converter.proto_to_execution_plan(node, self.ctx) + } + + fn decode_expr( + &self, + node: &protobuf::PhysicalExprNode, + input_schema: &Schema, + ) -> Result> { + self.proto_converter + .proto_to_physical_expr(node, input_schema, self.ctx) + } + + fn task_ctx(&self) -> &TaskContext { + self.ctx.task_ctx() + } +} From 281a95581ab9b461a5242082fb9cd4091a327ac5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sat, 11 Jul 2026 09:48:55 -0500 Subject: [PATCH 2/5] prototype: add typed bytes-only function serde to the ExecutionPlan proto ctx (#22419) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds encode_udf/udaf/udwf + decode_udf/udaf/udwf to ExecutionPlanEncodeCtx/ DecodeCtx (and the internal dispatch traits). These take datafusion-expr types + Vec and never name a proto type; the datafusion-proto adapter backs them over PhysicalExtensionCodec + the registry, owning the payload/registry lookup-order policy in one place. Wire-identical to the existing fun_definition semantics ((!buf.is_empty()).then_some(buf)). This lets function-carrying plans (Aggregate, window, ScalarSubquery) ride the polymorphic hook instead of staying as typed arms — possible because physical-plan sits above datafusion-expr (unlike the expr-side ctx). #23421's ScalarUdfCodec stays closed; ScalarFunctionExpr stays special-cased on the expr side where the crate cycle forces it. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01SfBPC2yYCwHUn4V8pQCVj3 --- datafusion/physical-plan/src/proto.rs | 103 +++++++++++++++++++--- datafusion/proto/src/physical_plan/mod.rs | 60 +++++++++++++ 2 files changed, 150 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs index 10aa72170a04f..e10c4d7c37eb9 100644 --- a/datafusion/physical-plan/src/proto.rs +++ b/datafusion/physical-plan/src/proto.rs @@ -38,19 +38,23 @@ //! `datafusion-physical-plan` depends on the pure prost types in //! `datafusion-proto-models` (feature `proto`), never on `datafusion-proto`. //! -//! # Scope: self-contained plans only +//! # Function-carrying plans //! -//! These hooks are for plans that are *self-contained*: everything they need to -//! round-trip is reachable from their child plans, their child expressions and -//! their input schema. Plans that reference session-scoped machinery only -//! `datafusion-proto` owns — the [`PhysicalExtensionCodec`] (for UDAF/UDWF and -//! extension nodes) — stay as typed dispatch arms inside `datafusion-proto`, -//! exactly as `ScalarFunctionExpr` does on the expression side. The decode -//! context intentionally exposes the [`TaskContext`] (for function-registry and -//! session-config lookups reachable from `datafusion-execution`) but *not* the -//! extension codec. +//! Plans that reference UD(A/W)Fs (`AggregateExec`, the window execs, …) also +//! ride the hook: the context exposes typed, *bytes-only* function serde — +//! [`encode_udaf`](ExecutionPlanEncodeCtx::encode_udaf) / +//! [`decode_udaf`](ExecutionPlanDecodeCtx::decode_udaf) and the udf/udwf +//! siblings. These take/return `datafusion-expr` types plus `Vec` and never +//! name a proto type, so the `PhysicalExtensionCodec` (which only +//! `datafusion-proto` can name) stays fully encapsulated behind the adapter that +//! backs these traits. The lookup-order policy (payload → codec; else registry → +//! codec fallback) lives once, in that adapter, rather than in every plan. +//! +//! This is possible because `datafusion-physical-plan` sits *above* +//! `datafusion-expr` in the crate graph; the expression-side ctx (in +//! `physical-expr-common`, *below* `datafusion-expr`) cannot do this, which is +//! why `ScalarFunctionExpr` remains special-cased there. //! -//! [`PhysicalExtensionCodec`]: (owned by `datafusion-proto`) //! [`ExecutionPlan`]: crate::ExecutionPlan use std::sync::Arc; @@ -58,6 +62,7 @@ use std::sync::Arc; use arrow::datatypes::Schema; use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::TaskContext; +use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::PhysicalExpr; use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; @@ -74,6 +79,18 @@ pub trait ExecutionPlanEncode { /// Serialize a physical expression owned by the plan. fn encode_expr(&self, expr: &Arc) -> Result; + + /// Serialize a scalar UDF to an opaque payload. `None` means "decodable by + /// name alone" (built-ins). Bytes-only: no proto types cross this boundary. + fn encode_udf(&self, udf: &ScalarUDF) -> Result>>; + + /// Serialize an aggregate UDF to an opaque payload. `None` means "decodable + /// by name alone". + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>>; + + /// Serialize a window UDF to an opaque payload. `None` means "decodable by + /// name alone". + fn encode_udwf(&self, udwf: &WindowUDF) -> Result>>; } /// Internal dispatch trait backing [`ExecutionPlanDecodeCtx`]. @@ -95,6 +112,21 @@ pub trait ExecutionPlanDecode { /// The session task context, used by plans that need the function registry /// or session configuration. Never exposes the proto extension codec. fn task_ctx(&self) -> &TaskContext; + + /// Reconstruct a scalar UDF from its name and optional payload. Encapsulates + /// the lookup-order policy (payload → codec; else registry → codec fallback) + /// so no plan re-derives it. Bytes-only: no proto types cross this boundary. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result>; + + /// Reconstruct an aggregate UDF from its name and optional payload. + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result>; + + /// Reconstruct a window UDF from its name and optional payload. + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result>; } /// Context handed to [`ExecutionPlan::try_to_proto`](crate::ExecutionPlan::try_to_proto). @@ -140,6 +172,23 @@ impl<'a> ExecutionPlanEncodeCtx<'a> { { exprs.into_iter().map(|e| self.encode_expr(e)).collect() } + + /// Serialize a scalar UDF to an opaque payload (`None` = built-in, decodable + /// by name). No proto types cross this boundary. + pub fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + self.encoder.encode_udf(udf) + } + + /// Serialize an aggregate UDF to an opaque payload (`None` = decodable by + /// name). + pub fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>> { + self.encoder.encode_udaf(udaf) + } + + /// Serialize a window UDF to an opaque payload (`None` = decodable by name). + pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { + self.encoder.encode_udwf(udwf) + } } /// Context handed to a plan's `try_from_proto` associated function. @@ -203,11 +252,39 @@ impl<'a> ExecutionPlanDecodeCtx<'a> { } /// The session task context (function registry + session config). Never - /// exposes the proto extension codec — session-codec-dependent plans stay as - /// typed dispatch arms in `datafusion-proto`. + /// exposes the proto extension codec. pub fn task_ctx(&self) -> &TaskContext { self.decoder.task_ctx() } + + /// Reconstruct a scalar UDF from its name and optional payload. The + /// lookup-order policy is owned by `datafusion-proto`; no proto types cross + /// this boundary. + pub fn decode_udf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udf(name, payload) + } + + /// Reconstruct an aggregate UDF from its name and optional payload. + pub fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udaf(name, payload) + } + + /// Reconstruct a window UDF from its name and optional payload. + pub fn decode_udwf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udwf(name, payload) + } } /// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 1879332b45454..9f0f16250e752 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -4365,6 +4365,26 @@ impl ExecutionPlanEncode for ConverterPlanEncoder<'_> { self.proto_converter .physical_expr_to_proto(expr, self.codec) } + + // Bytes-only function serde. `(!buf.is_empty()).then_some(buf)` preserves the + // existing `fun_definition` wire semantics (empty payload == encode-by-name). + fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udf(udf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } + + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udaf(udaf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } + + fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udwf(udwf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } } /// Adapter backing [`ExecutionPlanDecodeCtx`] for plans migrated to the @@ -4396,4 +4416,44 @@ impl ExecutionPlanDecode for ConverterPlanDecoder<'_, '_> { fn task_ctx(&self) -> &TaskContext { self.ctx.task_ctx() } + + // Lookup-order policy, owned here so no plan re-derives it: an explicit + // payload is decoded by the codec; otherwise resolve by name from the + // registry, falling back to the codec with an empty buffer. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udf(name, buf), + None => self + .ctx + .task_ctx() + .udf(name) + .or_else(|_| self.ctx.codec().try_decode_udf(name, &[])), + } + } + + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udaf(name, buf), + None => self + .ctx + .task_ctx() + .udaf(name) + .or_else(|_| self.ctx.codec().try_decode_udaf(name, &[])), + } + } + + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udwf(name, buf), + None => self + .ctx + .task_ctx() + .udwf(name) + .or_else(|_| self.ctx.codec().try_decode_udwf(name, &[])), + } + } } From 02108ebfaba3817119ee7ed05e96d1b6dd1344c4 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:18:49 -0500 Subject: [PATCH 3/5] fix: resolve downcast_delegate before the try_to_proto hook; fix rustdoc links Review follow-up (#23495): - A wrapper plan delegating its downcast identity to a migrated plan (e.g. ProjectionExec) hit the wrapper's default try_to_proto and then found no fallback arm for the delegate. Resolve the delegate chain before calling the hook, matching downcast_ref semantics; covered by a new serialize test that fails without the fix. - Drop redundant explicit rustdoc link targets that failed cargo doc with -D warnings. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01S4HUcyoHG8zRpWNrXRLmpg --- .../physical-plan/src/execution_plan.rs | 2 +- datafusion/physical-plan/src/proto.rs | 3 +- datafusion/proto/src/physical_plan/mod.rs | 10 ++++++- .../tests/cases/roundtrip_physical_plan.rs | 29 +++++++++++++++++++ 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index d03738b44e056..3b9d5d258a838 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -827,7 +827,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// Serialize this plan to its protobuf representation, if it knows how. /// /// This is the `ExecutionPlan` analog of - /// [`PhysicalExpr::try_to_proto`](datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto). + /// [`PhysicalExpr::try_to_proto`]. /// /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller /// (`datafusion-proto`) falls back to the central downcast chain. Every diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs index e10c4d7c37eb9..1731203f6c767 100644 --- a/datafusion/physical-plan/src/proto.rs +++ b/datafusion/physical-plan/src/proto.rs @@ -129,7 +129,8 @@ pub trait ExecutionPlanDecode { fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result>; } -/// Context handed to [`ExecutionPlan::try_to_proto`](crate::ExecutionPlan::try_to_proto). +/// Context handed to [`ExecutionPlan::try_to_proto`]. +/// /// /// Provides the primitives a plan needs to serialize its children and /// expressions without naming `datafusion-proto`. diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 9f0f16250e752..a66603a385275 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -454,7 +454,15 @@ pub trait PhysicalPlanNodeExt: Sized { proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { let plan_clone = Arc::clone(&plan); - let plan = plan.as_ref(); + let mut plan = plan.as_ref(); + // Resolve the downcast identity first so wrapper plans serialize as + // their delegate, matching how the `downcast_ref` chain below sees + // them. Without this a wrapper around a migrated plan would hit the + // wrapper's default `try_to_proto` (`Ok(None)`) and find no fallback + // arm for the delegate. + while let Some(delegate) = plan.downcast_delegate() { + plan = delegate; + } // Self-serializing plans handle themselves via the `try_to_proto` hook // (#22419). `Ok(None)` means "not migrated" and falls through to the diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 6ede6fc0e9ae3..1e1620ca8343a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -303,6 +303,35 @@ fn serialize_uses_downcast_delegate() -> Result<()> { Ok(()) } +/// A wrapper delegating to a plan that serializes itself via the +/// `try_to_proto` hook must serialize as its delegate: the wrapper's default +/// hook returns `Ok(None)` and the delegate has no downcast-chain fallback. +#[test] +fn serialize_uses_downcast_delegate_for_self_serializing_plan() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); + let inner: Arc = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: col("a", &schema)?, + alias: "a".to_string(), + }], + input, + )?); + let plan: Arc = Arc::new(DowncastDelegatingExec::new(inner)); + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; + + assert!(matches!( + proto.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Projection( + _ + )) + )); + + Ok(()) +} + #[test] fn roundtrip_date_time_interval() -> Result<()> { let schema = Schema::new(vec![ From 9a5ccc91af37efb36876f625b70953cdf4666d65 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 16 Jul 2026 12:03:48 -0500 Subject: [PATCH 4/5] test: unit-cover the bytes-only function serde and decode helpers The udf/udaf/udwf ctx methods have no in-tree plan caller until the function-carrying plans migrate (follow-ups of #23494), which showed up as missing patch coverage. Pin their contract now: `None` == encode by name, decode prefers an explicit payload, and a registry miss falls back to the codec with an empty buffer. Also cover the required-field decode helpers and the expect_plan_variant! mismatch error. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01S4HUcyoHG8zRpWNrXRLmpg --- datafusion/proto/src/physical_plan/mod.rs | 423 ++++++++++++++++++++++ 1 file changed, 423 insertions(+) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index a66603a385275..da0b7297f5246 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -177,6 +177,429 @@ mod tests { (display.as_str(), None) ); } + + /// Unit tests for the bytes-only function serde exposed on + /// [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] and backed by + /// [`ConverterPlanEncoder`] / [`ConverterPlanDecoder`]. Function-carrying + /// plans migrate in follow-up PRs, so these paths have no in-tree plan + /// caller yet; the tests pin the payload semantics (`None` == encode by + /// name) and the decode lookup order (payload → codec; else registry → + /// codec fallback with an empty buffer) that those migrations rely on. + mod function_serde { + use super::*; + use arrow::datatypes::{DataType, Field, FieldRef}; + use datafusion_common::plan_err; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnv; + use datafusion_expr::function::AccumulatorArgs; + use datafusion_expr::{ + Accumulator, AggregateUDFImpl, ColumnarValue, PartitionEvaluator, + ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, WindowUDFImpl, + }; + use datafusion_functions_window_common::field::WindowUDFFieldArgs; + use datafusion_functions_window_common::partition::PartitionEvaluatorArgs; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdf { + signature: Signature, + } + + impl TestUdf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl ScalarUDFImpl for TestUdf { + fn name(&self) -> &str { + "test_udf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _args: &[DataType]) -> Result { + Ok(DataType::Int64) + } + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> Result { + plan_err!("test only") + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdaf { + signature: Signature, + } + + impl TestUdaf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl AggregateUDFImpl for TestUdaf { + fn name(&self) -> &str { + "test_udaf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result> { + plan_err!("test only") + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdwf { + signature: Signature, + } + + impl TestUdwf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl WindowUDFImpl for TestUdwf { + fn name(&self) -> &str { + "test_udwf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn partition_evaluator( + &self, + _partition_evaluator_args: PartitionEvaluatorArgs, + ) -> Result> { + plan_err!("test only") + } + fn field(&self, field_args: WindowUDFFieldArgs) -> Result { + Ok(Field::new(field_args.name(), DataType::Int64, true).into()) + } + } + + /// Codec that encodes every function as its name bytes and decodes by + /// checking the payload it receives, so tests can observe exactly what + /// crosses the bytes-only boundary. + #[derive(Debug)] + struct PayloadCodec; + + impl PhysicalExtensionCodec for PayloadCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + assert_eq!(name, "test_udf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(ScalarUDF::from(TestUdf::new()))) + } + + fn try_encode_udaf( + &self, + node: &AggregateUDF, + buf: &mut Vec, + ) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udaf( + &self, + name: &str, + buf: &[u8], + ) -> Result> { + assert_eq!(name, "test_udaf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(AggregateUDF::from(TestUdaf::new()))) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + assert_eq!(name, "test_udwf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(WindowUDF::from(TestUdwf::new()))) + } + } + + /// Codec whose decode hooks only accept an empty payload, to pin the + /// by-name decode fallback (registry miss → codec with `&[]`). + #[derive(Debug)] + struct EmptyPayloadOnlyCodec; + + impl PhysicalExtensionCodec for EmptyPayloadOnlyCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_decode_udf(&self, _name: &str, buf: &[u8]) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(ScalarUDF::from(TestUdf::new()))) + } + + fn try_decode_udaf( + &self, + _name: &str, + buf: &[u8], + ) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(AggregateUDF::from(TestUdaf::new()))) + } + + fn try_decode_udwf(&self, _name: &str, buf: &[u8]) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(WindowUDF::from(TestUdwf::new()))) + } + } + + fn encode_ctx_over<'a>( + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, + ) -> ConverterPlanEncoder<'a> { + ConverterPlanEncoder { + codec, + proto_converter, + } + } + + #[test] + fn encode_by_name_functions_produce_no_payload() -> Result<()> { + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let encoder = encode_ctx_over(&codec, &converter); + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert!(ctx.encode_udf(&ScalarUDF::from(TestUdf::new()))?.is_none()); + assert!( + ctx.encode_udaf(&AggregateUDF::from(TestUdaf::new()))? + .is_none() + ); + assert!( + ctx.encode_udwf(&WindowUDF::from(TestUdwf::new()))? + .is_none() + ); + Ok(()) + } + + #[test] + fn encode_functions_surface_codec_payload() -> Result<()> { + let codec = PayloadCodec; + let converter = DefaultPhysicalProtoConverter {}; + let encoder = encode_ctx_over(&codec, &converter); + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert_eq!( + ctx.encode_udf(&ScalarUDF::from(TestUdf::new()))?.as_deref(), + Some(b"test_udf".as_slice()) + ); + assert_eq!( + ctx.encode_udaf(&AggregateUDF::from(TestUdaf::new()))? + .as_deref(), + Some(b"test_udaf".as_slice()) + ); + assert_eq!( + ctx.encode_udwf(&WindowUDF::from(TestUdwf::new()))? + .as_deref(), + Some(b"test_udwf".as_slice()) + ); + Ok(()) + } + + #[test] + fn decode_functions_prefer_explicit_payload() -> Result<()> { + let task_ctx = TaskContext::default(); + let codec = PayloadCodec; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert_eq!( + ctx.decode_udf("test_udf", Some(b"test_udf"))?.name(), + "test_udf" + ); + assert_eq!( + ctx.decode_udaf("test_udaf", Some(b"test_udaf"))?.name(), + "test_udaf" + ); + assert_eq!( + ctx.decode_udwf("test_udwf", Some(b"test_udwf"))?.name(), + "test_udwf" + ); + Ok(()) + } + + #[test] + fn decode_functions_by_name_resolve_from_registry() -> Result<()> { + let udf = Arc::new(ScalarUDF::from(TestUdf::new())); + let udaf = Arc::new(AggregateUDF::from(TestUdaf::new())); + let udwf = Arc::new(WindowUDF::from(TestUdwf::new())); + let task_ctx = TaskContext::new( + None, + "test".to_string(), + SessionConfig::new(), + HashMap::from([("test_udf".to_string(), Arc::clone(&udf))]), + HashMap::new(), + HashMap::from([("test_udaf".to_string(), Arc::clone(&udaf))]), + HashMap::from([("test_udwf".to_string(), Arc::clone(&udwf))]), + Arc::new(RuntimeEnv::default()), + ); + // The default codec fails any decode, so a success proves the + // registry satisfied the lookup without a codec fallback. + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert!(Arc::ptr_eq(&ctx.decode_udf("test_udf", None)?, &udf)); + assert!(Arc::ptr_eq(&ctx.decode_udaf("test_udaf", None)?, &udaf)); + assert!(Arc::ptr_eq(&ctx.decode_udwf("test_udwf", None)?, &udwf)); + assert_eq!(ctx.task_ctx().session_id(), "test"); + Ok(()) + } + + #[test] + fn decode_functions_by_name_fall_back_to_codec_on_registry_miss() -> Result<()> { + let task_ctx = TaskContext::default(); + let codec = EmptyPayloadOnlyCodec; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert_eq!(ctx.decode_udf("test_udf", None)?.name(), "test_udf"); + assert_eq!(ctx.decode_udaf("test_udaf", None)?.name(), "test_udaf"); + assert_eq!(ctx.decode_udwf("test_udwf", None)?.name(), "test_udwf"); + Ok(()) + } + + #[test] + fn decode_required_helpers_error_on_missing_fields() { + let task_ctx = TaskContext::default(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + let err = ctx + .decode_required_child(None, "FooExec", "input") + .unwrap_err(); + assert!( + err.to_string() + .contains("FooExec is missing required field 'input'"), + "unexpected error: {err}" + ); + + let schema = Schema::empty(); + let err = ctx + .decode_required_expr(None, &schema, "FooExec", "predicate") + .unwrap_err(); + assert!( + err.to_string() + .contains("FooExec is missing required field 'predicate'"), + "unexpected error: {err}" + ); + } + + #[test] + fn try_from_proto_rejects_wrong_plan_variant() { + let task_ctx = TaskContext::default(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + let node = protobuf::PhysicalPlanNode { + physical_plan_type: None, + }; + let err = ProjectionExec::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalPlanNode is not a ProjectionExec"), + "unexpected error: {err}" + ); + } + } } /// Context threaded through physical-plan deserialization. From 77dded51e3e77625bac484d441f7f0f394a237e3 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 16 Jul 2026 12:20:34 -0500 Subject: [PATCH 5/5] Restore projection PhysicalPlanNodeExt methods as deprecated Per the API health policy, keep try_into_projection_physical_plan / try_from_projection_exec with their original bodies under #[deprecated(since = "55.0.0")] instead of removing them outright, now that the hook path replaces them. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01S4HUcyoHG8zRpWNrXRLmpg --- datafusion/proto/src/physical_plan/mod.rs | 73 ++++++++++++++++++++++- 1 file changed, 72 insertions(+), 1 deletion(-) diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index da0b7297f5246..4c6c270994f38 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -85,7 +85,7 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; use datafusion_physical_plan::metrics::MetricCategory; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::proto::{ ExecutionPlanDecode, ExecutionPlanDecodeCtx, ExecutionPlanEncode, ExecutionPlanEncodeCtx, @@ -1178,6 +1178,40 @@ pub trait PhysicalPlanNodeExt: Sized { ))) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `ProjectionExec` deserializes itself via `ProjectionExec::try_from_proto`" + )] + fn try_into_projection_physical_plan( + &self, + projection: &protobuf::ProjectionExecNode, + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let input: Arc = + into_physical_plan(&projection.input, ctx, proto_converter)?; + let exprs = projection + .expr + .iter() + .zip(projection.expr_name.iter()) + .map(|(expr, name)| { + Ok(( + proto_converter.proto_to_physical_expr( + expr, + input.schema().as_ref(), + ctx, + )?, + name.to_string(), + )) + }) + .collect::, String)>>>()?; + let proj_exprs: Vec = exprs + .into_iter() + .map(|(expr, alias)| ProjectionExpr { expr, alias }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) + } + fn try_into_filter_physical_plan( &self, filter: &protobuf::FilterExecNode, @@ -2847,6 +2881,43 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `ProjectionExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] + fn try_from_projection_exec( + exec: &ProjectionExec, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result { + let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + exec.input().to_owned(), + codec, + proto_converter, + )?; + let expr = exec + .expr() + .iter() + .map(|proj_expr| { + proto_converter.physical_expr_to_proto(&proj_expr.expr, codec) + }) + .collect::>>()?; + let expr_name = exec + .expr() + .iter() + .map(|proj_expr| proj_expr.alias.clone()) + .collect(); + Ok(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Projection(Box::new( + protobuf::ProjectionExecNode { + input: Some(Box::new(input)), + expr, + expr_name, + }, + ))), + }) + } + fn try_from_analyze_exec( exec: &AnalyzeExec, codec: &dyn PhysicalExtensionCodec,