This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new b16bd3d874 Add ExecutionPlan try_to_proto / try_from_proto hooks +
ProjectionExec reference (#23495)
b16bd3d874 is described below
commit b16bd3d874a9e11877d7aa4e59e8845ff576c430
Author: Adrian Garcia Badaracco <[email protected]>
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) <[email protected]>
---
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: [email protected]
For additional commands, e-mail: [email protected]
