You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/21 12:35:39 UTC

(datafusion) branch gh-readonly-queue/main/pr-26353-6c1727fcd9230dda428f8c3a2397fce235544bd5 created (now 3bc3d0bb82)

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-26353-6c1727fcd9230dda428f8c3a2397fce235544bd5
in repository https://gitbox.apache.org/repos/asf/datafusion.git


      at 3bc3d0bb82 feat: add `SessionConfig::update_extension` so libraries can share one `ProtoDecoderRegistry` (#26353)

This branch includes the following new commits:

     new 3bc3d0bb82 feat: add `SessionConfig::update_extension` so libraries can share one `ProtoDecoderRegistry` (#26353)

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.



---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org


(datafusion) 01/01: feat: add `SessionConfig::update_extension` so libraries can share one `ProtoDecoderRegistry` (#26353)

Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26353-6c1727fcd9230dda428f8c3a2397fce235544bd5
in repository https://gitbox.apache.org/repos/asf/datafusion.git

commit 3bc3d0bb822f084a3195d8a2a42bc5e7e2c92010
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Wed Oct 21 12:35:13 2026 +0000

    feat: add `SessionConfig::update_extension` so libraries can share one `ProtoDecoderRegistry` (#26353)
    
    ## Which issue does this PR close?
    
    - Follow-up to https://github.com/apache/datafusion/pull/24631. No
    separate issue.
    
    ## Rationale for this change
    
    A session carries one `ProtoDecoderRegistry`, as a `SessionConfig`
    extension. Today the only way to attach it is
    `SessionConfig::with_extension`, which **replaces** the extension. If
    two libraries each register extension decoders, the second one silently
    discards the decoders of the first one. The failure shows only later, at
    decode time, as "No decoder is registered".
    
    We found this when we moved datafusion-distributed to the registry API.
    Distributed registers its plans in the worker `SessionConfig`, then
    gives the config to the user's session builder. A user who attaches
    their own registry there removes the distributed decoders.
    
    ```rust
    // Before: library B discards library A's decoders
    config = config.with_extension(Arc::new(registry_a));
    config = config.with_extension(Arc::new(registry_b));
    ```
    
    ## What changes are included in this PR?
    
    A new generic method, `SessionConfig::update_extension`. It changes an
    extension in place, and creates it with `Default` if it is not there. If
    another config shares the extension, it is cloned first
    (`Arc::make_mut`), so that config does not change.
    
    ```rust
    // After: each library adds to the same registry, in any order
    config.update_extension(|registry: &mut ProtoDecoderRegistry| {
        register_execution_plan::<MyExec>(registry)
    })?;
    ```
    
    The method is generic because `datafusion-execution` does not depend on
    `datafusion-proto-models`, where the registry is defined. It is useful
    for any extension that more than one part of an application writes to.
    
    The docs for `ProtoDecoderRegistry` and `register_execution_plan`, and
    the unreleased 56.0.0 upgrade guide, now recommend `update_extension`
    and say not to use `with_extension` for the registry.
    
    ## What is the testing strategy for this PR?
    
    - `proto::registry::tests::libraries_add_to_one_session_registry`: two
    "libraries" register different plans through `update_extension`. Both
    names stay registered. A config cloned between the two calls keeps only
    the first name.
    - A doc test on `update_extension`.
    
    ## Are there any user-facing changes?
    
    Yes, one new public method on `SessionConfig`. It is additive and does
    not break anything. No behavior changes for existing callers of
    `with_extension`.
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    Co-authored-by: Claude Opus 5.5 <no...@anthropic.com>
---
 datafusion/execution/src/config.rs                 | 33 ++++++++++++++++++
 datafusion/physical-plan/src/proto/registry.rs     | 40 ++++++++++++++++++++--
 datafusion/proto-models/src/registry.rs            | 26 +++++++-------
 docs/source/library-user-guide/upgrading/56.0.0.md | 22 ++++++------
 4 files changed, 96 insertions(+), 25 deletions(-)

diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs
index dbfd9bf1d8..a157ae2175 100644
--- a/datafusion/execution/src/config.rs
+++ b/datafusion/execution/src/config.rs
@@ -681,6 +681,39 @@ impl SessionConfig {
     {
         self.extensions.get_arc::<T>()
     }
+
+    /// Change the extension of type `T` in place, and return what `f` returns.
+    ///
+    /// If there is no extension of type `T`, `f` gets `T::default()`, which is
+    /// then attached. If the extension is shared with another config, it is
+    /// cloned first, so the other config does not change.
+    ///
+    /// Use this, not [`with_extension`](Self::with_extension), when more than
+    /// one part of an application adds to the same extension. `with_extension`
+    /// replaces the extension, so it discards what other parts added.
+    ///
+    /// # Example
+    /// ```
+    /// use datafusion_execution::config::SessionConfig;
+    ///
+    /// #[derive(Clone, Default)]
+    /// struct Names(Vec<&'static str>);
+    ///
+    /// let mut cfg = SessionConfig::default();
+    /// cfg.update_extension(|names: &mut Names| names.0.push("a"));
+    /// cfg.update_extension(|names: &mut Names| names.0.push("b"));
+    ///
+    /// assert_eq!(cfg.get_extension::<Names>().unwrap().0, ["a", "b"]);
+    /// ```
+    pub fn update_extension<T, R>(&mut self, f: impl FnOnce(&mut T) -> R) -> R
+    where
+        T: Clone + Default + Send + Sync + 'static,
+    {
+        let mut ext = self.get_extension::<T>().unwrap_or_default();
+        let result = f(Arc::make_mut(&mut ext));
+        self.set_extension(ext);
+        result
+    }
 }
 
 impl From<ConfigOptions> for SessionConfig {
diff --git a/datafusion/physical-plan/src/proto/registry.rs b/datafusion/physical-plan/src/proto/registry.rs
index 4f12f89e6f..4fa9ddda2e 100644
--- a/datafusion/physical-plan/src/proto/registry.rs
+++ b/datafusion/physical-plan/src/proto/registry.rs
@@ -122,9 +122,9 @@ type ExecutionPlanDecoder =
 /// decode. The name is scoped to `dyn ExecutionPlan`, so an expression or a
 /// data source may use the same name in the same registry.
 ///
-/// The application that owns the session builds the registry and attaches it
-/// once with `SessionConfig::with_extension`; a library exposes a function
-/// that fills a registry it is handed. See [`ProtoDecoderRegistry`].
+/// Add to the session's registry with `SessionConfig::update_extension`, so
+/// that the decoders other libraries added stay registered. See
+/// [`ProtoDecoderRegistry`].
 ///
 /// A built-in plan is not an extension and cannot be registered:
 ///
@@ -208,6 +208,40 @@ mod tests {
         OtherRegisteredExec, RegisteredExec, RegisteredExecClone, UnnamedExec,
     };
 
+    /// Two libraries add to one session's registry, and neither discards the
+    /// other's decoders. A config cloned before the second one adds is not
+    /// changed.
+    #[test]
+    fn libraries_add_to_one_session_registry() -> Result<()> {
+        use datafusion_execution::config::SessionConfig;
+
+        let mut config = SessionConfig::new();
+        config.update_extension(|registry: &mut ProtoDecoderRegistry| {
+            register_execution_plan::<RegisteredExec>(registry)
+        })?;
+        let before = config.clone();
+        config.update_extension(|registry: &mut ProtoDecoderRegistry| {
+            register_execution_plan::<OtherRegisteredExec>(registry)
+        })?;
+
+        let names = |config: &SessionConfig| {
+            let registry = config.get_extension::<ProtoDecoderRegistry>().unwrap();
+            let mut names: Vec<String> = execution_plan_names(&registry)
+                .map(str::to_string)
+                .collect();
+            names.sort();
+            names
+        };
+        let mut expected = vec![
+            RegisteredExec::NAME.to_string(),
+            OtherRegisteredExec::NAME.to_string(),
+        ];
+        expected.sort();
+        assert_eq!(names(&config), expected);
+        assert_eq!(names(&before), [RegisteredExec::NAME]);
+        Ok(())
+    }
+
     #[test]
     fn register_is_idempotent_for_the_same_type() -> Result<()> {
         let mut registry = ProtoDecoderRegistry::new();
diff --git a/datafusion/proto-models/src/registry.rs b/datafusion/proto-models/src/registry.rs
index e2852e57c6..8db237bd70 100644
--- a/datafusion/proto-models/src/registry.rs
+++ b/datafusion/proto-models/src/registry.rs
@@ -56,26 +56,28 @@ struct Registered {
 /// A store of extension decoders, keyed by the trait they produce and the name
 /// the encoder wrote on the wire.
 ///
-/// Build one, let every library that ships extension types register into it,
-/// and attach it to the decoding session with `SessionConfig::with_extension`.
-/// A session carries at most one, and attaching another replaces it, so the
-/// application that owns the session is the one that builds it. A library
-/// exposes a function that fills a registry it is handed:
+/// A session carries one registry, as a `SessionConfig` extension. Every
+/// library that ships extension types adds its decoders to that registry with
+/// `SessionConfig::update_extension`, which creates the registry if there is
+/// none:
 ///
 /// ```rust,ignore
 /// // in each library
-/// pub fn register(registry: &mut ProtoDecoderRegistry) -> Result<()> {
-///     register_execution_plan::<MyExec>(registry)?;
-///     register_physical_expr::<MyExpr>(registry)
+/// pub fn register(config: &mut SessionConfig) -> Result<()> {
+///     config.update_extension(|registry: &mut ProtoDecoderRegistry| {
+///         register_execution_plan::<MyExec>(registry)
+///     })
 /// }
 ///
 /// // in the application that owns the session
-/// let mut registry = ProtoDecoderRegistry::new();
-/// lib_one::register(&mut registry)?;
-/// lib_two::register(&mut registry)?;
-/// let config = SessionConfig::new().with_extension(Arc::new(registry));
+/// let mut config = SessionConfig::new();
+/// lib_one::register(&mut config)?;
+/// lib_two::register(&mut config)?;
 /// ```
 ///
+/// Do not attach a registry with `SessionConfig::with_extension`: it replaces
+/// the registry, so it discards the decoders that other libraries added.
+///
 /// The libraries need to know nothing about each other, and the order they are
 /// called in does not matter: lookup is by name, and a real collision is an
 /// error from [`register_decoder`](Self::register_decoder).
diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md
index 1e8d9e2daa..2ec01c9947 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -1170,28 +1170,30 @@ impl ExtensionPlanFromProto for MyExec {
 
 // On every session that decodes the plan — for a distributed engine, the
 // workers as well as the coordinator.
-let mut registry = ProtoDecoderRegistry::new();
-register_execution_plan::<MyExec>(&mut registry)?;
-let config = SessionConfig::new().with_extension(Arc::new(registry));
+let mut config = SessionConfig::new();
+config.update_extension(|registry: &mut ProtoDecoderRegistry| {
+    register_execution_plan::<MyExec>(registry)
+})?;
 ```
 
 Only extension plans implement `ExtensionPlanFromProto`, and `register_execution_plan` takes that trait as its bound. A built-in plan has a wire variant of its own and keeps its inherent `try_from_proto`, so `register_execution_plan::<FilterExec>(&mut registry)` does not compile.
 
 `ProtoDecoderRegistry` is shared by every kind of extension DataFusion can decode, not only plans. Its key is the trait plus the name, so a plan and an expression may use the same name, and one session carries one registry rather than one per kind.
 
-The application that owns the session builds the registry. A session carries at most one, and `with_extension` replaces what is there, so a library must not 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:
+A session carries one registry. Each library adds its decoders to it with `SessionConfig::update_extension`, which creates the registry if there is none. Do not attach a registry with `with_extension`: it replaces the registry, so it discards the decoders that other libraries added.
 
 ```rust,ignore
 // In each library that ships extension types:
-pub fn register(registry: &mut ProtoDecoderRegistry) -> Result<()> {
-    register_execution_plan::<MyExec>(registry)
+pub fn register(config: &mut SessionConfig) -> Result<()> {
+    config.update_extension(|registry: &mut ProtoDecoderRegistry| {
+        register_execution_plan::<MyExec>(registry)
+    })
 }
 
 // In the application that owns the session:
-let mut registry = ProtoDecoderRegistry::new();
-shuffle_lib::register(&mut registry)?;
-sampling_lib::register(&mut registry)?;
-let config = SessionConfig::new().with_extension(Arc::new(registry));
+let mut config = SessionConfig::new();
+shuffle_lib::register(&mut config)?;
+sampling_lib::register(&mut config)?;
 ```
 
 The libraries need to know nothing about each other, and the order they are called in does not matter. Dispatch is by name, and a real name collision is an error from `register_execution_plan` rather than a wrong decode later.


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org