feat: decode extension PhysicalExprs through a session-scoped registry

`PhysicalExtensionExprNode` carried no type discriminator, so the
`PhysicalExtensionCodec` *was* the discriminator, with the same problems the
plan side has: composition is by registration order, and a name collision
between two independent crates is undetectable.

An extension expression now declares its wire name with `ExtensionExprName`,
implements `PhysicalExprFromProto` like a built-in, writes itself with
`PhysicalExprEncodeCtx::extension_expr_node` (which stamps the name, so the
encoded name and the registry key cannot drift apart), and is registered with
`register_physical_expr` into the `ProtoDecoderRegistry` that the decoding
session carries — the same registry extension `ExecutionPlan`s use, keyed on
the trait as well as the name.

The name is a separate trait from the decode contract so that only extension
expressions can be registered: a built-in has an `ExprType` variant of its
own, and `register_physical_expr::<Column>(..)` does not compile. It lives in
`physical-expr-common`, beside the encode context that stamps it, because
that crate sits below the one that owns the decode contract;
`datafusion_physical_expr::proto` re-exports it.

Decode rule: a named node the session has a decoder for goes to that decoder,
and a failure there is fatal (falling through would let a codec decode a
payload it never wrote); anything else takes the `PhysicalExtensionCodec`
chain exactly as before. If the codec fails too, the missing registration is
added as *context* on the codec's own error, so the kind the codec chose
still reaches a caller that matches on it.

`decode_physical_expr` reads the name off the node rather than taking it as
an argument, so a caller cannot pair a node with a name it does not carry.

Documented in the 56.0.0 upgrade guide; tests in
`datafusion/proto/tests/cases/plans/expr_registry.rs`.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: adriangb

datafusion/physical-expr-common/src/physical_expr.rs

@@ -555,10 +555,36 @@ pub mod proto_encode {
     use std::sync::Arc;
 
     use datafusion_common::Result;
-    use datafusion_proto_models::protobuf::PhysicalExprNode;
+    use datafusion_proto_models::protobuf::physical_expr_node::ExprType;
+    use datafusion_proto_models::protobuf::{
+        PhysicalExprNode, PhysicalExtensionExprNode,
+    };
 
     use super::PhysicalExpr;
 
+    /// The wire name of an extension [`PhysicalExpr`], and the key it is
+    /// registered and dispatched by.
+    ///
+    /// Only extension expressions implement this. A built-in expression has an
+    /// `ExprType` variant of its own and is dispatched by that variant, so it
+    /// needs no name and cannot be registered.
+    ///
+    /// It lives here, beside [`PhysicalExprEncodeCtx`], because the encode
+    /// helper that stamps the name needs it and this crate sits below the one
+    /// that owns the decode contract. `datafusion_physical_expr::proto`
+    /// re-exports it, next to `PhysicalExprFromProto`, which is where an
+    /// expression author should import it from.
+    pub trait ExtensionExprName: PhysicalExpr {
+        /// The name this expression type is written and registered under.
+        ///
+        /// Namespace it with the owning crate — `"my_crate.MyExpr"`, not
+        /// `"MyExpr"` — so that two independent crates registering into the
+        /// same session collide at registration time instead of silently
+        /// decoding each other's nodes. Never write it by hand on the wire:
+        /// [`PhysicalExprEncodeCtx::extension_expr_node`] stamps it for you.
+        const NAME: &'static str;
+    }
+
     /// Encoder context handed to [`super::PhysicalExpr::try_to_proto`].
     ///
     /// Wraps an internal [`PhysicalExprEncode`] trait object so callers see a
@@ -584,6 +610,40 @@ pub mod proto_encode {
             self.encoder.encode(expr)
         }
 
+        /// Build the `PhysicalExprNode` for an extension expression `T`: an
+        /// `Extension` variant carrying `payload`, the encoded `children`, and
+        /// `T`'s [`NAME`](ExtensionExprName::NAME).
+        ///
+        /// The one supported way for an extension expression to write itself.
+        /// The name on the wire is taken from the same const the registry is
+        /// keyed by, so an encoder cannot write a name no decoder answers to:
+        ///
+        /// ```ignore
+        /// fn try_to_proto(
+        ///     &self,
+        ///     ctx: &PhysicalExprEncodeCtx<'_>,
+        /// ) -> Result<Option<PhysicalExprNode>> {
+        ///     Ok(Some(ctx.extension_expr_node::<Self>(
+        ///         self.state.to_bytes()?,
+        ///         self.children(),
+        ///     )?))
+        /// }
+        /// ```
+        pub fn extension_expr_node<T: ExtensionExprName>(
+            &self,
+            payload: Vec<u8>,
+            children: Vec<&Arc<dyn PhysicalExpr>>,
+        ) -> Result<PhysicalExprNode> {
+            Ok(PhysicalExprNode {
+                expr_type: Some(ExprType::Extension(PhysicalExtensionExprNode {
+                    expr: payload,
+                    inputs: self.encode_children_expressions(children)?,
+                    expr_name: Some(T::NAME.to_string()),
+                })),
+                ..Default::default()
+            })
+        }
+
         /// Encode a sequence of child expressions, preserving order.
         ///
         /// Convenience wrapper over [`Self::encode_child`] for expressions

datafusion/physical-expr/src/proto.rs

@@ -16,7 +16,13 @@
 // under the License.
 
 //! The decode half of [`PhysicalExpr::try_to_proto`]: [`PhysicalExprFromProto`],
-//! the contract every self-serializing expression implements.
+//! the contract every self-serializing expression implements, plus the
+//! `PhysicalExpr` face of the session-scoped [`ProtoDecoderRegistry`] that
+//! routes extension expressions to theirs by name —
+//! [`register_physical_expr`], [`decode_physical_expr`] and
+//! [`physical_expr_names`], all keyed on `dyn PhysicalExpr`. The store itself
+//! is shared with every other extension kind, so a session carries one
+//! registry.
 //!
 //! These live here rather than next to the encode context in
 //! `datafusion-physical-expr-common` because a decoder may need the session,
@@ -33,7 +39,12 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_expr::registry::FunctionRegistry;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 pub use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+pub use datafusion_physical_expr_common::physical_expr::proto_encode::{
+    ExtensionExprName, PhysicalExprEncodeCtx,
+};
+pub use datafusion_proto_models::ProtoDecoderRegistry;
 use datafusion_proto_models::protobuf::PhysicalExprNode;
+use datafusion_proto_models::protobuf::physical_expr_node::ExprType;
 
 /// What an expression decoder can see of the session it runs under: the
 /// function registry, for resolving UDFs by name, and the configuration
@@ -81,6 +92,10 @@ impl std::fmt::Debug for ExprDecodeSession<'_> {
 /// Every self-serializing expression implements this. Built-in expressions
 /// are dispatched to their `try_from_proto` by their `ExprType` variant.
 ///
+/// An extension expression has no variant of its own, so it is dispatched by
+/// name instead and implements [`ExtensionExprName`] as well.
+/// [`register_physical_expr`] carries the worked example.
+///
 /// ```ignore
 /// impl PhysicalExprFromProto for MyExpr {
 ///     fn try_from_proto(
@@ -114,3 +129,105 @@ pub trait PhysicalExprFromProto: PhysicalExpr + Sized {
         ctx: &PhysicalExprDecodeCtx<'_, ExprDecodeSession<'_>>,
     ) -> Result<Arc<dyn PhysicalExpr>>;
 }
+
+/// How this facade stores a decoder in the shared registry: a function pointer
+/// to the monomorphized [`PhysicalExprFromProto::try_from_proto`].
+///
+/// Deliberately private, and the same type on both the
+/// [`register_physical_expr`] and the [`decode_physical_expr`] side.
+/// [`PhysicalExprFromProto`] is the public contract and
+/// [`decode_physical_expr`] is the public way to invoke one, so this can become
+/// something else — a `dyn` decoder object, to admit stateful or closure
+/// decoders, which is what an FFI decoder needs — without a breaking change.
+type PhysicalExprDecoder = fn(
+    &PhysicalExprNode,
+    &PhysicalExprDecodeCtx<'_, ExprDecodeSession<'_>>,
+) -> Result<Arc<dyn PhysicalExpr>>;
+
+/// Register `T` in `registry` under its [`ExtensionExprName::NAME`].
+///
+/// Registering the same type twice is a no-op. Registering a *different* type
+/// under a name already taken is an error, so collisions surface here rather
+/// than as a wrong decode later. The name is scoped to `dyn PhysicalExpr`, so a
+/// plan or a data source may use the same name in the same registry.
+///
+/// # Who builds the registry
+///
+/// The application that owns the session builds it. A session carries at most
+/// one registry, and `SessionConfig::with_extension` replaces what is there, so
+/// a library must never attach a registry of its own: it would silently discard
+/// another library's. A library exposes a function that fills a registry it is
+/// handed, and the application composes them. One registry holds every kind, so
+/// a library registers its plans and its expressions into the same object.
+///
+/// # A built-in expression cannot be registered
+///
+/// The bound is [`ExtensionExprName`], which only extension expressions
+/// implement. A built-in implements the decode contract:
+///
+/// ```
+/// use datafusion_physical_expr::expressions::Column;
+/// use datafusion_physical_expr::proto::PhysicalExprFromProto;
+///
+/// let _ = <Column as PhysicalExprFromProto>::try_from_proto;
+/// ```
+///
+/// but it has no wire name, and registering it does not compile:
+///
+/// ```compile_fail
+/// use datafusion_physical_expr::expressions::Column;
+/// use datafusion_physical_expr::proto::register_physical_expr;
+/// use datafusion_proto_models::ProtoDecoderRegistry;
+///
+/// let mut registry = ProtoDecoderRegistry::new();
+/// register_physical_expr::<Column>(&mut registry).unwrap();
+/// ```
+pub fn register_physical_expr<T: ExtensionExprName + PhysicalExprFromProto>(
+    registry: &mut ProtoDecoderRegistry,
+) -> Result<()> {
+    registry.register_decoder::<dyn PhysicalExpr, T, PhysicalExprDecoder>(
+        T::NAME,
+        T::try_from_proto,
+    )
+}
+
+/// Decode `node` with the extension expression decoder registered under the
+/// name `node` carries.
+///
+/// The name is read from the node's `Extension` variant rather than passed in,
+/// so a caller cannot pair a node with a name it does not carry.
+///
+/// `None` means "this node names no extension expression decoder of ours": it
+/// is not an extension node, it carries no name, or no registered name matches.
+/// The caller then falls back to the `PhysicalExtensionCodec` chain.
+/// `Some(Err(..))` means the decoder that *does* own the name failed, which is
+/// fatal: falling back there would let a codec decode the payload wrongly, the
+/// very thing the name exists to prevent.
+pub fn decode_physical_expr(
+    registry: &ProtoDecoderRegistry,
+    node: &PhysicalExprNode,
+    ctx: &PhysicalExprDecodeCtx<'_, ExprDecodeSession<'_>>,
+) -> Option<Result<Arc<dyn PhysicalExpr>>> {
+    let name = expr_name(node)?;
+    let decoder = registry.decoder::<dyn PhysicalExpr, PhysicalExprDecoder>(name)?;
+    Some(decoder(node, ctx))
+}
+
+/// Every extension expression name registered in `registry`, in arbitrary
+/// order.
+///
+/// Names registered for another kind are not included.
+pub fn physical_expr_names(
+    registry: &ProtoDecoderRegistry,
+) -> impl Iterator<Item = &str> {
+    registry.names::<dyn PhysicalExpr>()
+}
+
+/// The registry name `node` was written with, if it is an extension node that
+/// carries one.
+fn expr_name(node: &PhysicalExprNode) -> Option<&str> {
+    match node.expr_type.as_ref()? {
+        ExprType::Extension(extension) => extension.expr_name.as_deref(),
+        _ => None,
+    }
+}

datafusion/physical-plan/src/proto/mod.rs

@@ -79,7 +79,10 @@ use datafusion_execution::TaskContext;
 use datafusion_expr::physical_planning_context::ScalarSubqueryResults;
 use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
 use datafusion_physical_expr::PhysicalExpr;
-pub use datafusion_physical_expr::proto::{ExprDecodeSession, PhysicalExprFromProto};
+pub use datafusion_physical_expr::proto::{
+    ExprDecodeSession, ExtensionExprName, PhysicalExprFromProto, decode_physical_expr,
+    physical_expr_names, register_physical_expr,
+};
 use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecode;
 pub use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
 use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncode;

datafusion/proto/src/physical_plan/from_proto.rs

@@ -23,7 +23,7 @@ use arrow::array::RecordBatch;
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Field, Schema};
 use arrow::ipc::reader::StreamReader;
-use datafusion_common::{Result, internal_datafusion_err, not_impl_err};
+use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err};
 use datafusion_datasource::TableSchema;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
@@ -41,7 +41,10 @@ use datafusion_physical_plan::expressions::{
 };
 use datafusion_physical_plan::joins::HashExpr;
 use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx;
-use datafusion_physical_plan::proto::{ExprDecodeSession, PhysicalExprFromProto};
+use datafusion_physical_plan::proto::{
+    ExprDecodeSession, PhysicalExprFromProto, ProtoDecoderRegistry, decode_physical_expr,
+    physical_expr_names,
+};
 use datafusion_physical_plan::repartition::RangeExpr;
 use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field};
 use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr};
@@ -364,16 +367,42 @@ pub fn parse_physical_expr_with_converter(
             SqlSimilarToPattern::try_from_proto(proto, &decode_ctx)?
         }
         ExprType::Extension(extension) => {
+            // Lookup-order policy, the same one extension plans use: an
+            // expression that named itself on the wire and is registered on
+            // this session decodes itself, no codec involved. Anything else —
+            // no name, or a name this session has no decoder for — takes the
+            // codec path, which is also every node written before `expr_name`
+            // existed. `None` from the registry means no decoder claims the
+            // name; a decode *failure* is returned as-is rather than falling
+            // through to the codec.
+            let registry = ctx
+                .task_ctx()
+                .session_config()
+                .get_extension::<ProtoDecoderRegistry>();
+            if let Some(registry) = registry.as_ref()
+                && let Some(decoded) = decode_physical_expr(registry, proto, &decode_ctx)
+            {
+                return decoded;
+            }
+
             let inputs: Vec<Arc<dyn PhysicalExpr>> = extension
                 .inputs
                 .iter()
                 .map(|e| proto_converter.proto_to_physical_expr(e, input_schema, ctx))
                 .collect::<Result<_>>()?;
-            ctx.codec().try_decode_expr(
-                extension.expr.as_slice(),
-                &inputs,
-                &decode_ctx,
-            )? as _
+            ctx.codec()
+                .try_decode_expr(extension.expr.as_slice(), &inputs, &decode_ctx)
+                .map_err(|e| match extension.expr_name.as_deref() {
+                    // The writer named the expression but this session has no
+                    // decoder for it: say so, rather than leaving only the
+                    // codec's error to explain a missing registration.
+                    Some(expr_name) => unregistered_extension_expr_context(
+                        expr_name,
+                        registry.as_deref(),
+                        e,
+                    ),
+                    None => e,
+                })?
         }
         ExprType::Lambda(_) => LambdaExpr::try_from_proto(proto, &decode_ctx)?,
         ExprType::LambdaVariable(_) => {
@@ -483,6 +512,36 @@ struct ConverterDecoder<'a, 'b> {
     proto_converter: &'a dyn PhysicalProtoConverterExtension,
 }
 
+/// Error for a `PhysicalExtensionExprNode` that names its expression type but
+/// finds no decoder registered for that name.
+///
+/// Lists what *is* registered, because the usual cause is a session configured
+/// on the writing side but not on the reading one. The codec's own error is
+/// kept as the cause rather than reworded: the codec may have failed for a
+/// reason that has nothing to do with the missing registration, and a caller
+/// matching on the error kind must still see the kind the codec chose.
+fn unregistered_extension_expr_context(
+    expr_name: &str,
+    registry: Option<&ProtoDecoderRegistry>,
+    codec_error: DataFusionError,
+) -> DataFusionError {
+    let mut registered: Vec<&str> = registry
+        .map(|registry| physical_expr_names(registry).collect())
+        .unwrap_or_default();
+    registered.sort_unstable();
+    let registered = if registered.is_empty() {
+        "none".to_string()
+    } else {
+        registered.join(", ")
+    };
+    codec_error.context(format!(
+        "No decoder is registered for the extension PhysicalExpr '{expr_name}'. Register \
+         the expression in the ProtoDecoderRegistry attached to the decoding session's \
+         SessionConfig, or supply a PhysicalExtensionCodec that handles it. Registered \
+         extension expressions: {registered}"
+    ))
+}
+
 impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecode
     for ConverterDecoder<'_, '_>
 {

datafusion/proto/src/physical_plan/mod.rs

@@ -1745,6 +1745,13 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any {
 
     /// Decode a custom extension expression from `buf`.
     ///
+    /// Expressions can instead name themselves on the wire and be decoded by a
+    /// per-type decoder registered on the session — see
+    /// [`PhysicalExprFromProto`] and [`register_physical_expr`]. That path
+    /// resolves by name rather than by codec registration order, so two crates
+    /// claiming the same name collide at registration. This method stays the
+    /// fallback for everything unnamed or unregistered.
+    ///
     /// `inputs` holds the already-decoded children carried in the
     /// `PhysicalExtensionExprNode.inputs` field. If the codec instead embeds
     /// nested `PhysicalExprNode`s *inside* `buf`, decode them through
@@ -1757,6 +1764,8 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any {
     /// cache-hits on its `expr_id` and re-shares one `Arc<dyn PhysicalExpr>`.
     ///
     /// [`parse_physical_expr`]: crate::physical_plan::from_proto::parse_physical_expr
+    /// [`PhysicalExprFromProto`]: datafusion_physical_plan::proto::PhysicalExprFromProto
+    /// [`register_physical_expr`]: datafusion_physical_plan::proto::register_physical_expr
     fn try_decode_expr(
         &self,
         _buf: &[u8],