From 38a02b7aa4de525a1fbd0c33432e4490e870b3a5 Mon Sep 17 00:00:00 2001 From: Olawoyin007 Date: Thu, 9 Jul 2026 09:47:11 +0100 Subject: [PATCH] Port ScalarFunctionExpr to use try_to_proto / try_from_proto Sub-issue of #22418; closes #22430. ScalarFunctionExpr needs the extension codec and session config, so this also adds the ctx helpers discussed on the epic: encode_udf on PhysicalExprEncodeCtx, decode_udf and config_options on PhysicalExprDecodeCtx. The UDF signatures are type-erased because datafusion-expr depends on datafusion-physical-expr-common, so the ctx traits cannot name ScalarUDF without a dependency cycle; the downcasts live in ConverterEncoder and ScalarFunctionExpr::try_from_proto and fail with clean internal errors. Both central arms (to_proto.rs downcast, from_proto.rs inline decode) are deleted in this change, so the roundtrip tests only pass through the new hooks. Wire format is unchanged. --- .../physical-expr-common/src/physical_expr.rs | 79 +++- .../physical-expr/src/proto_test_util.rs | 9 + .../physical-expr/src/scalar_function.rs | 344 ++++++++++++++++++ .../proto/src/physical_plan/from_proto.rs | 55 ++- .../proto/src/physical_plan/to_proto.rs | 38 +- 5 files changed, 468 insertions(+), 57 deletions(-) diff --git a/datafusion/physical-expr-common/src/physical_expr.rs b/datafusion/physical-expr-common/src/physical_expr.rs index 679a44e85ee9a..e2b10210603a8 100644 --- a/datafusion/physical-expr-common/src/physical_expr.rs +++ b/datafusion/physical-expr-common/src/physical_expr.rs @@ -519,10 +519,10 @@ pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash { /// `datafusion-proto`, which is what lets `physical-expr-common` stay free /// of `datafusion-proto` as a dep. /// -/// More specialized helpers (e.g. encoding UDFs/UDAFs/UDWFs through the -/// extension codec) can be added to the context as expressions migrate; -/// today they're not required because the encoder forwards to the existing -/// codec via the proto converter. +/// More specialized helpers can be added to the context as expressions +/// migrate; [`proto_encode::PhysicalExprEncodeCtx::encode_udf`] (routing +/// scalar UDFs through the extension codec) is the first such helper — +/// UDAFs/UDWFs will follow the same shape when their expressions migrate. #[cfg(feature = "proto")] pub mod proto_encode { use std::sync::Arc; @@ -574,6 +574,23 @@ pub mod proto_encode { .map(|expr| self.encode_child(expr)) .collect() } + + /// Encode a scalar UDF through the registered extension codec, + /// returning the codec's opaque payload for the proto + /// `fun_definition` field — `None` when the codec wrote nothing + /// (built-in functions that decode by name alone). + /// + /// `udf` is type-erased: the concrete type must be + /// `datafusion_expr::ScalarUDF`. It cannot be named here because + /// `datafusion-expr` depends on this crate, so a typed signature + /// would create a dependency cycle; the encoder in + /// `datafusion-proto` downcasts and errors on any other type. + pub fn encode_udf( + &self, + udf: &(dyn std::any::Any + Send + Sync), + ) -> Result>> { + self.encoder.encode_udf(udf) + } } /// Internal dispatch trait. Implementors live in `datafusion-proto` and @@ -583,6 +600,17 @@ pub mod proto_encode { pub trait PhysicalExprEncode { /// Encode an expression to a protobuf node. fn encode(&self, expr: &Arc) -> Result; + + /// Encode a scalar UDF (type-erased `datafusion_expr::ScalarUDF`; + /// see [`PhysicalExprEncodeCtx::encode_udf`]) into its opaque + /// `fun_definition` payload. The default errors so encoders that + /// never see UDF expressions need not implement it. + fn encode_udf( + &self, + _udf: &(dyn std::any::Any + Send + Sync), + ) -> Result>> { + datafusion_common::internal_err!("this encoder does not support scalar UDFs") + } } } @@ -733,6 +761,29 @@ pub mod proto_decode { { nodes.into_iter().map(|node| self.decode(node)).collect() } + + /// Resolve a scalar UDF by name, preferring the opaque + /// `fun_definition` payload written by the encode-side extension + /// codec and falling back to the session's function registry. + /// + /// The returned value is type-erased: the concrete type is + /// `Arc`. It cannot be named here + /// because `datafusion-expr` depends on this crate, so a typed + /// signature would create a dependency cycle; callers downcast with + /// [`Arc::downcast`]. + pub fn decode_udf( + &self, + name: &str, + fun_definition: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udf(name, fun_definition) + } + + /// Session configuration options for constructing expressions that + /// carry them (e.g. `ScalarFunctionExpr`). + pub fn config_options(&self) -> Arc { + self.decoder.config_options() + } } /// Unwrap a required non-expression proto field. @@ -774,6 +825,26 @@ pub mod proto_decode { node: &PhysicalExprNode, schema: &Schema, ) -> Result>; + + /// Resolve a scalar UDF by name (type-erased + /// `Arc`; see + /// [`PhysicalExprDecodeCtx::decode_udf`]). The default errors so + /// decoders that never see UDF expressions need not implement it. + fn decode_udf( + &self, + name: &str, + _fun_definition: Option<&[u8]>, + ) -> Result> { + datafusion_common::internal_err!( + "this decoder does not support scalar UDFs (attempted to decode '{name}')" + ) + } + + /// Session configuration options handed to decoded expressions. The + /// default returns fresh defaults for decoders without a session. + fn config_options(&self) -> Arc { + Arc::new(datafusion_common::config::ConfigOptions::default()) + } } } diff --git a/datafusion/physical-expr/src/proto_test_util.rs b/datafusion/physical-expr/src/proto_test_util.rs index ab280335800b9..6ab8f9919b4bd 100644 --- a/datafusion/physical-expr/src/proto_test_util.rs +++ b/datafusion/physical-expr/src/proto_test_util.rs @@ -138,4 +138,13 @@ impl PhysicalExprEncode for StubEncoder { } Ok(column_node("child")) } + + fn encode_udf( + &self, + _udf: &(dyn std::any::Any + Send + Sync), + ) -> Result>> { + // Behaves like a codec that writes nothing (built-in functions + // decodable by name alone). + Ok(None) + } } diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 418d005c971ea..058f9a96ae00c 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -349,6 +349,90 @@ impl PhysicalExpr for ScalarFunctionExpr { self.args.iter().map(|arg| arg.placement()).collect(); self.fun.placement(&arg_placements) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarUdf( + protobuf::PhysicalScalarUdfNode { + name: self.name.clone(), + args: ctx.encode_children_expressions(&self.args)?, + fun_definition: ctx.encode_udf(self.fun.as_ref())?, + return_type: Some(self.return_type().try_into()?), + nullable: self.nullable(), + return_field_name: self.return_field.name().to_string(), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl ScalarFunctionExpr { + /// Reconstruct a [`ScalarFunctionExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`] so the decode signature matches + /// other migrated expressions and can inspect outer-node metadata if + /// needed in the future. The UDF itself is resolved through + /// [`PhysicalExprDecodeCtx::decode_udf`]: the extension codec's opaque + /// `fun_definition` payload wins, falling back to the session's function + /// registry by name. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + /// [`PhysicalExprDecodeCtx::decode_udf`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode_udf + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use arrow::datatypes::Field; + use datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_physical_expr_common::physical_expr::proto_decode::require_proto_field; + use datafusion_proto_models::protobuf; + + let scalar_udf = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::ScalarUdf, + "ScalarFunctionExpr", + ); + + let udf = ctx + .decode_udf(&scalar_udf.name, scalar_udf.fun_definition.as_deref())? + .downcast::() + .map_err(|_| { + internal_datafusion_err!( + "decode_udf returned a non-ScalarUDF value for '{}'", + scalar_udf.name + ) + })?; + + let args = ctx.decode_children_expressions(&scalar_udf.args)?; + let return_type = require_proto_field( + scalar_udf.return_type.as_ref(), + "ScalarFunctionExpr", + "return_type", + )?; + let return_field = + Field::new(&scalar_udf.return_field_name, return_type.try_into()?, true) + .into(); + + Ok(Arc::new( + ScalarFunctionExpr::new( + &scalar_udf.name, + udf, + args, + return_field, + ctx.config_options(), + ) + .with_nullable(scalar_udf.nullable), + )) + } } #[cfg(test)] @@ -426,3 +510,263 @@ mod tests { assert!(!is_volatile(&stable_arc)); } } + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::Column; + use crate::proto_test_util::{StubDecoder, StubEncoder, column_node}; + use arrow::datatypes::Field; + use datafusion_expr::Signature; + use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, + }; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::datafusion_common::ArrowType; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalScalarUdfNode, physical_expr_node, + }; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdf { + signature: Signature, + } + + 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(42)))) + } + } + + fn test_udf() -> Arc { + Arc::new(ScalarUDF::from(TestUdf { + signature: Signature::uniform( + 1, + vec![DataType::Int32], + Volatility::Immutable, + ), + })) + } + + /// A `ScalarFunctionExpr` calling [`test_udf`] over one `Int32` column. + fn proto_scalar_fn_fixture() -> ScalarFunctionExpr { + ScalarFunctionExpr::new( + "test_udf", + test_udf(), + vec![Arc::new(Column::new("a", 0)) as Arc], + Field::new("test_udf", DataType::Int32, true).into(), + Arc::new(ConfigOptions::default()), + ) + } + + fn proto_int32_arrow_type() -> ArrowType { + (&DataType::Int32).try_into().unwrap() + } + + /// Build a `ScalarUdf` proto node with the given fields. + fn proto_scalar_udf_node( + args: Vec, + return_type: Option, + nullable: bool, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::ScalarUdf( + PhysicalScalarUdfNode { + name: "test_udf".to_string(), + args, + fun_definition: None, + return_type, + nullable, + return_field_name: "test_udf".to_string(), + }, + )), + } + } + + /// Decoder stub that resolves every UDF name to [`test_udf`], delegating + /// child decoding to an inner [`StubDecoder`]. + struct UdfStubDecoder { + inner: StubDecoder, + } + + impl PhysicalExprDecode for UdfStubDecoder { + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.inner.decode(node, schema) + } + + fn decode_udf( + &self, + _name: &str, + _fun_definition: Option<&[u8]>, + ) -> Result> { + Ok(test_udf() as _) + } + } + + /// Decoder stub whose `decode_udf` returns a value that is not a + /// `ScalarUDF`, to exercise the downcast reject path. + struct NonUdfStubDecoder; + + impl PhysicalExprDecode for NonUdfStubDecoder { + fn decode( + &self, + _node: &PhysicalExprNode, + _schema: &Schema, + ) -> Result> { + unreachable!("child decoding must not be reached") + } + + fn decode_udf( + &self, + _name: &str, + _fun_definition: Option<&[u8]>, + ) -> Result> { + Ok(Arc::new("not a udf") as _) + } + } + + #[test] + fn try_to_proto_encodes_scalar_function_expr() { + let expr = proto_scalar_fn_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = expr + .try_to_proto(&ctx) + .unwrap() + .expect("ScalarFunctionExpr must serialize itself"); + + assert_eq!(node.expr_id, None); + let Some(physical_expr_node::ExprType::ScalarUdf(udf_node)) = node.expr_type + else { + panic!("expected ScalarUdf node"); + }; + assert_eq!(udf_node.name, "test_udf"); + assert_eq!(udf_node.args, vec![column_node("child")]); + assert_eq!(udf_node.fun_definition, None); + assert_eq!(udf_node.return_type, Some(proto_int32_arrow_type())); + assert!(udf_node.nullable); + assert_eq!(udf_node.return_field_name, "test_udf"); + } + + #[test] + fn try_to_proto_propagates_child_encode_error() { + let expr = proto_scalar_fn_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let err = expr.try_to_proto(&ctx).unwrap_err(); + assert!(err.to_string().contains("stub encode failure"), "{err}"); + } + + #[test] + fn try_from_proto_decodes_scalar_function_expr() { + let node = proto_scalar_udf_node( + vec![column_node("a")], + Some(proto_int32_arrow_type()), + false, + ); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = UdfStubDecoder { + inner: StubDecoder::ok(), + }; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = ScalarFunctionExpr::try_from_proto(&node, &ctx).unwrap(); + let expr = decoded + .downcast_ref::() + .expect("expected a ScalarFunctionExpr"); + + assert_eq!(expr.name(), "test_udf"); + assert_eq!(expr.fun().name(), "test_udf"); + assert_eq!(expr.args().len(), 1); + assert_eq!(expr.return_type(), &DataType::Int32); + assert!(!expr.nullable()); + } + + #[test] + fn try_from_proto_rejects_non_scalar_udf_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = NonUdfStubDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = ScalarFunctionExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalExprNode is not a ScalarFunctionExpr"), + "{err}" + ); + } + + #[test] + fn try_from_proto_rejects_missing_return_type() { + let node = proto_scalar_udf_node(vec![column_node("a")], None, true); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = UdfStubDecoder { + inner: StubDecoder::ok(), + }; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = ScalarFunctionExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("missing required field 'return_type'"), + "{err}" + ); + } + + #[test] + fn try_from_proto_rejects_non_scalar_udf_value() { + let node = proto_scalar_udf_node( + vec![column_node("a")], + Some(proto_int32_arrow_type()), + true, + ); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = NonUdfStubDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = ScalarFunctionExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("decode_udf returned a non-ScalarUDF value"), + "{err}" + ); + } + + #[test] + fn try_from_proto_propagates_child_decode_error() { + let node = proto_scalar_udf_node( + vec![column_node("a")], + Some(proto_int32_arrow_type()), + true, + ); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = UdfStubDecoder { + inner: StubDecoder::failing_on(1), + }; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = ScalarFunctionExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(err.to_string().contains("stub decode failure"), "{err}"); + } +} diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 53ff4a41d466e..5cdf049c28c3d 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::{ @@ -304,36 +304,7 @@ pub fn parse_physical_expr_with_converter( ExprType::Case(_) => CaseExpr::try_from_proto(proto, &decode_ctx)?, 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, &[]))?, - }; - 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), - ) - } + ExprType::ScalarUdf(_) => ScalarFunctionExpr::try_from_proto(proto, &decode_ctx)?, ExprType::LikeExpr(_) => LikeExpr::try_from_proto(proto, &decode_ctx)?, ExprType::HashExpr(_) => HashExpr::try_from_proto(proto, &decode_ctx)?, ExprType::ScalarSubquery(sq) => { @@ -784,6 +755,28 @@ impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprD self.proto_converter .proto_to_physical_expr(node, schema, self.ctx) } + + fn decode_udf( + &self, + name: &str, + fun_definition: Option<&[u8]>, + ) -> Result> { + // Returned type-erased to avoid a dependency cycle (see the trait + // docs); the concrete type is always `Arc`. + let udf = 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, &[]))?, + }; + Ok(udf) + } + + fn config_options(&self) -> Arc { + Arc::clone(self.ctx.task_ctx().session_config().options()) + } } #[cfg(test)] diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 4614c4f002169..e027bd0d67da5 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -30,8 +30,7 @@ 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_common::sort_expr::PhysicalSortExpr; @@ -269,6 +268,20 @@ impl datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprE self.proto_converter .physical_expr_to_proto(expr, self.codec) } + + fn encode_udf( + &self, + udf: &(dyn std::any::Any + Send + Sync), + ) -> Result>> { + // Type-erased at the trait to avoid a dependency cycle (see the + // trait docs); the concrete type is always `ScalarUDF`. + let udf = udf + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("encode_udf expects a ScalarUDF"))?; + 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. @@ -299,26 +312,7 @@ 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(), - }, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { + if let Some(expr) = expr.downcast_ref::() { Ok(protobuf::PhysicalExprNode { expr_id, expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarSubquery(