feat: allow replacing the global SessionContext

Promote the previously immutable global context slot in `datafusion-python-util`
from `OnceLock<Arc<SessionContext>>` to a `RwLock<Arc<SessionContext>>` and
expose `set_global_ctx` (Rust) / `SessionContext.set_as_global` (Python). Users
who register UDFs or otherwise customize a context can now make it the default
seen by `SessionContext.global_ctx()` and the module-level `read_*` helpers.

Existing snapshots returned by `get_global_ctx()` are unaffected — the swap only
changes what subsequent readers see.

Also fixes a pre-existing clippy `uninlined_format_args` nit in
`dataframe.rs` that was tripping the pre-commit hook.

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

Author: ntjohnson1

crates/core/src/context.rs

@@ -56,7 +56,7 @@ use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
 use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
 use datafusion_python_util::{
     create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx,
-    get_tokio_runtime, spawn_future, wait_for_future,
+    get_tokio_runtime, set_global_ctx, spawn_future, wait_for_future,
 };
 use object_store::ObjectStore;
 use pyo3::IntoPyObjectExt;
@@ -407,11 +407,22 @@ impl PySessionContext {
     #[staticmethod]
     #[pyo3(signature = ())]
     pub fn global_ctx() -> PyResult<Self> {
-        let ctx = get_global_ctx().clone();
+        let ctx = get_global_ctx();
         let logical_codec = Self::default_logical_codec(&ctx);
         Ok(Self { ctx, logical_codec })
     }
 
+    /// Replace the process-wide global `SessionContext` with this one.
+    ///
+    /// All subsequent callers of `SessionContext.global_ctx()` (and Rust
+    /// helpers that fall back to the global context, such as the
+    /// `read_parquet` / `read_csv` / etc. module-level helpers) will see this
+    /// context. Existing references already obtained from `global_ctx()` are
+    /// not affected.
+    pub fn set_as_global(&self) {
+        set_global_ctx(self.ctx.clone());
+    }
+
     /// Register an object store with the given name
     #[pyo3(signature = (scheme, store, host=None))]
     pub fn register_object_store(

crates/core/src/dataframe.rs

@@ -851,7 +851,7 @@ impl PyDataFrame {
             Some(f) => f
                 .parse::<datafusion::common::format::ExplainFormat>()
                 .map_err(|e| {
-                    PyDataFusionError::Common(format!("Invalid explain format '{}': {}", f, e))
+                    PyDataFusionError::Common(format!("Invalid explain format '{f}': {e}"))
                 })?,
             None => datafusion::common::format::ExplainFormat::Indent,
         };

crates/util/src/lib.rs

@@ -17,7 +17,7 @@
 
 use std::future::Future;
 use std::ptr::NonNull;
-use std::sync::{Arc, OnceLock};
+use std::sync::{Arc, OnceLock, RwLock};
 use std::time::Duration;
 
 use datafusion::datasource::TableProvider;
@@ -59,11 +59,29 @@ pub fn is_ipython_env(py: Python) -> &'static bool {
     })
 }
 
-/// Utility to get the Global Datafussion CTX
+fn global_ctx_slot() -> &'static RwLock<Arc<SessionContext>> {
+    static CTX: OnceLock<RwLock<Arc<SessionContext>>> = OnceLock::new();
+    CTX.get_or_init(|| RwLock::new(Arc::new(SessionContext::new())))
+}
+
+/// Utility to get the Global DataFusion CTX.
+///
+/// Returns an owned `Arc<SessionContext>` snapshot. The underlying slot can be
+/// replaced via [`set_global_ctx`]; existing snapshots are unaffected.
 #[inline]
-pub fn get_global_ctx() -> &'static Arc<SessionContext> {
-    static CTX: OnceLock<Arc<SessionContext>> = OnceLock::new();
-    CTX.get_or_init(|| Arc::new(SessionContext::new()))
+pub fn get_global_ctx() -> Arc<SessionContext> {
+    global_ctx_slot()
+        .read()
+        .expect("global SessionContext lock poisoned")
+        .clone()
+}
+
+/// Replace the Global DataFusion CTX. Subsequent calls to [`get_global_ctx`]
+/// will return the new context. Already-cloned `Arc`s are not affected.
+pub fn set_global_ctx(ctx: Arc<SessionContext>) {
+    *global_ctx_slot()
+        .write()
+        .expect("global SessionContext lock poisoned") = ctx;
 }
 
 /// Utility to collect rust futures with GIL released and respond to
@@ -224,3 +242,40 @@ pub fn ffi_logical_codec_from_pycapsule(obj: Bound<PyAny>) -> PyResult<FFI_Logic
 
     Ok(codec.clone())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The global slot must round-trip a custom `SessionContext`. Since the
+    /// global is process-wide, this test only asserts identity through a
+    /// single set/get cycle and restores the prior value at the end so the
+    /// test is independent of ordering with other tests in the binary.
+    #[test]
+    fn set_global_ctx_replaces_default() {
+        let prior = get_global_ctx();
+        let custom = Arc::new(SessionContext::new());
+        let custom_ptr = Arc::as_ptr(&custom);
+
+        set_global_ctx(custom.clone());
+        let observed = get_global_ctx();
+        assert_eq!(
+            Arc::as_ptr(&observed),
+            custom_ptr,
+            "get_global_ctx should return the context installed by set_global_ctx",
+        );
+
+        // A snapshot taken before the swap should be unaffected after another
+        // set_global_ctx call, because get_global_ctx clones the Arc.
+        let snapshot = get_global_ctx();
+        let replacement = Arc::new(SessionContext::new());
+        set_global_ctx(replacement);
+        assert_eq!(
+            Arc::as_ptr(&snapshot),
+            custom_ptr,
+            "previously cloned snapshots must not be invalidated by set_global_ctx",
+        );
+
+        set_global_ctx(prior);
+    }
+}

python/datafusion/context.py

@@ -566,6 +566,22 @@ def global_ctx(cls) -> SessionContext:
         wrapper.ctx = internal_ctx
         return wrapper
 
+    def set_as_global(self) -> None:
+        """Install this context as the process-wide global ``SessionContext``.
+
+        After this call, :meth:`SessionContext.global_ctx` (and the module-level
+        helpers in :mod:`datafusion.io` that fall back to the global context)
+        will return this context. Existing references already obtained from
+        ``global_ctx()`` are not invalidated.
+
+        Example::
+
+            ctx = SessionContext()
+            ctx.register_udf(my_udf)
+            ctx.set_as_global()
+        """
+        self.ctx.set_as_global()
+
     def enable_url_table(self) -> SessionContext:
         """Control if local files can be queried as tables.