diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index e58acee4ce6bd..3b9d5d258a838 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -823,6 +823,27 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { ) -> Option> { None } + + /// Serialize this plan to its protobuf representation, if it knows how. + /// + /// This is the `ExecutionPlan` analog of + /// [`PhysicalExpr::try_to_proto`]. + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller + /// (`datafusion-proto`) falls back to the central downcast chain. Every + /// un-migrated plan keeps its existing behavior. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// * `Err(_)` — a real failure (e.g. a child failed to serialize). + /// + /// Only *self-contained* plans should override this — see [`crate::proto`] + /// for the session-dependency boundary. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn ExecutionPlan { diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 0ab232bc102bc..8cba650b79770 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -88,6 +88,8 @@ pub mod metrics; pub mod operator_statistics; pub mod placeholder_row; pub mod projection; +#[cfg(feature = "proto")] +pub mod proto; pub mod recursive_query; pub mod repartition; pub mod scalar_subquery; diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index d55363297bd52..42501f22395b4 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -500,6 +500,71 @@ impl ExecutionPlan for ProjectionExec { .ok() }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expressions(self.expr().iter().map(|p| &p.expr))?; + let expr_name = self.expr().iter().map(|p| p.alias.clone()).collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Projection(Box::new( + protobuf::ProjectionExecNode { + input: Some(Box::new(input)), + expr, + expr_name, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ProjectionExec { + /// Reconstruct a [`ProjectionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. Child plans and expressions are decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let projection = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Projection, + "ProjectionExec", + ); + let input = ctx.decode_required_child( + projection.input.as_deref(), + "ProjectionExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = projection + .expr + .iter() + .zip(projection.expr_name.iter()) + .map(|(expr, name)| { + Ok(ProjectionExpr { + expr: ctx.decode_expr(expr, input_schema.as_ref())?, + alias: name.to_string(), + }) + }) + .collect::>>()?; + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) + } } impl ProjectionStream { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs new file mode 100644 index 0000000000000..1731203f6c767 --- /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` 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) -> Result; + + /// 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`]. +/// +/// 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; + + /// 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`]. +/// +/// +/// 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() + } + + /// 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. +/// +/// 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. + 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` +/// variant, returning a reference to the inner payload, else an `internal_err!`. +/// Mirrors `expect_expr_variant!` on the expression side. Field access on the +/// result auto-derefs through the `Box` that boxed variants use. +#[macro_export] +macro_rules! expect_plan_variant { + ($node:expr, $variant:path, $plan_name:literal $(,)?) => {{ + match &$node.physical_plan_type { + Some($variant(inner)) => inner, + _ => { + return ::datafusion_common::internal_err!(concat!( + "PhysicalPlanNode is not a ", + $plan_name + )); + } + } + }}; +} diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 0744a94dcebd1..4c6c270994f38 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 { + Ok(DataType::Int64) + } + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> Result { + plan_err!("test only") + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdaf { + signature: Signature, + } + + impl TestUdaf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl AggregateUDFImpl for TestUdaf { + fn name(&self) -> &str { + "test_udaf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result> { + plan_err!("test only") + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdwf { + signature: Signature, + } + + impl TestUdwf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl WindowUDFImpl for TestUdwf { + fn name(&self) -> &str { + "test_udwf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn partition_evaluator( + &self, + _partition_evaluator_args: PartitionEvaluatorArgs, + ) -> Result> { + plan_err!("test only") + } + fn field(&self, field_args: WindowUDFFieldArgs) -> Result { + Ok(Field::new(field_args.name(), DataType::Int64, true).into()) + } + } + + /// Codec that encodes every function as its name bytes and decodes by + /// checking the payload it receives, so tests can observe exactly what + /// crosses the bytes-only boundary. + #[derive(Debug)] + struct PayloadCodec; + + impl PhysicalExtensionCodec for PayloadCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + assert_eq!(name, "test_udf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(ScalarUDF::from(TestUdf::new()))) + } + + fn try_encode_udaf( + &self, + node: &AggregateUDF, + buf: &mut Vec, + ) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udaf( + &self, + name: &str, + buf: &[u8], + ) -> Result> { + assert_eq!(name, "test_udaf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(AggregateUDF::from(TestUdaf::new()))) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + assert_eq!(name, "test_udwf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(WindowUDF::from(TestUdwf::new()))) + } + } + + /// Codec whose decode hooks only accept an empty payload, to pin the + /// by-name decode fallback (registry miss → codec with `&[]`). + #[derive(Debug)] + struct EmptyPayloadOnlyCodec; + + impl PhysicalExtensionCodec for EmptyPayloadOnlyCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_decode_udf(&self, _name: &str, buf: &[u8]) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(ScalarUDF::from(TestUdf::new()))) + } + + fn try_decode_udaf( + &self, + _name: &str, + buf: &[u8], + ) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(AggregateUDF::from(TestUdaf::new()))) + } + + fn try_decode_udwf(&self, _name: &str, buf: &[u8]) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(WindowUDF::from(TestUdwf::new()))) + } + } + + fn encode_ctx_over<'a>( + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, + ) -> ConverterPlanEncoder<'a> { + ConverterPlanEncoder { + codec, + proto_converter, + } + } + + #[test] + fn encode_by_name_functions_produce_no_payload() -> Result<()> { + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let encoder = encode_ctx_over(&codec, &converter); + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert!(ctx.encode_udf(&ScalarUDF::from(TestUdf::new()))?.is_none()); + assert!( + ctx.encode_udaf(&AggregateUDF::from(TestUdaf::new()))? + .is_none() + ); + assert!( + ctx.encode_udwf(&WindowUDF::from(TestUdwf::new()))? + .is_none() + ); + Ok(()) + } + + #[test] + fn encode_functions_surface_codec_payload() -> Result<()> { + let codec = PayloadCodec; + let converter = DefaultPhysicalProtoConverter {}; + let encoder = encode_ctx_over(&codec, &converter); + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert_eq!( + ctx.encode_udf(&ScalarUDF::from(TestUdf::new()))?.as_deref(), + Some(b"test_udf".as_slice()) + ); + assert_eq!( + ctx.encode_udaf(&AggregateUDF::from(TestUdaf::new()))? + .as_deref(), + Some(b"test_udaf".as_slice()) + ); + assert_eq!( + ctx.encode_udwf(&WindowUDF::from(TestUdwf::new()))? + .as_deref(), + Some(b"test_udwf".as_slice()) + ); + Ok(()) + } + + #[test] + fn decode_functions_prefer_explicit_payload() -> Result<()> { + let task_ctx = TaskContext::default(); + let codec = PayloadCodec; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert_eq!( + ctx.decode_udf("test_udf", Some(b"test_udf"))?.name(), + "test_udf" + ); + assert_eq!( + ctx.decode_udaf("test_udaf", Some(b"test_udaf"))?.name(), + "test_udaf" + ); + assert_eq!( + ctx.decode_udwf("test_udwf", Some(b"test_udwf"))?.name(), + "test_udwf" + ); + Ok(()) + } + + #[test] + fn decode_functions_by_name_resolve_from_registry() -> Result<()> { + let udf = Arc::new(ScalarUDF::from(TestUdf::new())); + let udaf = Arc::new(AggregateUDF::from(TestUdaf::new())); + let udwf = Arc::new(WindowUDF::from(TestUdwf::new())); + let task_ctx = TaskContext::new( + None, + "test".to_string(), + SessionConfig::new(), + HashMap::from([("test_udf".to_string(), Arc::clone(&udf))]), + HashMap::new(), + HashMap::from([("test_udaf".to_string(), Arc::clone(&udaf))]), + HashMap::from([("test_udwf".to_string(), Arc::clone(&udwf))]), + Arc::new(RuntimeEnv::default()), + ); + // The default codec fails any decode, so a success proves the + // registry satisfied the lookup without a codec fallback. + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert!(Arc::ptr_eq(&ctx.decode_udf("test_udf", None)?, &udf)); + assert!(Arc::ptr_eq(&ctx.decode_udaf("test_udaf", None)?, &udaf)); + assert!(Arc::ptr_eq(&ctx.decode_udwf("test_udwf", None)?, &udwf)); + assert_eq!(ctx.task_ctx().session_id(), "test"); + Ok(()) + } + + #[test] + fn decode_functions_by_name_fall_back_to_codec_on_registry_miss() -> Result<()> { + let task_ctx = TaskContext::default(); + let codec = EmptyPayloadOnlyCodec; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert_eq!(ctx.decode_udf("test_udf", None)?.name(), "test_udf"); + assert_eq!(ctx.decode_udaf("test_udaf", None)?.name(), "test_udaf"); + assert_eq!(ctx.decode_udwf("test_udwf", None)?.name(), "test_udwf"); + Ok(()) + } + + #[test] + fn decode_required_helpers_error_on_missing_fields() { + let task_ctx = TaskContext::default(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + let err = ctx + .decode_required_child(None, "FooExec", "input") + .unwrap_err(); + assert!( + err.to_string() + .contains("FooExec is missing required field 'input'"), + "unexpected error: {err}" + ); + + let schema = Schema::empty(); + let err = ctx + .decode_required_expr(None, &schema, "FooExec", "predicate") + .unwrap_err(); + assert!( + err.to_string() + .contains("FooExec is missing required field 'predicate'"), + "unexpected error: {err}" + ); + } + + #[test] + fn try_from_proto_rejects_wrong_plan_variant() { + let task_ctx = TaskContext::default(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + let node = protobuf::PhysicalPlanNode { + physical_plan_type: None, + }; + let err = ProjectionExec::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalPlanNode is not a ProjectionExec"), + "unexpected error: {err}" + ); + } + } } /// Context threaded through physical-plan deserialization. @@ -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 { 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::() { - 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,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, @@ -2430,6 +2881,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, @@ -4387,3 +4842,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, + ) -> 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) + } + + // 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 +/// `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() + } + + // 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, &[])), + } + } +} diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 6ede6fc0e9ae3..1e1620ca8343a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -303,6 +303,35 @@ fn serialize_uses_downcast_delegate() -> Result<()> { Ok(()) } +/// A wrapper delegating to a plan that serializes itself via the +/// `try_to_proto` hook must serialize as its delegate: the wrapper's default +/// hook returns `Ok(None)` and the delegate has no downcast-chain fallback. +#[test] +fn serialize_uses_downcast_delegate_for_self_serializing_plan() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); + let inner: Arc = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: col("a", &schema)?, + alias: "a".to_string(), + }], + input, + )?); + let plan: Arc = Arc::new(DowncastDelegatingExec::new(inner)); + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; + + assert!(matches!( + proto.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Projection( + _ + )) + )); + + Ok(()) +} + #[test] fn roundtrip_date_time_interval() -> Result<()> { let schema = Schema::new(vec![