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 implements `ExtensionExprFromProto`, which pairs
the name it is dispatched by with the constructor that rebuilds it. Its
`try_to_proto` writes itself with `extension_expr_node`, which stamps
`Self::NAME`, so the encoded name and the registry key cannot drift apart.
`register_physical_expr` puts its decoder in the `ProtoDecoderRegistry` the
decoding session carries — the same registry extension `ExecutionPlan`s use,
keyed on the trait as well as the name.

Only extension expressions implement the trait. A built-in has an `ExprType`
variant of its own, keeps its inherent `try_from_proto`, and has no wire name
to be registered under — so `register_physical_expr::<Column>(..)` does not
compile. Nothing about built-in serialization changes, and no caller needs a
new import.

`extension_expr_node` is a free function rather than a method on
`PhysicalExprEncodeCtx` because that context lives in
`datafusion-physical-expr-common`, below the crate that can name
`ExtensionExprFromProto` — the same layering that makes the decode context
generic over its session.

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/src/proto.rs

@@ -15,7 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! The session side of decoding a `PhysicalExpr` from protobuf.
+//! The session an expression decode runs under, and the `PhysicalExpr`
+//! face of the session-scoped [`ProtoDecoderRegistry`] that routes extension
+//! expressions to their decoders 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.
 //!
 //! The decode context in `datafusion-physical-expr-common` is generic over its
 //! session type: a decoder may need the session, and what it needs —
@@ -25,9 +30,17 @@
 //! expressions decode under it. `datafusion-proto` builds one from the
 //! `TaskContext` it decodes under.
 
+use std::sync::Arc;
+
+use datafusion_common::Result;
 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::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
@@ -68,3 +81,206 @@ impl std::fmt::Debug for ExprDecodeSession<'_> {
         f.debug_struct("ExprDecodeSession").finish_non_exhaustive()
     }
 }
+/// The wire name of an extension [`PhysicalExpr`], and the constructor that
+/// rebuilds it.
+///
+/// Only extension expressions implement this. A built-in expression has an
+/// `ExprType` variant of its own, keeps the inherent `try_from_proto`
+/// dispatched from that variant, and has no wire name to be registered under —
+/// so [`register_physical_expr`] will not accept one.
+///
+/// One impl block says what the expression is called and how it decodes; its
+/// `try_to_proto` on [`PhysicalExpr`] writes the matching node with
+/// [`extension_expr_node`]:
+///
+/// ```ignore
+/// impl PhysicalExpr for MyExpr {
+///     // ...
+///     fn try_to_proto(
+///         &self,
+///         ctx: &PhysicalExprEncodeCtx<'_>,
+///     ) -> Result<Option<PhysicalExprNode>> {
+///         // Stamps `Self::NAME`, so the encoded name and the registry key
+///         // cannot drift apart.
+///         Ok(Some(extension_expr_node::<Self>(
+///             ctx,
+///             self.state.to_bytes()?,
+///             self.children(),
+///         )?))
+///     }
+/// }
+///
+/// impl ExtensionExprFromProto for MyExpr {
+///     // Namespace the name with the owning crate, so a collision with
+///     // another crate is an error at registration and not a wrong decode.
+///     const NAME: &'static str = "my_crate.MyExpr";
+///
+///     fn try_from_proto(
+///         node: &PhysicalExprNode,
+///         ctx: &PhysicalExprDecodeCtx<'_, ExprDecodeSession<'_>>,
+///     ) -> Result<Arc<dyn PhysicalExpr>> {
+///         let extension = expect_expr_variant!(
+///             node,
+///             physical_expr_node::ExprType::Extension,
+///             "Extension",
+///         );
+///         let state = MyState::from_bytes(&extension.expr)?;
+///         let children = ctx.decode_children_expressions(&extension.inputs)?;
+///         Ok(Arc::new(MyExpr::new(state, children)))
+///     }
+/// }
+///
+/// let mut registry = ProtoDecoderRegistry::new();
+/// register_physical_expr::<MyExpr>(&mut registry)?;
+/// let config = SessionConfig::new().with_extension(Arc::new(registry));
+/// ```
+pub trait ExtensionExprFromProto: PhysicalExpr + Sized {
+    /// 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:
+    /// [`extension_expr_node`] stamps it for you.
+    const NAME: &'static str;
+
+    /// Rebuild the expression from the `PhysicalExprNode` its
+    /// [`PhysicalExpr::try_to_proto`] wrote.
+    ///
+    /// Takes the whole node — the exact inverse of what `try_to_proto`
+    /// returns — so the constructor can also see outer-node fields such as
+    /// `expr_id`. `ctx.session()` gives the function registry and the
+    /// configuration options.
+    fn try_from_proto(
+        node: &PhysicalExprNode,
+        ctx: &PhysicalExprDecodeCtx<'_, ExprDecodeSession<'_>>,
+    ) -> Result<Arc<dyn PhysicalExpr>>;
+}
+
+/// Build the `PhysicalExprNode` for an extension expression `T`: an
+/// `Extension` variant carrying `payload`, the encoded `children`, and `T`'s
+/// [`NAME`](ExtensionExprFromProto::NAME).
+///
+/// The one supported way for an extension expression to write itself. It is a
+/// free function rather than a method on [`PhysicalExprEncodeCtx`] because the
+/// context lives in `datafusion-physical-expr-common`, below the crate that
+/// can name [`ExtensionExprFromProto`] — the same layering that makes the
+/// decode context generic over its session.
+pub fn extension_expr_node<T: ExtensionExprFromProto>(
+    ctx: &PhysicalExprEncodeCtx<'_>,
+    payload: Vec<u8>,
+    children: Vec<&Arc<dyn PhysicalExpr>>,
+) -> Result<PhysicalExprNode> {
+    Ok(PhysicalExprNode {
+        expr_type: Some(ExprType::Extension(
+            datafusion_proto_models::protobuf::PhysicalExtensionExprNode {
+                expr: payload,
+                inputs: ctx.encode_children_expressions(children)?,
+                expr_name: Some(T::NAME.to_string()),
+            },
+        )),
+        ..Default::default()
+    })
+}
+
+/// How this facade stores a decoder in the shared registry: a function pointer
+/// to the monomorphized [`ExtensionExprFromProto::try_from_proto`].
+///
+/// Deliberately private, and the same type on both the
+/// [`register_physical_expr`] and the [`decode_physical_expr`] side.
+/// [`ExtensionExprFromProto`] 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 [`ExtensionExprFromProto::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 [`ExtensionExprFromProto`], which only extension expressions
+/// implement. A built-in keeps an inherent `try_from_proto`, dispatched from
+/// its own `ExprType` variant:
+///
+/// ```
+/// use datafusion_physical_expr::expressions::Column;
+///
+/// let _ = Column::try_from_proto;
+/// ```
+///
+/// but it is not an extension, 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: ExtensionExprFromProto>(
+    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;
+pub use datafusion_physical_expr::proto::{
+    ExprDecodeSession, ExtensionExprFromProto, decode_physical_expr, extension_expr_node,
+    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,9 @@ use datafusion_physical_plan::expressions::{
 };
 use datafusion_physical_plan::joins::HashExpr;
 use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx;
-use datafusion_physical_plan::proto::ExprDecodeSession;
+use datafusion_physical_plan::proto::{
+    ExprDecodeSession, 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 +366,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 +511,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

@@ -1749,6 +1749,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
@@ -1761,6 +1768,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],