You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/07/17 15:36:42 UTC
(datafusion) branch gh-readonly-queue/main/pr-23495-d4160119db8c6b348928af100393c23783caf05e created (now b16bd3d874)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-23495-d4160119db8c6b348928af100393c23783caf05e
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at b16bd3d874 Add ExecutionPlan try_to_proto / try_from_proto hooks + ProjectionExec reference (#23495)
This branch includes the following new commits:
new b16bd3d874 Add ExecutionPlan try_to_proto / try_from_proto hooks + ProjectionExec reference (#23495)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Add ExecutionPlan try_to_proto / try_from_proto hooks + ProjectionExec reference (#23495)
Posted by gi...@apache.org.
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-23495-d4160119db8c6b348928af100393c23783caf05e
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b16bd3d874a9e11877d7aa4e59e8845ff576c430
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Fri Jul 17 10:36:20 2026 -0500
Add ExecutionPlan try_to_proto / try_from_proto hooks + ProjectionExec reference (#23495)
## Which issue does this PR close?
- Part of #23494.
## Rationale for this change
This is the foundation for porting `ExecutionPlan` serialization off the
central `downcast_ref` chain in
`datafusion/proto/src/physical_plan/mod.rs`, mirroring what #21929 did
for `PhysicalExpr` (tracked in #22418). Everything else in #23494
depends on this PR; it can't be parallelized because it defines the ctx
API and wiring that every per-plan follow-up uses.
## What changes are included in this PR?
In `datafusion-physical-plan` (feature-gated `#[cfg(feature =
"proto")]`):
- `ExecutionPlan::try_to_proto(&self, ctx: &ExecutionPlanEncodeCtx) ->
Result<Option<PhysicalPlanNode>>` — a new trait method defaulting to
`Ok(None)` (so un-migrated plans are unchanged).
- `ExecutionPlanEncodeCtx` / `ExecutionPlanDecodeCtx` plus the internal
`ExecutionPlanEncode` / `ExecutionPlanDecode` dispatch traits (defined
here, implemented in `datafusion-proto`) — the same dependency-inversion
the `PhysicalExpr` ctx uses. The ctx exposes
`encode_child`/`decode_child` (+ `_children`),
`encode_expr`/`decode_expr`, and typed **bytes-only**
`encode_udaf`/`decode_udaf` (+ udf/udwf) for function-carrying plans. No
proto type is ever named by physical-plan beyond the pure prost types in
`datafusion-proto-models`.
- An `expect_plan_variant!` macro mirroring `expect_expr_variant!`.
In `datafusion-proto`:
- `ConverterPlanEncoder` / `ConverterPlanDecoder` implement the dispatch
traits over the existing `PhysicalExtensionCodec` + converter (the
function serde reuses today's `fun_definition` byte semantics exactly).
The encode dispatch resolves `downcast_delegate()` first (so wrapper
plans serialize as their delegate, exactly like the `downcast_ref` chain
sees them), then calls the hook, then falls back to the existing
downcast chain.
`ProjectionExec` is migrated as the reference implementation: its old
encode arm is deleted and its decode arm reduced to
`ProjectionExec::try_from_proto`, so a green roundtrip proves the hook
is the only path.
## Are these changes tested?
Yes — the existing `proto_integration` roundtrip suite (including TPC-H)
passes unchanged, which exercises the `ProjectionExec` migration
end-to-end, plus a new test asserting that a wrapper plan delegating its
downcast identity to `ProjectionExec` still serializes as a projection
through the hook (it fails without the `downcast_delegate()`
resolution). The bytes-only function serde and the required-field decode
helpers — which have no in-tree plan caller until the function-carrying
plans migrate — are pinned by dedicated unit tests (payload semantics,
decode lookup order, error paths), addressing the patch-coverage gap.
`cargo clippy --all-targets -- -D warnings`, `cargo fmt`, and
`RUSTDOCFLAGS="-D warnings" cargo doc` are clean.
## Are there any user-facing changes?
New public API: the `try_to_proto` trait method (with a default, so
existing `ExecutionPlan` impls are unaffected) and the ctx types. No
wire-format change.
No API is removed: per the [API health
policy](https://datafusion.apache.org/contributor-guide/api-health.html),
the two per-plan `PhysicalPlanNodeExt` methods this migration replaces
(`try_into_projection_physical_plan` / `try_from_projection_exec`) are
kept with their original bodies under `#[deprecated(since = "55.0.0")]`.
(They were also never released: `PhysicalPlanNodeExt` was introduced by
#21929 after 54.0.0 shipped.) The per-plan follow-ups in #23494 will
deprecate their corresponding methods the same way, with removal after
the policy window.
---------
Co-authored-by: Claude Opus 4.8 (1M context) <no...@anthropic.com>
---
datafusion/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 | 308 +++++++++++
datafusion/proto/src/physical_plan/mod.rs | 594 ++++++++++++++++++++-
.../proto/tests/cases/roundtrip_physical_plan.rs | 29 +
6 files changed, 1008 insertions(+), 11 deletions(-)
diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs
index e58acee4ce..3b9d5d258a 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<Arc<dyn ExecutionPlan>> {
None
}
+
+ /// Serialize this plan to its protobuf representation, if it knows how.
+ ///
+ /// This is the `ExecutionPlan` analog of
+ /// [`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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ Ok(None)
+ }
}
impl dyn ExecutionPlan {
diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs
index 0ab232bc10..8cba650b79 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 d55363297b..42501f2239 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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+ 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<Arc<dyn ExecutionPlan>> {
+ 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::<Result<Vec<ProjectionExpr>>>()?;
+ 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 0000000000..1731203f6c
--- /dev/null
+++ b/datafusion/physical-plan/src/proto.rs
@@ -0,0 +1,308 @@
+// 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`.
+//!
+//! # Function-carrying plans
+//!
+//! 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<u8>` 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.
+//!
+//! [`ExecutionPlan`]: crate::ExecutionPlan
+
+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};
+
+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<dyn ExecutionPlan>) -> Result<PhysicalPlanNode>;
+
+ /// Serialize a physical expression owned by the plan.
+ fn encode_expr(&self, expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode>;
+
+ /// 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<Option<Vec<u8>>>;
+
+ /// Serialize an aggregate UDF to an opaque payload. `None` means "decodable
+ /// by name alone".
+ fn encode_udaf(&self, udaf: &AggregateUDF) -> Result<Option<Vec<u8>>>;
+
+ /// Serialize a window UDF to an opaque payload. `None` means "decodable by
+ /// name alone".
+ fn encode_udwf(&self, udwf: &WindowUDF) -> Result<Option<Vec<u8>>>;
+}
+
+/// 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<Arc<dyn ExecutionPlan>>;
+
+ /// Deserialize a physical expression against `input_schema`.
+ fn decode_expr(
+ &self,
+ node: &PhysicalExprNode,
+ input_schema: &Schema,
+ ) -> Result<Arc<dyn PhysicalExpr>>;
+
+ /// 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<Arc<ScalarUDF>>;
+
+ /// Reconstruct an aggregate UDF from its name and optional payload.
+ fn decode_udaf(
+ &self,
+ name: &str,
+ payload: Option<&[u8]>,
+ ) -> Result<Arc<AggregateUDF>>;
+
+ /// Reconstruct a window UDF from its name and optional payload.
+ fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result<Arc<WindowUDF>>;
+}
+
+/// Context handed to [`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<dyn ExecutionPlan>,
+ ) -> Result<PhysicalPlanNode> {
+ self.encoder.encode_plan(plan)
+ }
+
+ /// Serialize an iterator of child plans.
+ pub fn encode_children<'b, I>(&self, plans: I) -> Result<Vec<PhysicalPlanNode>>
+ where
+ I: IntoIterator<Item = &'b Arc<dyn ExecutionPlan>>,
+ {
+ plans.into_iter().map(|p| self.encode_child(p)).collect()
+ }
+
+ /// Serialize a single physical expression.
+ pub fn encode_expr(&self, expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> {
+ self.encoder.encode_expr(expr)
+ }
+
+ /// Serialize an iterator of physical expressions.
+ pub fn encode_expressions<'b, I>(&self, exprs: I) -> Result<Vec<PhysicalExprNode>>
+ where
+ I: IntoIterator<Item = &'b Arc<dyn PhysicalExpr>>,
+ {
+ 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<Option<Vec<u8>>> {
+ 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<Option<Vec<u8>>> {
+ 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<Option<Vec<u8>>> {
+ self.encoder.encode_udwf(udwf)
+ }
+}
+
+/// 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<Arc<dyn ExecutionPlan>> {
+ 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<Arc<dyn ExecutionPlan>> {
+ 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<Arc<dyn PhysicalExpr>> {
+ 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<Arc<dyn PhysicalExpr>> {
+ 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.
+ 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<Arc<ScalarUDF>> {
+ 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<Arc<AggregateUDF>> {
+ 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<Arc<WindowUDF>> {
+ self.decoder.decode_udwf(name, payload)
+ }
+}
+
+/// 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 f3d14ec53f..e5f8aa072d 100644
--- a/datafusion/proto/src/physical_plan/mod.rs
+++ b/datafusion/proto/src/physical_plan/mod.rs
@@ -86,6 +86,10 @@ 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::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;
@@ -173,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<DataType> {
+ Ok(DataType::Int64)
+ }
+ fn invoke_with_args(
+ &self,
+ _args: ScalarFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ 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<DataType> {
+ Ok(DataType::Int64)
+ }
+ fn accumulator(
+ &self,
+ _acc_args: AccumulatorArgs,
+ ) -> Result<Box<dyn Accumulator>> {
+ 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<Box<dyn PartitionEvaluator>> {
+ plan_err!("test only")
+ }
+ fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+ 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<dyn ExecutionPlan>],
+ _ctx: &TaskContext,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ internal_err!("not needed for these tests")
+ }
+
+ fn try_encode(
+ &self,
+ _node: Arc<dyn ExecutionPlan>,
+ _buf: &mut Vec<u8>,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<()> {
+ internal_err!("not needed for these tests")
+ }
+
+ fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) -> Result<()> {
+ buf.extend_from_slice(node.name().as_bytes());
+ Ok(())
+ }
+
+ fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+ 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<u8>,
+ ) -> Result<()> {
+ buf.extend_from_slice(node.name().as_bytes());
+ Ok(())
+ }
+
+ fn try_decode_udaf(
+ &self,
+ name: &str,
+ buf: &[u8],
+ ) -> Result<Arc<AggregateUDF>> {
+ 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<u8>) -> Result<()> {
+ buf.extend_from_slice(node.name().as_bytes());
+ Ok(())
+ }
+
+ fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> {
+ 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<dyn ExecutionPlan>],
+ _ctx: &TaskContext,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ internal_err!("not needed for these tests")
+ }
+
+ fn try_encode(
+ &self,
+ _node: Arc<dyn ExecutionPlan>,
+ _buf: &mut Vec<u8>,
+ _proto_converter: &dyn PhysicalProtoConverterExtension,
+ ) -> Result<()> {
+ internal_err!("not needed for these tests")
+ }
+
+ fn try_decode_udf(&self, _name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+ assert!(buf.is_empty());
+ Ok(Arc::new(ScalarUDF::from(TestUdf::new())))
+ }
+
+ fn try_decode_udaf(
+ &self,
+ _name: &str,
+ buf: &[u8],
+ ) -> Result<Arc<AggregateUDF>> {
+ assert!(buf.is_empty());
+ Ok(Arc::new(AggregateUDF::from(TestUdaf::new())))
+ }
+
+ fn try_decode_udwf(&self, _name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> {
+ 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.
@@ -312,12 +739,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)
@@ -442,18 +877,30 @@ pub trait PhysicalPlanNodeExt: Sized {
proto_converter: &dyn PhysicalProtoConverterExtension,
) -> Result<protobuf::PhysicalPlanNode> {
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;
+ }
- if let Some(exec) = plan.downcast_ref::<ExplainExec>() {
- 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::<ProjectionExec>() {
- return protobuf::PhysicalPlanNode::try_from_projection_exec(
- exec,
- codec,
- proto_converter,
- );
+ if let Some(exec) = plan.downcast_ref::<ExplainExec>() {
+ return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec);
}
if let Some(exec) = plan.downcast_ref::<AnalyzeExec>() {
@@ -731,6 +1178,10 @@ 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,
@@ -2438,6 +2889,10 @@ 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,
@@ -4399,3 +4854,120 @@ 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<dyn ExecutionPlan>,
+ ) -> Result<protobuf::PhysicalPlanNode> {
+ self.proto_converter
+ .execution_plan_to_proto(plan, self.codec)
+ }
+
+ fn encode_expr(
+ &self,
+ expr: &Arc<dyn PhysicalExpr>,
+ ) -> Result<protobuf::PhysicalExprNode> {
+ 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<Option<Vec<u8>>> {
+ 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<Option<Vec<u8>>> {
+ 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<Option<Vec<u8>>> {
+ 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
+/// `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<Arc<dyn ExecutionPlan>> {
+ self.proto_converter.proto_to_execution_plan(node, self.ctx)
+ }
+
+ fn decode_expr(
+ &self,
+ node: &protobuf::PhysicalExprNode,
+ input_schema: &Schema,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ self.proto_converter
+ .proto_to_physical_expr(node, input_schema, self.ctx)
+ }
+
+ 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<Arc<ScalarUDF>> {
+ 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<Arc<AggregateUDF>> {
+ 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<Arc<WindowUDF>> {
+ 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, &[])),
+ }
+ }
+}
diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
index 1dc4580302..cbd6fd912a 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -355,6 +355,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<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
+ vec![ProjectionExpr {
+ expr: col("a", &schema)?,
+ alias: "a".to_string(),
+ }],
+ input,
+ )?);
+ let plan: Arc<dyn ExecutionPlan> = 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![
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org