fix: enable URL tables on the shared session

.ai/skills/ffi-capsule-protocol/SKILL.md

@@ -187,13 +187,9 @@ an instruction to derive one first. It is not: the factories are handed the
 receiver, and the returned handle shares its allocation. There is nothing to
 keep alive separately and nothing to garbage-collect out from under a provider.
 
-`SessionContext.enable_url_table` is the one method that mints a second
-allocation for a session. Its result must not outlive the receiver, and it also
-forks the session's `SessionState` while keeping its id, so two handles report
-one `session_id()` with divergent configuration. That is a bug rather than a
-design — tracked in
-[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708)
-— so do not cite it as precedent for deriving a replacement context.
+`SessionContext.enable_url_table` follows the same rule: it replaces only the
+catalog list through `state_ref()` and returns a handle sharing the original
+allocation. Its idempotence check and catalog replacement share one write lock.
 
 ## Rule 7 — installing a planner mutates the session, and says so
 

crates/core/src/context.rs

@@ -27,8 +27,11 @@ use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory};
+use datafusion::catalog::{
+    CatalogProvider, CatalogProviderList, DynamicFileCatalog, TableProviderFactory, UrlTableFactory,
+};
 use datafusion::common::{DFSchema, ScalarValue, TableReference, exec_err};
+use datafusion::datasource::dynamic_file::DynamicListTableFactory;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::listing::{
@@ -423,15 +426,29 @@ impl PySessionContext {
     }
 
     pub fn enable_url_table(&self) -> PyResult<Self> {
-        // Pre-existing caveat, unrelated to query planners: this is the one
-        // method that mints a second `Arc<SessionContext>` for a session, and
-        // it also forks the session's state while keeping its id. Any weak
-        // `FFI_TaskContextProvider` handed out by the receiver stays bound to
-        // the receiver, so the returned context must not outlive it. See
-        // `set_session_query_planner` for why everything else mutates in place.
-        // Tracked as a bug in <https://github.com/apache/datafusion-python/issues/1708>.
+        let state_ref = self.ctx.state_ref();
+        {
+            // Check and replace under one lock so concurrent calls cannot nest
+            // wrappers or overwrite a newer catalog list.
+            let mut state = state_ref.write();
+            if !state.catalog_list().is::<DynamicFileCatalog>() {
+                let factory = Arc::new(DynamicListTableFactory::default());
+                // Bind before publishing the catalog: a reader must never see
+                // a factory whose session store has not been initialized.
+                factory
+                    .session_store()
+                    .with_state(self.ctx.state_weak_ref());
+                let catalog_list = Arc::new(DynamicFileCatalog::new(
+                    Arc::clone(state.catalog_list()),
+                    factory as Arc<dyn UrlTableFactory>,
+                ));
+                // Only the catalog changes. In particular, preserve the state
+                // and context allocations targeted by weak FFI providers.
+                state.register_catalog_list(catalog_list);
+            }
+        }
         Ok(PySessionContext {
-            ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()),
+            ctx: Arc::clone(&self.ctx),
             logical_codec: Arc::clone(&self.logical_codec),
             physical_codec: Arc::clone(&self.physical_codec),
         })
@@ -1434,14 +1451,13 @@ impl PySessionContext {
     /// time, and a payload written through one would resolve to the other on
     /// decode. See [`SESSION_CODEC_ID_PREFIX`].
     ///
-    /// Handles derived from one session — `with_python_udf_inlining`,
-    /// `with_logical_extension_codec`, [`Self::_install_extension_codecs`] —
-    /// report the same id even though their codec chains differ, so installing
-    /// two of them on one target is refused. That is the intended answer: they
-    /// share a `state_ref`, so their payloads would resolve against the same
-    /// session and are indistinguishable on decode. Every derivation shares the
-    /// session for exactly this reason; `enable_url_table` is the one that does
-    /// not, and it is tracked as a bug.
+    /// Handles derived from one session — `enable_url_table`,
+    /// `with_python_udf_inlining`, `with_logical_extension_codec`,
+    /// [`Self::_install_extension_codecs`] — report the same id even though
+    /// their codec chains may differ, so installing two of them on one target
+    /// is refused. That is the intended answer: they share a `state_ref`, so
+    /// their payloads would resolve against the same session and are
+    /// indistinguishable on decode.
     #[getter]
     pub fn __datafusion_codec_id__(&self) -> String {
         format!("{SESSION_CODEC_ID_PREFIX}{}", self.ctx.session_id())

docs/source/contributor-guide/ffi-internals.md

@@ -64,13 +64,10 @@ would close the cycle
 `SessionContext -> catalog -> FFI provider -> FFI codec -> SessionContext` and
 leak it.
 
-`SessionContext.enable_url_table` is the one exception. It clones the
-underlying `SessionContext`, so the returned context has an allocation of its
-own and must not outlive the receiver. It also forks the session's state while
-keeping its id, so two handles report one `session_id()` with divergent
-configuration. That is a bug rather than a design, tracked in
-[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708);
-do not copy the pattern.
+`SessionContext.enable_url_table` follows the same rule. It replaces only the
+catalog list through `state_ref()` and returns a handle sharing the original
+allocation. The idempotence check and catalog replacement share one write lock,
+so concurrent calls cannot nest wrappers or overwrite a newer catalog list.
 
 (ffi_internals_rebinding)=
 

docs/source/extension-guide/sessions.md

@@ -52,6 +52,7 @@ divergence.
 
 ## What a derived context shares
 
+{py:meth}`~datafusion.SessionContext.enable_url_table`,
 {py:meth}`~datafusion.SessionContext.with_logical_extension_codec`,
 {py:meth}`~datafusion.SessionContext.with_physical_extension_codec`,
 {py:meth}`~datafusion.SessionContext.with_python_udf_inlining`, and
@@ -105,13 +106,6 @@ The same rule applies to a capsule you take off a context inside your own code:
 a codec capsule taken from a throwaway `SessionContext()` names a session that
 is already gone and fails on first use.
 
-:::{warning}
-{py:meth}`~datafusion.SessionContext.enable_url_table` is an exception to the
-one-session-one-allocation rule above: it clones the underlying
-`SessionContext`, so the returned context has an allocation of its own and must
-not outlive the receiver. It also forks the session's state while keeping its
-id, so two handles report one `session_id()` with divergent configuration. That
-is a bug rather than a design, tracked in
-[apache/datafusion-python#1708](https://github.com/apache/datafusion-python/issues/1708);
-do not build on the behaviour.
-:::
+Enabling URL tables takes effect on the shared session even if the returned
+handle is discarded. The returned handle keeps the receiver's codec settings
+unchanged, and repeated calls do not nest catalog wrappers.

docs/source/user-guide/upgrade-guides.md

@@ -21,6 +21,34 @@
 
 ## DataFusion 55.0.0
 
+### URL tables share the original session
+
+`SessionContext.enable_url_table()` now enables URL tables on the existing session
+and returns another handle on it. Previously it copied session state into a separate
+session while retaining the same session id. Configuration and function registrations
+could then diverge, and replacing the original handle could invalidate FFI providers.
+
+Before, callers had to use the returned context to query file paths:
+
+```python
+ctx = SessionContext()
+enabled = ctx.enable_url_table()
+# Only enabled could query local file paths as tables.
+```
+
+After, both handles share configuration, registrations, and URL table support:
+
+```python
+ctx = SessionContext()
+ctx.enable_url_table()  # Takes effect even when the returned handle is discarded.
+```
+
+Existing `ctx = ctx.enable_url_table()` calls continue to work and now retain FFI
+providers bound to the original session. Repeated calls have no effect. To keep a
+session without URL table support, create a separate `SessionContext` explicitly.
+
+### FFI codec hooks receive the session
+
 This release extends the change made in 52.0.0 to the remaining
 {ref}`extension_capsule_protocol` hook methods. Users who contribute their own
 `LogicalExtensionCodec` or

examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py

@@ -93,6 +93,43 @@ def probe_context(
     return ctx, logical_codec, physical_codec
 
 
+@pytest.mark.parametrize("keep_returned", [False, True])
+def test_enable_url_table_preserves_ffi_providers(keep_returned):
+    """A catalog's unreachable weak codec remains valid through URL enabling."""
+    ctx, logical_codec, physical_codec = probe_context(
+        logical_requires=HOST_ONLY_UDF, physical_requires=HOST_ONLY_UDF, max_rows=100
+    )
+    ctx.register_catalog_provider("ffi_catalog", MyCatalogProvider())
+    session_id = ctx.session_id()
+    alias = ctx.with_python_udf_inlining(enabled=True)
+    if keep_returned:
+        ctx = alias.enable_url_table()
+    else:
+        alias.enable_url_table()
+    del alias
+    gc.collect()
+
+    for _ in range(2):
+        # Force filter pushdown through a real foreign catalog before planning.
+        batches = ctx.sql(
+            "SELECT units FROM ffi_catalog.my_schema.my_table WHERE units > 5"
+        ).collect()
+        assert sorted(v for b in batches for v in b.column(0).to_pylist()) == [
+            7,
+            10,
+            20,
+            30,
+        ]
+        assert logical_codec.table_provider_decode_calls() > 0
+        assert physical_codec.execution_plan_decode_calls() > 0
+        assert logical_codec.last_task_context_session_id() == session_id
+        assert physical_codec.last_task_context_session_id() == session_id
+        assert logical_codec.task_context_udf_resolutions() > 0
+        assert physical_codec.task_context_udf_resolutions() > 0
+        ctx = ctx.enable_url_table()
+        gc.collect()
+
+
 def test_logical_codec_resolves_a_host_registered_udf():
     """``try_decode_table_provider`` sees the host session's registry.
 

python/datafusion/context.py

@@ -551,7 +551,8 @@ class SessionContext:
     See :ref:`user_guide_concepts` in the online documentation for more information.
 
     **A context is a handle on a session, not the session itself.** The
-    ``with_*`` methods — :py:meth:`with_logical_extension_codec`,
+    derivation methods — :py:meth:`enable_url_table`,
+    :py:meth:`with_logical_extension_codec`,
     :py:meth:`with_physical_extension_codec`,
     :py:meth:`with_python_udf_inlining`, and :py:meth:`with_extensions` —
     return a new context wrapping the *same* underlying session. Only the
@@ -560,10 +561,11 @@ class SessionContext:
     either handle is visible to both.
 
     A few things therefore belong to the session rather than to a handle, and
-    take effect even if the handle that set them is discarded: the query
-    planner (see :py:meth:`set_query_planner`), and the rebuild of an installed
-    foreign planner that follows installing a codec. :ref:`extension_sessions`
-    in the online documentation works through when that matters.
+    take effect even if the handle that set them is discarded: URL table
+    support, the query planner (see :py:meth:`set_query_planner`), and the
+    rebuild of an installed foreign planner that follows installing a codec.
+    :ref:`extension_sessions` in the online documentation works through when
+    that matters.
 
     **Keep a context alive for as long as anything derived from it is in use.**
     A :py:class:`~datafusion.DataFrame`, logical plan, or exported capsule does
@@ -619,10 +621,26 @@ def global_ctx(cls) -> SessionContext:
         return wrapper
 
     def enable_url_table(self) -> SessionContext:
-        """Control if local files can be queried as tables.
+        """Enable querying local files as tables on the shared session.
+
+        The receiver and all handles sharing its session gain URL table support,
+        even if the returned handle is discarded. Repeated calls have no effect.
+        Registered catalogs, functions, configuration, and session identity are
+        preserved, as are FFI providers bound to the session.
 
         Returns:
-            A new :py:class:`SessionContext` object with url table enabled.
+            A new :py:class:`SessionContext` handle wrapping the same session.
+
+        Examples:
+            >>> ctx = SessionContext()
+            >>> enabled = ctx.enable_url_table()
+            >>> enabled.session_id() == ctx.session_id()
+            True
+            >>> ctx.sql("SET datafusion.execution.batch_size = 111").collect()
+            []
+            >>> batches = enabled.sql("SHOW datafusion.execution.batch_size").collect()
+            >>> batches[0].column(1).to_pylist()
+            ['111']
         """
         klass = self.__class__
         obj = klass.__new__(klass)

python/tests/test_context.py

@@ -20,6 +20,7 @@
 import gzip
 import pathlib
 import shutil
+from concurrent.futures import ThreadPoolExecutor
 
 import pyarrow as pa
 import pyarrow.dataset as ds
@@ -43,6 +44,91 @@ def test_create_context_no_args():
     SessionContext()
 
 
+@pytest.mark.parametrize("discard_returned", [False, True])
+def test_enable_url_table_shares_session(tmp_path, discard_returned):
+    """URL tables, configuration, and registrations belong to all aliases."""
+    ctx = SessionContext()
+    original_config = (
+        ctx.sql("SHOW datafusion.catalog.create_default_catalog_and_schema")
+        .collect()[0]
+        .column(1)
+        .to_pylist()
+    )
+    alias = ctx.with_python_udf_inlining(enabled=False)
+    session_id = ctx.session_id()
+    ctx.sql("CREATE SCHEMA existing").collect()
+    ctx.register_record_batches("existing.numbers", [[pa.record_batch({"n": [7]})]])
+    ctx.register_udf(
+        udf(lambda x: x, [pa.int64()], pa.int64(), "immutable", "identity")
+    )
+    path = tmp_path / "numbers.csv"
+    path.write_text("n\n7\n")
+
+    if discard_returned:
+        alias.enable_url_table()
+        returned = ctx
+    else:
+        returned = alias.enable_url_table()
+
+    # Both directions must see subsequent SETs, not only the initial snapshot.
+    for writer, reader, value in [(ctx, returned, 111), (returned, alias, 222)]:
+        writer.sql(f"SET datafusion.execution.batch_size = {value}").collect()
+        assert reader.sql("SHOW datafusion.execution.batch_size").collect()[0].column(
+            1
+        ).to_pylist() == [str(value)]
+
+    returned.register_udf(
+        udf(lambda x: x, [pa.int64()], pa.int64(), "immutable", "later_identity")
+    )
+    for handle in (ctx, alias, returned):
+        assert handle.session_id() == session_id
+        assert (
+            handle.sql("SHOW datafusion.catalog.create_default_catalog_and_schema")
+            .collect()[0]
+            .column(1)
+            .to_pylist()
+            == original_config
+        )
+        assert handle.sql("SELECT identity(n) FROM existing.numbers").collect()[
+            0
+        ].column(0).to_pylist() == [7]
+        assert handle.sql("SELECT later_identity(9)").collect()[0].column(
+            0
+        ).to_pylist() == [9]
+        assert handle.sql(f'SELECT n FROM "{path}"').collect()[0].column(
+            0
+        ).to_pylist() == [7]
+        handle.enable_url_table()
+        handle.enable_url_table()
+        assert handle.sql(f'SELECT n FROM "{path}"').collect()[0].column(
+            0
+        ).to_pylist() == [7]
+
+
+def test_enable_url_table_from_multiple_aliases(tmp_path):
+    """Enabling through multiple handles preserves concurrent registrations."""
+    ctx = SessionContext()
+    path = tmp_path / "numbers.csv"
+    path.write_text("n\n7\n")
+    aliases = [ctx.with_python_udf_inlining(enabled=False) for _ in range(4)]
+
+    def enable_and_register(index):
+        alias = aliases[index]
+        for _ in range(4):
+            alias.enable_url_table()
+            alias.sql(f'SELECT n FROM "{path}"').collect()
+        alias.register_udf(
+            udf(lambda x: x, [pa.int64()], pa.int64(), "immutable", f"identity_{index}")
+        )
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        list(executor.map(enable_and_register, range(4)))
+    for index in range(4):
+        assert ctx.sql(f"SELECT identity_{index}(7)").collect()[0].column(
+            0
+        ).to_pylist() == [7]
+
+
 def test_create_context_session_config_only():
     SessionContext(config=SessionConfig())