adriangb opened a new pull request, #24631: URL: https://github.com/apache/datafusion/pull/24631
## Which issue does this PR close? - Closes #24625. ## Rationale for this change #23494 moved built-in `ExecutionPlan` serialization onto per-type `try_to_proto` / `try_from_proto` hooks, but third-party plans could only use half of it. `ExecutionPlan::try_to_proto` is callable by them; there is no way back, because decoding an extension node always routes through `PhysicalExtensionCodec::try_decode`. The blocker is the wire type. `PhysicalExtensionNode` carries `node: Vec<u8>` and `inputs`, and nothing else — no type discriminator, so **the codec *is* the discriminator**. That is why `ComposedPhysicalExtensionCodec` has to try each registered codec in sequence and read a decode error as "not mine": resolution depends on registration order and on error strings, and a name collision between two independent crates is undetectable. The issue's worked example is [datafusion-distributed](https://github.com/datafusion-contrib/datafusion-distributed), which ships six extension plans and serializes every query plan across the network. Its `distributed_codec.rs` is 833 lines holding six `downcast_ref` arms and a matching six-arm `match` — the central dispatch chain #23494 set out to remove, rebuilt downstream — plus a 33-line `user_codec.rs` whose only job is accumulating codecs into a `ComposedPhysicalExtensionCodec` so the library's codec and the user's can coexist. ## What changes are included in this PR? Four commits, each of which builds and tests green on its own. **1. `optional string plan_name = 3;` on `PhysicalExtensionNode`.** Additive and proto3-compatible: old writers omit it, old readers ignore it. **2. A session-scoped decode registry** (`datafusion-physical-plan`, feature `proto`): - `ExecutionPlanEncodeCtx::encode_extension::<T>(payload, children)` — builds the node, stamps the name, and recurses into the children. The name comes from `T::PLAN_NAME`, so it cannot drift from the registration; `encode_extension_named` keeps the `&str` form for the escape hatch below. - `ExecutionPlanDecodeCtx::decode_extension` / `decode_children` — the inverse. - `ExtensionExecutionPlan` — the per-type contract: a namespaced `PLAN_NAME` plus `try_from_proto`. Because `try_from_proto` is a plain inherent function rather than a trait method, there is no `Self`-return object-safety problem and the registry is just `name -> fn(&PhysicalPlanNode, &ExecutionPlanDecodeCtx) -> Result<Arc<dyn ExecutionPlan>>`. - `ExecutionPlanRegistry` + `SessionConfig::register_execution_plan::<T>()`. Registering the same type twice is a no-op (keyed by `TypeId`); registering a *different* type under a taken name is an error, so collisions surface at registration instead of as a mis-decode. `register_decoder(name, fn)` is an escape hatch for names that are not a compile-time constant of one type. The registry is session scoped, matching the `FunctionRegistry` precedent that the issue asked for rather than a global static. It lives in the `SessionConfig` extension map because `datafusion-execution` sits *below* `datafusion-physical-plan` in the crate graph and so cannot name the decoder fn type — the extension map is the one hole through that layering, and it means the registry reaches every `TaskContext` derived from the session with no new plumbing. **3. Decode dispatch.** The rule is the one function decode already uses a layer down (payload → codec; else registry → codec fallback): - `plan_name` present **and** registered on the decoding session → the plan's own `try_from_proto`, no codec involved. - anything else → the `PhysicalExtensionCodec` chain, unchanged. A failure *inside* a registered decoder is fatal rather than falling through to the codec — falling through is what reintroduces the mis-decode this replaces. When a name is present but unregistered and the codec also fails, the error names the missing registration and lists what the session does know, instead of leaving only the codec's generic "unsupported plan". **4. Docs** — an upgrade-guide entry. `PhysicalExtensionCodec` is **not** retired: extension `PhysicalExpr`s and UDF payloads still need it, and the codec fallback stays for unmigrated plans. Companion issue for the expression side is #24626. ## Are these changes tested? Yes. 20 unit tests in `datafusion-physical-plan` and 10 integration tests in `datafusion/proto/tests/cases/plans/extensions.rs`: - Registration: idempotent for the same type, an error for a different type under the same name, empty names rejected, registrations accumulate on and survive a clone of `SessionConfig`. - A plan round-tripping with `DefaultPhysicalExtensionCodec` — which decodes nothing — so whatever survives came from the registry. - A **session-dependent** plan modeled on `NetworkShuffleExec`: its worker pool is rebuilt from the *decoding* session, proving `ExecutionPlanDecodeCtx::task_ctx()` covers real-world extension plans and that the pool never rode the wire. - Two independent crates' plans nested inside each other, decoding with no `ComposedPhysicalExtensionCodec`. - The registry taking precedence over a codec that claims every payload, and a registered decoder's failure *not* falling back to that codec. - Back-compat both ways: codec-encoded plans stay anonymous on the wire and decode as before; a name the session does not know still decodes if a codec understands the payload; a node re-encoded without the field reads back as unnamed. - A bytes-level `encode_to_vec` / `decode` hop, so tag 3 is exercised as wire bytes and not just as a struct field. Full extended suite (`--features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`), workspace `clippy -D warnings` and `cargo fmt --check` are all green. ## Are there any user-facing changes? New public API, all opt-in; no behavior changes for existing code. The one mechanical break is that `PhysicalExtensionNode` gained a field, so exhaustive struct literals must set `plan_name` (`None` reproduces the previous behavior). Documented in the 55.0.0 upgrade guide. -- 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]
