From 015428d32e1ecc405aadb14d98a60c22328a9a7c 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/2] 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 76abf73e0ebbe..13c413d174c19 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -757,6 +757,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 6cc6e44c32cc3..0658b4df9fcb7 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -84,6 +84,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 16b0a5ad7e4b5..39444132565ec 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -494,6 +494,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 72f6e5af5bff2..5948127402b9c 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, @@ -4369,3 +4322,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 5ca0d6376154879f1bf3eae012dcd85d1407833b 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/2] 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 5948127402b9c..96e81c3caeb49 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -4347,6 +4347,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 @@ -4378,4 +4398,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, &[])), + } + } }