From 4edb0053e80f38c380bebef2eec7deb661783503 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 10 Jul 2026 12:27:56 -0500 Subject: [PATCH 1/2] feat: typed erasure-free proto bridge for ScalarFunctionExpr `ScalarFunctionExpr` is the one built-in expression whose ser/de needs session-level objects: the embedded `ScalarUDF` round-trips through the `datafusion-proto` extension codec + function registry, and reconstruction needs the session `ConfigOptions`. Those types live *above* `physical-expr-common` in the crate graph, so they cannot appear on the crate-wide `PhysicalExpr::try_to_proto` encode/decode contexts without a dependency cycle. Rather than erase the function type to `dyn Any` on the crate-wide context, the expression declares its own ser/de against two fully-typed bridge traits (`ScalarUdfProtoEncoder` / `ScalarUdfProtoDecoder`) implemented in `datafusion-proto`, where `ScalarUDF` and the codec are both nameable. No `dyn Any` is involved anywhere. Adds the traits, the inherent `try_to_proto` / `try_from_proto` on `ScalarFunctionExpr`, and 6 direct tests over mock bridges. Wiring in `datafusion-proto` follows in the next commit. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_018LmymuRdTRmGthjMhtCijn --- datafusion/physical-expr/src/lib.rs | 2 + .../physical-expr/src/scalar_function.rs | 379 ++++++++++++++++++ 2 files changed, 381 insertions(+) diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 67419944cfde6..21f4ef39484aa 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -80,6 +80,8 @@ pub use datafusion_physical_expr_common::sort_expr::{ pub use higher_order_function::HigherOrderFunctionExpr; pub use planner::{create_physical_expr, create_physical_exprs}; pub use scalar_function::ScalarFunctionExpr; +#[cfg(feature = "proto")] +pub use scalar_function::{ScalarUdfProtoDecoder, ScalarUdfProtoEncoder}; pub use simplifier::PhysicalExprSimplifier; pub use utils::{conjunction, conjunction_opt, split_conjunction}; diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 418d005c971ea..49d9727b8e543 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -351,6 +351,385 @@ impl PhysicalExpr for ScalarFunctionExpr { } } +/// Typed protobuf serialization for [`ScalarFunctionExpr`]. +/// +/// Most built-in expressions serialize themselves through the crate-wide +/// [`PhysicalExpr::try_to_proto`] hook, whose encode/decode contexts live in +/// `physical-expr-common`. `ScalarFunctionExpr` cannot: its ser/de needs +/// session-level objects — the embedded [`ScalarUDF`] round-trips through +/// `datafusion-proto`'s extension codec + function registry, and +/// reconstruction needs the session [`ConfigOptions`]. Those types are +/// defined *above* `physical-expr-common` in the crate graph +/// (`datafusion-expr` depends on `physical-expr-common`, not the other way +/// round), so they cannot appear on the crate-wide contexts without a +/// dependency cycle. +/// +/// The expression still declares its own ser/de here, but against the +/// **fully-typed** [`ScalarUdfProtoEncoder`] / [`ScalarUdfProtoDecoder`] +/// bridge traits — implemented in `datafusion-proto`, which can name +/// [`ScalarUDF`] and the codec directly. Dispatch stays a downcast match in +/// `datafusion-proto` (symmetric with the decode side), and no type erasure +/// (`dyn Any`) is involved anywhere. The trade versus routing through the +/// hook is one surviving downcast arm in `to_proto.rs` in exchange for a +/// fully-typed, erasure-free surface. See #22430 / the epic in #22418. +#[cfg(feature = "proto")] +mod proto { + use super::*; + use arrow::datatypes::Field; + use datafusion_common::internal_datafusion_err; + use datafusion_proto_models::protobuf; + + /// Encode-side serialization bridge for [`ScalarFunctionExpr`]. + /// + /// Implemented by `datafusion-proto`, where the extension codec and + /// [`ScalarUDF`] are both nameable. Keeps this crate free of a + /// `datafusion-proto` dependency while staying fully typed. + pub trait ScalarUdfProtoEncoder { + /// Encode a child expression to a protobuf node (dedup-aware, via the + /// same converter used for the rest of the plan). + fn encode_child( + &self, + expr: &Arc, + ) -> Result; + + /// Encode a [`ScalarUDF`]'s custom payload via the extension codec. + /// Returns `None` when the codec writes no bytes (the function is then + /// resolved by name alone on decode). + fn encode_udf(&self, udf: &ScalarUDF) -> Result>>; + } + + /// Decode-side counterpart to [`ScalarUdfProtoEncoder`]. + pub trait ScalarUdfProtoDecoder { + /// Decode a child expression node, recursing through the converter. + fn decode_child( + &self, + node: &protobuf::PhysicalExprNode, + ) -> Result>; + + /// Resolve a [`ScalarUDF`] by name, using the optional custom payload. + /// Implementors reproduce the registry-then-codec lookup order. + fn decode_udf( + &self, + name: &str, + fun_definition: Option<&[u8]>, + ) -> Result>; + + /// Session configuration used to reconstruct the expression. + fn config_options(&self) -> Arc; + } + + impl ScalarFunctionExpr { + /// Serialize this expression to a [`protobuf::PhysicalExprNode`] via a + /// typed encoder. + /// + /// Called from `datafusion-proto`'s downcast dispatch rather than the + /// crate-wide `try_to_proto` hook — see the [module docs](self) for + /// why. The wire shape is identical to the pre-existing inline arm. + pub fn try_to_proto( + &self, + encoder: &dyn ScalarUdfProtoEncoder, + ) -> Result { + let args = self + .args() + .iter() + .map(|arg| encoder.encode_child(arg)) + .collect::>>()?; + Ok(protobuf::PhysicalExprNode { + // `ScalarFunctionExpr` never overrides `expression_id`, so this + // is always `None` (matching the old inline arm). + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarUdf( + protobuf::PhysicalScalarUdfNode { + name: self.name().to_string(), + args, + fun_definition: encoder.encode_udf(self.fun())?, + return_type: Some(self.return_type().try_into()?), + nullable: self.nullable(), + return_field_name: self + .return_field(&Schema::empty())? + .name() + .to_string(), + }, + )), + }) + } + + /// Reconstruct a [`ScalarFunctionExpr`] from its protobuf + /// representation via a typed decoder. The inverse of + /// [`Self::try_to_proto`]. + pub fn try_from_proto( + node: &protobuf::PhysicalScalarUdfNode, + decoder: &dyn ScalarUdfProtoDecoder, + ) -> Result> { + let udf = decoder.decode_udf(&node.name, node.fun_definition.as_deref())?; + let args = node + .args + .iter() + .map(|arg| decoder.decode_child(arg)) + .collect::>>()?; + let return_type = node.return_type.as_ref().ok_or_else(|| { + internal_datafusion_err!( + "ScalarFunctionExpr is missing required field 'return_type'" + ) + })?; + Ok(Arc::new( + ScalarFunctionExpr::new( + node.name.as_str(), + udf, + args, + Field::new(&node.return_field_name, return_type.try_into()?, true) + .into(), + decoder.config_options(), + ) + .with_nullable(node.nullable), + )) + } + } +} + +#[cfg(feature = "proto")] +pub use proto::{ScalarUdfProtoDecoder, ScalarUdfProtoEncoder}; + +/// Direct tests for the typed `try_to_proto` / `try_from_proto` bridge. +/// +/// These drive the inherent methods against tiny mock encoder/decoder +/// implementations, so they cover the marshalling and error paths without a +/// running session. End-to-end round-trips (through `datafusion-proto`'s real +/// codec) are covered by `datafusion-proto`'s `proto_integration` suite. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::proto::{ScalarUdfProtoDecoder, ScalarUdfProtoEncoder}; + use super::*; + use crate::expressions::Column; + use arrow::datatypes::Field; + use datafusion_expr::{ScalarUDFImpl, Signature, Volatility}; + use datafusion_proto_models::protobuf; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdf { + signature: Signature, + } + + impl TestUdf { + fn new() -> Self { + Self { + signature: Signature::variadic_any(Volatility::Immutable), + } + } + } + + impl ScalarUDFImpl for TestUdf { + fn name(&self) -> &str { + "test_udf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int32) + } + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(1)))) + } + } + + fn test_udf() -> Arc { + Arc::new(ScalarUDF::from(TestUdf::new())) + } + + fn sample_expr() -> ScalarFunctionExpr { + let args: Vec> = + vec![Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1))]; + ScalarFunctionExpr::new( + "test_udf", + test_udf(), + args, + Field::new("out", DataType::Int32, true).into(), + Arc::new(ConfigOptions::new()), + ) + .with_nullable(true) + } + + /// Encoder that yields a fixed child node and a configurable UDF payload. + struct MockEncoder { + udf_payload: Option>, + fail_child: bool, + } + + impl ScalarUdfProtoEncoder for MockEncoder { + fn encode_child( + &self, + _expr: &Arc, + ) -> Result { + if self.fail_child { + return internal_err!("child encode failed"); + } + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: None, + }) + } + + fn encode_udf(&self, _udf: &ScalarUDF) -> Result>> { + Ok(self.udf_payload.clone()) + } + } + + /// Decoder that resolves a fixed UDF, asserting the payload it was handed. + struct MockDecoder { + expected_payload: Option>, + fail_child: bool, + } + + impl ScalarUdfProtoDecoder for MockDecoder { + fn decode_child( + &self, + _node: &protobuf::PhysicalExprNode, + ) -> Result> { + if self.fail_child { + return internal_err!("child decode failed"); + } + Ok(Arc::new(Column::new("a", 0))) + } + + fn decode_udf( + &self, + _name: &str, + fun_definition: Option<&[u8]>, + ) -> Result> { + if fun_definition.map(|b| b.to_vec()) != self.expected_payload { + return internal_err!("decoder received unexpected fun_definition"); + } + Ok(test_udf()) + } + + fn config_options(&self) -> Arc { + Arc::new(ConfigOptions::new()) + } + } + + fn unwrap_udf_node( + node: protobuf::PhysicalExprNode, + ) -> protobuf::PhysicalScalarUdfNode { + match node.expr_type { + Some(protobuf::physical_expr_node::ExprType::ScalarUdf(n)) => *Box::new(n), + other => panic!("expected ScalarUdf node, got {other:?}"), + } + } + + #[test] + fn try_to_proto_encodes_fields() { + let node = sample_expr() + .try_to_proto(&MockEncoder { + udf_payload: None, + fail_child: false, + }) + .unwrap(); + assert_eq!(node.expr_id, None); + let udf = unwrap_udf_node(node); + assert_eq!(udf.name, "test_udf"); + assert_eq!(udf.args.len(), 2); + assert_eq!(udf.fun_definition, None); + assert!(udf.nullable); + assert_eq!(udf.return_field_name, "out"); + assert!(udf.return_type.is_some()); + } + + #[test] + fn try_to_proto_passes_codec_payload() { + let node = sample_expr() + .try_to_proto(&MockEncoder { + udf_payload: Some(vec![1, 2, 3]), + fail_child: false, + }) + .unwrap(); + assert_eq!(unwrap_udf_node(node).fun_definition, Some(vec![1, 2, 3])); + } + + #[test] + fn try_to_proto_propagates_child_error() { + let err = sample_expr() + .try_to_proto(&MockEncoder { + udf_payload: None, + fail_child: true, + }) + .unwrap_err(); + assert!(err.to_string().contains("child encode failed")); + } + + #[test] + fn try_from_proto_roundtrip() { + let node = sample_expr() + .try_to_proto(&MockEncoder { + udf_payload: Some(vec![9]), + fail_child: false, + }) + .unwrap(); + let udf_node = unwrap_udf_node(node); + let decoded = ScalarFunctionExpr::try_from_proto( + &udf_node, + &MockDecoder { + expected_payload: Some(vec![9]), + fail_child: false, + }, + ) + .unwrap(); + let decoded = decoded + .downcast_ref::() + .expect("decoded to ScalarFunctionExpr"); + assert_eq!(decoded.name(), "test_udf"); + assert_eq!(decoded.args().len(), 2); + assert!(decoded.nullable()); + assert_eq!( + decoded.return_field(&Schema::empty()).unwrap().name(), + "out" + ); + } + + #[test] + fn try_from_proto_errors_on_missing_return_type() { + let mut udf_node = unwrap_udf_node( + sample_expr() + .try_to_proto(&MockEncoder { + udf_payload: None, + fail_child: false, + }) + .unwrap(), + ); + udf_node.return_type = None; + let err = ScalarFunctionExpr::try_from_proto( + &udf_node, + &MockDecoder { + expected_payload: None, + fail_child: false, + }, + ) + .unwrap_err(); + assert!(err.to_string().contains("return_type")); + } + + #[test] + fn try_from_proto_propagates_child_error() { + let udf_node = unwrap_udf_node( + sample_expr() + .try_to_proto(&MockEncoder { + udf_payload: None, + fail_child: false, + }) + .unwrap(), + ); + let err = ScalarFunctionExpr::try_from_proto( + &udf_node, + &MockDecoder { + expected_payload: None, + fail_child: true, + }, + ) + .unwrap_err(); + assert!(err.to_string().contains("child decode failed")); + } +} + #[cfg(test)] mod tests { use super::*; From dde0f3f43c30afdcfc1e4a280a192b00bb4502e8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 10 Jul 2026 12:27:56 -0500 Subject: [PATCH 2/2] feat: route ScalarFunctionExpr proto through the typed bridge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement `ScalarUdfProtoEncoder` on `ConverterEncoder` and add a `ScalarUdfConverterDecoder` implementing `ScalarUdfProtoDecoder` (reproducing the exact registry-then-codec UDF lookup order). The encode downcast arm and the decode `ExprType::ScalarUdf` arm now delegate to `ScalarFunctionExpr::{try_to_proto, try_from_proto}` — the marshalling logic lives with the expression, with no type erasure. Wire format is byte-for-byte unchanged; `proto_integration` (190 tests) still passes, exercising the scalar-UDF round-trip through the new bridge. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_018LmymuRdTRmGthjMhtCijn --- .../proto/src/physical_plan/from_proto.rs | 85 ++++++++++++------- .../proto/src/physical_plan/to_proto.rs | 49 ++++++----- 2 files changed, 83 insertions(+), 51 deletions(-) diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 53ff4a41d466e..870b976d36027 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use arrow::array::RecordBatch; use arrow::compute::SortOptions; -use arrow::datatypes::{Field, Schema}; +use arrow::datatypes::Schema; use arrow::ipc::reader::StreamReader; use chrono::{TimeZone, Utc}; use datafusion_common::{ @@ -38,12 +38,14 @@ use datafusion_datasource_json::file_format::JsonSink; use datafusion_datasource_parquet::file_format::ParquetSink; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; -use datafusion_expr::WindowFunctionDefinition; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::SubqueryIndex; +use datafusion_expr::{ScalarUDF, WindowFunctionDefinition}; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; -use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, ScalarFunctionExpr}; +use datafusion_physical_expr::{ + LexOrdering, PhysicalSortExpr, ScalarFunctionExpr, ScalarUdfProtoDecoder, +}; use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, @@ -305,34 +307,16 @@ pub fn parse_physical_expr_with_converter( ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?, ExprType::TryCast(_) => TryCastExpr::try_from_proto(proto, &decode_ctx)?, ExprType::ScalarUdf(e) => { - let udf = match &e.fun_definition { - Some(buf) => ctx.codec().try_decode_udf(&e.name, buf)?, - None => ctx - .task_ctx() - .udf(e.name.as_str()) - .or_else(|_| ctx.codec().try_decode_udf(&e.name, &[]))?, + // `ScalarFunctionExpr` decodes itself against a fully-typed bridge + // (no `dyn Any`); this struct supplies the codec, registry, and + // session config it needs. See its proto module docs for why it + // can't ride the crate-wide decode context. + let decoder = ScalarUdfConverterDecoder { + ctx, + proto_converter, + input_schema, }; - let scalar_fun_def = Arc::clone(&udf); - - let args = parse_physical_exprs(&e.args, ctx, input_schema, proto_converter)?; - - let config_options = Arc::clone(ctx.task_ctx().session_config().options()); - - Arc::new( - ScalarFunctionExpr::new( - e.name.as_str(), - scalar_fun_def, - args, - Field::new( - &e.return_field_name, - convert_required!(e.return_type)?, - true, - ) - .into(), - config_options, - ) - .with_nullable(e.nullable), - ) + ScalarFunctionExpr::try_from_proto(e, &decoder)? } ExprType::LikeExpr(_) => LikeExpr::try_from_proto(proto, &decode_ctx)?, ExprType::HashExpr(_) => HashExpr::try_from_proto(proto, &decode_ctx)?, @@ -786,6 +770,47 @@ impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprD } } +/// Typed decode bridge for [`ScalarFunctionExpr`]. Bundles the codec, +/// function registry, and session config the expression needs to +/// reconstruct itself — none of which the crate-wide decode context can +/// carry, because [`ScalarUDF`] and the codec live above +/// `physical-expr-common` in the crate graph. Reproduces the exact +/// registry-then-codec lookup order of the old inline arm. +struct ScalarUdfConverterDecoder<'a, 'b, 'c> { + ctx: &'a PhysicalPlanDecodeContext<'b>, + proto_converter: &'a dyn PhysicalProtoConverterExtension, + input_schema: &'c Schema, +} + +impl ScalarUdfProtoDecoder for ScalarUdfConverterDecoder<'_, '_, '_> { + fn decode_child( + &self, + node: &protobuf::PhysicalExprNode, + ) -> Result> { + self.proto_converter + .proto_to_physical_expr(node, self.input_schema, self.ctx) + } + + fn decode_udf( + &self, + name: &str, + fun_definition: Option<&[u8]>, + ) -> Result> { + match fun_definition { + 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 config_options(&self) -> Arc { + Arc::clone(self.ctx.task_ctx().session_config().options()) + } +} + #[cfg(test)] mod tests { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 4614c4f002169..11485ceba89a5 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -30,10 +30,10 @@ use datafusion_datasource_csv::file_format::CsvSink; use datafusion_datasource_json::file_format::JsonSink; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::ParquetSink; -use datafusion_expr::WindowFrame; -use datafusion_physical_expr::ScalarFunctionExpr; +use datafusion_expr::{ScalarUDF, WindowFrame}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; +use datafusion_physical_expr::{ScalarFunctionExpr, ScalarUdfProtoEncoder}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; @@ -271,6 +271,26 @@ impl datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprE } } +/// Typed encode bridge for [`ScalarFunctionExpr`], whose `ScalarUDF` cannot be +/// carried by the crate-wide encode context (crate-graph cycle — see +/// `ScalarFunctionExpr`'s proto module docs). Routes the UDF through the same +/// [`PhysicalExtensionCodec::try_encode_udf`] path as the old inline arm. +impl ScalarUdfProtoEncoder for ConverterEncoder<'_> { + fn encode_child( + &self, + expr: &Arc, + ) -> Result { + self.proto_converter + .physical_expr_to_proto(expr, self.codec) + } + + fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + let mut buf = Vec::new(); + self.codec.try_encode_udf(udf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } +} + /// Serialize a `PhysicalExpr` to default protobuf representation. /// /// If required, a [`PhysicalExtensionCodec`] can be provided which can handle @@ -299,25 +319,12 @@ pub fn serialize_physical_expr_with_converter( return Ok(node); } - if let Some(expr) = expr.downcast_ref::() { - let mut buf = Vec::new(); - codec.try_encode_udf(expr.fun(), &mut buf)?; - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarUdf( - protobuf::PhysicalScalarUdfNode { - name: expr.name().to_string(), - args: serialize_physical_exprs(expr.args(), codec, proto_converter)?, - fun_definition: (!buf.is_empty()).then_some(buf), - return_type: Some(expr.return_type().try_into()?), - nullable: expr.nullable(), - return_field_name: expr - .return_field(&Schema::empty())? - .name() - .to_string(), - }, - )), - }) + if let Some(scalar_fn) = expr.downcast_ref::() { + // `ScalarFunctionExpr` declares its own ser/de against a fully-typed + // bridge (no `dyn Any`); we hand it the encoder that already backs the + // crate-wide hook. See its proto module docs for why it can't ride the + // hook itself. + scalar_fn.try_to_proto(&encoder) } else if let Some(expr) = expr.downcast_ref::() { Ok(protobuf::PhysicalExprNode { expr_id,