Port ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities#23421
Port ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities#23421adriangb wants to merge 3 commits into
ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities#23421Conversation
…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
ScalarFunctionExpr proto serde to a typed ScalarUdfCodec (no type erasure)ScalarFunctionExpr to try_to_proto/try_from_proto via typed ctx function-codec capabilities
1872465 to
e51e167
Compare
|
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
There was a problem hiding this comment.
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_functionandconfig_optionshelpers on the proto (de)serialization contexts, backed by internal type-erased dispatch. - Migrated
ScalarFunctionExprto implementtry_to_proto(trait override) andtry_from_proto(inherent constructor), with comprehensive direct hook tests. - Removed the special-case
ScalarFunctionExprdowncast arm fromto_proto.rsand replaced the inline ScalarUDF decode logic infrom_proto.rswith 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
left a comment
There was a problem hiding this comment.
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?
Yes, pretty much. The issue in this case is that
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 :| |
Which issue does this PR close?
ScalarFunctionExprto usetry_to_proto/try_from_proto#22430.Part of epic #22418. Alternative to #23415 — see "Rationale" for the relationship.
Rationale for this change
ScalarFunctionExprwas the last expression on the central proto downcast chain whose serialization needs session-level objects: the embeddedScalarUDFmust go throughPhysicalExtensionCodecon encode and resolve via codec + function registry on decode, and the reconstructed expression needs the session'sConfigOptions.The tension: the
try_to_proto/try_from_protocontexts live indatafusion-physical-expr-common, which sits belowdatafusion-exprin the crate graph (datafusion-exprdepends on it forArc<dyn PhysicalExpr>in the UDAF/UDWF APIs), so the contexts cannot nameScalarUDFwithout 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 alwaysScalarUDF" and downcasts at every use — plus aconfig_options()default that silently fabricatesConfigOptions::default().This PR takes a different shape: generic, fully-typed capability accessors on the contexts, with the erasure confined to the internal dispatch traits.
PhysicalExprEncodeCtx::encode_function<F: Any>(&F) -> Result<Option<Vec<u8>>>andPhysicalExprDecodeCtx::decode_function<F: Any>(name, payload) -> Result<Arc<F>>are generic over the function type. Call sites — including third-partyPhysicalExprimplementations that wrap aScalarUDF(or later anAggregateUDF/WindowUDF) — are fully typed and never seeAnyor write a downcast.PhysicalExprEncode/PhysicalExprDecode, documented as implemented only by serialization layers) carry monomorphic erased channels (encode_function_erased,decode_function_erased) that route onTypeIdand receivetype_namefor 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()returnsResultand the dispatch default errors — a layer without a session must opt in explicitly rather than silently substituting defaults (avoiding the footgun in theConfigOptions::default()fallback).With that in place,
ScalarFunctionExprrides the standard hooks like every other migrated expression:try_to_protoinsideimpl PhysicalExpr,try_from_proto(node, ctx)with the standard signature, one-line dispatch infrom_proto.rs, no downcast arm into_proto.rsat all — the encode chain shrinks by one special case, anddatafusion-protono longer importsScalarFunctionExprfor encoding.datafusion-proto'sConverterEncoder/ConverterDecoderimplement the erased channels by routingScalarUDFthroughtry_encode_udf/try_decode_udfwith 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
Fa 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 needScalarUDFbelowphysical-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).Commit 2:
datafusion/physical-expr/src/scalar_function.rs:try_to_protoinimpl PhysicalExpr(feature-gatedproto),try_from_protoinherent fn with the standard signature, plus aproto_testsmodule.datafusion/proto/src/physical_plan/to_proto.rs:ConverterEncoder::encode_function_erased; theScalarFunctionExprdowncast arm is deleted.datafusion/proto/src/physical_plan/from_proto.rs:ConverterDecoder::decode_function_erased+config_options; the inlineExprType::ScalarUdfarm becomes the standard one-line dispatch.The wire format is byte-for-byte unchanged: same
PhysicalScalarUdfNodefields,fun_definitiononly set when the codec writes bytes,return_field_namefrom the stored return field (identical to the oldreturn_field(&Schema::empty())?.name()),expr_id: None(ScalarFunctionExprnever overridesexpression_id, so the old arm always wroteNonehere too).Are these changes tested?
scalar_function.rs::proto_tests, drivingtry_to_protothroughdyn PhysicalExpr(so inherent-method shadowing cannot silently pass): encode happy path (field-by-field), codec payload passthrough in both directions, decode happy path, wrongExprTypevariant, missingreturn_type, child + function-codec error propagation on both sides, the ctx downcast guard against a contract-violating layer, and the no-sessionconfig_optionserror.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 --allandcargo clippy --all-targets --all-features -- -D warningsare 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