adriangb opened a new pull request, #24628: URL: https://github.com/apache/datafusion/pull/24628
## Which issue does this PR close? - Closes #24626. ## Rationale for this change Extension `PhysicalExpr`s reach the wire as a `PhysicalExtensionExprNode`, which carries an opaque payload and the encoded children but no type discriminator. Because of that, **the codec *is* the discriminator**: `ComposedPhysicalExtensionCodec` tries each registered codec in turn and reads a decode error as "not mine". Resolution therefore depends on registration order and on error strings, and a name collision between two independent crates' codecs is undetectable. #23494 moved built-in expressions onto per-type `try_to_proto` / `try_from_proto` hooks, which are nicer to implement — encode and decode live next to the expression, and private state stays private. Third parties can use the encode half today (`PhysicalExpr::try_to_proto` is callable by anyone) but there is no way back, so they stay on `PhysicalExtensionCodec` for decode. This closes that asymmetry for expressions. This does **not** retire `PhysicalExtensionCodec`. UDF payloads and unmigrated expressions still need it, and it remains the fallback. ## What changes are included in this PR? Three commits, each green on its own. **1. Expose the session on `PhysicalExprDecodeCtx`.** The context had `schema()` and `decode()` but no way to reach session state, unlike `ExecutionPlanDecodeCtx::task_ctx()` on the plan side, so a decoder could not touch the function registry or session configuration. The issue describes this as three lines adding `fn task_ctx(&self) -> &TaskContext` to the dispatch trait. That turns out not to work: `PhysicalExprDecode` lives in `datafusion-physical-expr-common`, and `datafusion-execution` — home of `TaskContext` — **depends on** that crate, so naming `TaskContext` there is a dependency cycle. `datafusion-physical-plan` is the lowest crate that can name it, which is exactly why the plan-side accessor works and the expression-side one cannot be written the same way. The dispatch trait therefore hands the session over type-erased, and the public accessor downcasts: ```rust let task_ctx = ctx.task_ctx::<TaskContext>()?; ``` The new dispatch-trait method has a `None` default so downstream implementors keep compiling. Both expression dispatch traits also gained `#[doc(hidden)]` and the same "not public API" note the plan-side traits carry, reserving room for future additions — the issue asks for this on `PhysicalExprDecode`; I did both for symmetry. This half stands on its own merits even if the registry is deferred. **2. Add `optional string expr_name = 3;` to `PhysicalExtensionExprNode`.** Additive and proto3-compatible. Nothing reads it in this commit, and the codec path keeps writing `None`, since there the codec remains the discriminator. **3. The registry.** - `ExtensionPhysicalExpr` pairs a namespaced `EXPR_NAME` with a `try_from_proto` whose signature matches what built-in expressions already use, so the decoder is a plain `fn` pointer — no object-safety problem. - `PhysicalExprDecoderRegistry` maps names to decoders and rejects a duplicate name at registration. - `PhysicalExprEncodeCtx::encode_extension(expr, payload)` builds the wrapper node. It takes the expression itself rather than loose arguments, so the wire name (`T::EXPR_NAME`), the children (`expr.children()`, encoded through the context so recursion stays dedup-aware), and `expr_id` (`expr.expression_id()`) are all read off it and cannot drift from the decode half. - `SessionConfig::with_physical_expr::<T>()` in `datafusion-proto` registers on a session. It merges into whatever registry the config already holds, so one library cannot silently drop another's registrations. Decode rule, as proposed in the issue: name present **and** registered → registry; otherwise → `PhysicalExtensionCodec::try_decode_expr`, unchanged. There is exactly one `ExprType::Extension` decode site and the fallback arm is byte-for-byte what it was, pre-decoded `inputs` included. Registry types live in `datafusion-physical-expr-common` next to the contexts they are expressed in terms of, so a crate defining an extension expression does not need to depend on `datafusion-proto`. Only the session glue lives in `datafusion-proto`. Session scoping goes through `SessionConfig` extensions rather than a new field on `TaskContext`/`SessionState`. A real field would force `datafusion-physical-expr-common/proto` on for everyone, and `datafusion-execution` takes that dependency with `default-features = false` precisely so crates that never serialize plans pay nothing. ## Are these changes tested? Yes — `datafusion/proto/tests/cases/plans/expr_registry.rs`, 11 tests around a `TagExpr` that decodes through the registry. The two codecs in the file make the paths distinguishable: `RefusingCodec` errors on every method, so any codec involvement is fatal, and `TagCodec` can decode the same payload, so a test that should reach it is not merely observing a failure. - The registry path end to end, including `ctx.task_ctx::<TaskContext>()` reading the session's batch size — the decoded value is what distinguishes the two decode paths everywhere else in the file. - A whole plan through `physical_plan_to_bytes_with_extension_codec` / `physical_plan_from_bytes_with_extension_codec`. - Deduplication under `DeduplicatingProtoConverter`: one expression referenced twice must resolve to one deduplicated expression. Note that `Arc::ptr_eq` on the two decoded expressions is the wrong assertion, because the deserializer returns `cached.with_new_children(..)` — a fresh `Arc`; the test instead ptr-compares shared state that a fresh decode mints and `with_new_children` carries over. Counterfactually validated: reverting `expr_id` to `None` inside `encode_extension` fails this test on exactly that assertion. - Nested extension expressions, proving the registry is re-entered for children. - The three fallback rules: unnamed node, named-but-unregistered name, and no registry at all. - A registered decoder that errors must not silently retry through the codec — the codec in play *can* decode that node, so a fall-through would mask the failure. - Registry unit behavior: duplicate name, empty name, and additive registration on `SessionConfig`. Full extended workspace suite, clippy, fmt, rustdoc, typos, prettier, license headers and toml fmt all pass locally. ## Are there any user-facing changes? Additive; nothing changes for anyone who does not opt in. Documented in the 55.0.0 upgrade guide. Two things worth a reviewer's attention: - **The wire field is additive for binary protobuf, but not for the optional JSON serde** — generated JSON deserializers reject unknown fields, so a node carrying `exprName` needs a reader built after this change. Called out in the upgrade guide. - **The codec fallback is only equivalent for nodes written the old way.** Once an expression migrates its encoding to `encode_extension`, the payload on the wire is that expression's own message. A reader that fails to register the type hands that payload to `try_decode_expr`, and if the codec still recognizes the expression it may decode the new payload as the old message rather than failing outright, since protobuf skips unknown fields and defaults missing ones. Both the rustdoc and the upgrade guide say so, and say to register migrated expressions everywhere their plans are read. Open questions I would like input on: - **Naming.** `PhysicalExprRegistration` for the `SessionConfig` extension trait is provisional; it should probably be settled together with whatever the plan-side companion (#24625) ends up calling its equivalent. - **The turbofish on `task_ctx::<TaskContext>()`.** A typed helper could live in `datafusion-physical-plan`, which can name both types, but that puts a third crate into a story whose selling point is that expression authors only need the leaf crate. Left as the documented turbofish. - The plan side (#24625) is not touched here. Per that issue's own scoping note, the two are worth prioritizing independently. 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
