Skip to content

Port ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities#23421

Open
adriangb wants to merge 3 commits into
apache:mainfrom
pydantic:scalar-udf-codec
Open

Port ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities#23421
adriangb wants to merge 3 commits into
apache:mainfrom
pydantic:scalar-udf-codec

Conversation

@adriangb

@adriangb adriangb commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Part of epic #22418. Alternative to #23415 — see "Rationale" for the relationship.

Rationale for this change

ScalarFunctionExpr was the last expression on the central proto downcast chain whose serialization needs session-level objects: the embedded ScalarUDF must go through PhysicalExtensionCodec on encode and resolve via codec + function registry on decode, and the reconstructed expression needs the session's ConfigOptions.

The tension: the try_to_proto / try_from_proto contexts live in datafusion-physical-expr-common, which sits below datafusion-expr in the crate graph (datafusion-expr depends on it for Arc<dyn PhysicalExpr> in the UDAF/UDWF APIs), so the contexts cannot name ScalarUDF without a dependency cycle. #23415 resolves this with caller-visible type erasure — encode_udf(&dyn Any) / decode_udf(..) -> Arc<dyn Any> on the public ctx API, with docs saying "the concrete type is always ScalarUDF" and downcasts at every use — plus a config_options() default that silently fabricates ConfigOptions::default().

This PR takes a different shape: generic, fully-typed capability accessors on the contexts, with the erasure confined to the internal dispatch traits.

// encode side — call site in ScalarFunctionExpr::try_to_proto:
fun_definition: ctx.encode_function(self.fun.as_ref())?,

// decode side — call site in ScalarFunctionExpr::try_from_proto:
let udf: Arc<ScalarUDF> = ctx.decode_function(&node.name, node.fun_definition.as_deref())?;
let config = ctx.config_options()?;
  • PhysicalExprEncodeCtx::encode_function<F: Any>(&F) -> Result<Option<Vec<u8>>> and PhysicalExprDecodeCtx::decode_function<F: Any>(name, payload) -> Result<Arc<F>> are generic over the function type. Call sites — including third-party PhysicalExpr implementations that wrap a ScalarUDF (or later an AggregateUDF/WindowUDF) — are fully typed and never see Any or write a downcast.
  • The internal dispatch traits (PhysicalExprEncode/PhysicalExprDecode, documented as implemented only by serialization layers) carry monomorphic erased channels (encode_function_erased, decode_function_erased) that route on TypeId and receive type_name for error messages. The ctx verifies the returned type and converts any contract violation into a clean internal error. Which function types a layer supports is an explicit, documented part of its contract; requesting an unsupported type errors naming it.
  • PhysicalExprDecodeCtx::config_options() returns Result and the dispatch default errors — a layer without a session must opt in explicitly rather than silently substituting defaults (avoiding the footgun in the ConfigOptions::default() fallback).

With that in place, ScalarFunctionExpr rides the standard hooks like every other migrated expression: try_to_proto inside impl PhysicalExpr, try_from_proto(node, ctx) with the standard signature, one-line dispatch in from_proto.rs, no downcast arm in to_proto.rs at all — the encode chain shrinks by one special case, and datafusion-proto no longer imports ScalarFunctionExpr for encoding.

datafusion-proto's ConverterEncoder/ConverterDecoder implement the erased channels by routing ScalarUDF through try_encode_udf/try_decode_udf with the task-context registry fallback (identical lookup order to the old inline arm). Aggregate/window UDFs slot into the same channels when their expressions migrate.

Honest limitation: which F a layer supports is a runtime contract, not a compile-time one. That is the floor imposed by the crate graph — a compile-time signature would need ScalarUDF below physical-expr-common, which its logical half (simplify, call, preimage) rules out.

What changes are included in this PR?

Commit 1 — datafusion/physical-expr-common/src/physical_expr.rs:

  • PhysicalExprEncodeCtx::encode_function / PhysicalExprDecodeCtx::decode_function / PhysicalExprDecodeCtx::config_options (typed public surface).
  • PhysicalExprEncode::encode_function_erased / PhysicalExprDecode::decode_function_erased / PhysicalExprDecode::config_options (internal dispatch, erroring defaults so existing implementors keep compiling).
  • Module docs describing the capability contract.

Commit 2:

  • datafusion/physical-expr/src/scalar_function.rs: try_to_proto in impl PhysicalExpr (feature-gated proto), try_from_proto inherent fn with the standard signature, plus a proto_tests module.
  • datafusion/proto/src/physical_plan/to_proto.rs: ConverterEncoder::encode_function_erased; the ScalarFunctionExpr downcast arm is deleted.
  • datafusion/proto/src/physical_plan/from_proto.rs: ConverterDecoder::decode_function_erased + config_options; the inline ExprType::ScalarUdf arm becomes the standard one-line dispatch.

The wire format is byte-for-byte unchanged: same PhysicalScalarUdfNode fields, fun_definition only set when the codec writes bytes, return_field_name from the stored return field (identical to the old return_field(&Schema::empty())?.name()), expr_id: None (ScalarFunctionExpr never overrides expression_id, so the old arm always wrote None here too).

Are these changes tested?

  • 11 new direct tests in scalar_function.rs::proto_tests, driving try_to_proto through dyn PhysicalExpr (so inherent-method shadowing cannot silently pass): encode happy path (field-by-field), codec payload passthrough in both directions, decode happy path, wrong ExprType variant, missing return_type, child + function-codec error propagation on both sides, the ctx downcast guard against a contract-violating layer, and the no-session config_options error.
  • cargo test -p datafusion-proto --test proto_integration: 190 passed. With both inline arms deleted, scalar-UDF roundtrip coverage can only pass through the new hooks.
  • cargo fmt --all and cargo clippy --all-targets --all-features -- -D warnings are clean.

Are there any user-facing changes?

No wire-format or behavior changes. New public API: the three typed ctx methods and the three defaulted dispatch-trait methods (all feature proto, all additive/non-breaking). No existing signatures change.

🤖 Generated with Claude Code

https://claude.ai/code/session_018MmNo1SfPHXZNaMW26KmaG

@github-actions github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates proto Related to proto crate labels Jul 9, 2026
adriangb and others added 2 commits July 9, 2026 09:42
…to expr contexts

PhysicalExprEncodeCtx::encode_function<F> and
PhysicalExprDecodeCtx::decode_function<F> expose the serialization layer's
UDF codec through generic, fully-typed accessors: call sites name the
function type (e.g. datafusion_expr::ScalarUDF) and never see erasure or
write downcasts. Function types live above this crate in the dependency
graph, so the internal dispatch traits carry monomorphic type-erased
channels routed by TypeId; the ctx verifies the returned type and turns
contract violations into clean internal errors. Which types a layer
supports is an explicit part of its contract.

PhysicalExprDecodeCtx::config_options provides session configuration with
an erroring default: a layer without a session must opt in explicitly
rather than silently substituting ConfigOptions::default().

All new dispatch-trait methods are defaulted, so existing implementors
keep compiling.

Part of apache#22418; groundwork for apache#22430.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018MmNo1SfPHXZNaMW26KmaG
Rides the standard hooks like every other migrated expression: the
embedded ScalarUDF serializes through the encode ctx's typed
encode_function capability and resolves on decode via decode_function
(payload-first through the extension codec, then the task context's
function registry, then the codec with an empty payload — the same lookup
order as the old inline arm), with session configuration from
ctx.config_options(). The ScalarFunctionExpr downcast arm in to_proto.rs
is deleted and the ExprType::ScalarUdf decode arm becomes the standard
one-line dispatch, so the roundtrip suite can only pass through the new
hooks. Wire format is byte-for-byte unchanged.

Closes apache#22430. Part of apache#22418.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018MmNo1SfPHXZNaMW26KmaG
@adriangb adriangb changed the title Port ScalarFunctionExpr proto serde to a typed ScalarUdfCodec (no type erasure) Port ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities Jul 9, 2026
@adriangb adriangb force-pushed the scalar-udf-codec branch from 1872465 to e51e167 Compare July 9, 2026 14:42
@adriangb adriangb marked this pull request as ready for review July 9, 2026 23:43
@adriangb adriangb requested a review from Copilot July 9, 2026 23:43
@adriangb

adriangb commented Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Hi @timsaucer , would you mind giving this beast a look? Thanks!

A serialization layer that implements only the required encode/decode
hooks must reject encode_function / decode_function / config_options
requests with errors naming the requested type, exercised directly
through the public ctx surface.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PzhFvJM2cEXzy9tixZeRYg
@github-actions github-actions Bot removed the logical-expr Logical plan and expressions label Jul 9, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Ports ScalarFunctionExpr onto the standard PhysicalExpr::try_to_proto / try_from_proto hooks by extending the encode/decode contexts with typed “function codec” capabilities (plus session ConfigOptions access), allowing scalar UDF payloads to be routed through the extension codec without reintroducing crate dependency cycles.

Changes:

  • Added typed encode_function / decode_function and config_options helpers on the proto (de)serialization contexts, backed by internal type-erased dispatch.
  • Migrated ScalarFunctionExpr to implement try_to_proto (trait override) and try_from_proto (inherent constructor), with comprehensive direct hook tests.
  • Removed the special-case ScalarFunctionExpr downcast arm from to_proto.rs and replaced the inline ScalarUDF decode logic in from_proto.rs with standard dispatch.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File Description
datafusion/physical-expr-common/src/physical_expr.rs Introduces typed context capabilities for function (de)serialization and session config access, backed by erased dispatch traits.
datafusion/physical-expr/src/scalar_function.rs Implements ScalarFunctionExpr proto hooks and adds direct hook/unit tests.
datafusion/proto/src/physical_plan/to_proto.rs Implements encode_function_erased and removes the ScalarFunctionExpr downcast serialization arm.
datafusion/proto/src/physical_plan/from_proto.rs Implements decode_function_erased + config_options and routes ScalarUdf decoding through ScalarFunctionExpr::try_from_proto.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@timsaucer timsaucer left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This adds a bit of complexity, but not a ton, to the encoding/decoding approach. I'm not a huge fan of having to jump through the erasing to an Arc<dyn Any>. I think as we discussed this is entirely because the circular dependency that would be added, right? Since I'm dealing with a related problem in #23348 where I have to either hoist traits to a higher level crate or do type erasure like you're doing.

The solution here is tactical and it solves your immediate problem, but I'm wondering if we should take a more strategic approach and think about reorganizing some of our traits & crates so that you never had to do that in the first place. What do you think?

@adriangb

Copy link
Copy Markdown
Contributor Author

I think as we discussed this is entirely because the circular dependency that would be added, right?

Yes, pretty much. The issue in this case is that PhysicalExpr declares the trait method. For ScalarUDF to serialize itself it needs to call PhysicalExtensionCodec::try_encode_udf. But in order to get a PhysicalExtensionCodec that needs to be passed in via PhysicalExpr::to_proto which results in the requirement that PhysicalExpr reference PhysicalExtensionCodec which then references ScalarUDF. But ScalarUDF also references PhysicalExpr, creating a circular dependency.

The solution here is tactical and it solves your immediate problem, but I'm wondering if we should take a more strategic approach and think about reorganizing some of our traits & crates so that you never had to do that in the first place.

That sounds great. I spent quite some time trying something along these lines, but kept hitting dead ends. Do you have a feeling for this being possible?

@timsaucer

Copy link
Copy Markdown
Member

That sounds great. I spent quite some time trying something along these lines, but kept hitting dead ends. Do you have a feeling for this being possible?

I just spent some time looking at it again and I don't have a great solution at this time :|

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-expr Changes to the physical-expr crates proto Related to proto crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Port ScalarFunctionExpr to use try_to_proto / try_from_proto

3 participants