Fix enable_url_table session sharing

Author: dajiaohuang

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

@@ -187,13 +187,11 @@ 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 this rule: it rebuilds the catalog
+inside the receiver's existing `SessionState` and returns another handle
+sharing the same `Arc<SessionContext>`. It preserves the session id and keeps
+weak FFI providers valid. An idempotence guard avoids wrapping the dynamic-file
+catalog more than once.
 
 ## Rule 7 — installing a planner mutates the session, and says so
 

Cargo.lock

@@ -53,6 +53,7 @@ datafusion = { workspace = true, features = ["avro", "unicode_expressions"] }
 datafusion-substrait = { workspace = true, optional = true }
 datafusion-proto = { workspace = true }
 datafusion-ffi = { workspace = true }
+datafusion-session = { workspace = true }
 datafusion-spark = { workspace = true, features = ["core"] }
 prost = { workspace = true } # keep in line with `datafusion-substrait`
 serde_json = { workspace = true }

crates/core/Cargo.toml

@@ -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::{
@@ -65,6 +68,7 @@ use datafusion_python_util::{
     ffi_query_planner_from_pycapsule, get_global_ctx, get_tokio_runtime,
     physical_optimizer_rule_from_pycapsule, spawn_future, validate_pycapsule, wait_for_future,
 };
+use datafusion_session::SessionStore;
 use object_store::ObjectStore;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyTypeError, PyValueError};
@@ -452,15 +456,30 @@ 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();
+        let mut guard = state_ref.write();
+        if guard
+            .catalog_list()
+            .downcast_ref::<DynamicFileCatalog>()
+            .is_none()
+        {
+            let factory = Arc::new(DynamicListTableFactory::new(SessionStore::new()));
+            let catalog_list = Arc::new(DynamicFileCatalog::new(
+                Arc::clone(guard.catalog_list()),
+                Arc::clone(&factory) as Arc<dyn UrlTableFactory>,
+            ));
+            factory
+                .session_store()
+                .with_state(self.ctx.state_weak_ref());
+
+            *guard = SessionStateBuilder::new_from_existing(guard.clone())
+                .with_session_id(guard.session_id().to_string())
+                .with_catalog_list(catalog_list)
+                .build();
+        }
+
         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),
         })
@@ -1471,14 +1490,12 @@ 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 — including `with_python_udf_inlining`,
+    /// `with_logical_extension_codec`, [`Self::_install_extension_codecs`], and
+    /// `enable_url_table` — 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.
     #[getter]
     pub fn __datafusion_codec_id__(&self) -> String {
         format!("{SESSION_CODEC_ID_PREFIX}{}", self.ctx.session_id())

crates/core/src/context.rs

@@ -64,13 +64,12 @@ 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 installs the URL
+catalog wrapper by rebuilding `SessionState` through the existing context's
+`state_ref()`, then returns another handle sharing the same
+`Arc<SessionContext>`. The session id and weak FFI task-context providers remain
+valid across the derivation. Repeated calls detect the existing dynamic-file
+catalog and do not nest another wrapper.
 
 (ffi_internals_rebinding)=
 

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

@@ -105,13 +105,8 @@ 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.
+:::{note}
+{py:meth}`~datafusion.SessionContext.enable_url_table` returns another handle
+to the same session and enables URL tables for every handle on it. Repeated
+calls are safe and do not wrap the catalog more than once.
 :::

docs/source/extension-guide/sessions.md

@@ -21,6 +21,22 @@
 
 ## DataFusion 56.0.0
 
+### `SessionContext.enable_url_table` now shares the existing session
+
+`enable_url_table()` now returns a handle on the same session rather than a
+separate session with copied state. The original handle sees URL tables enabled
+too, and state changes remain visible through both handles. Code that reassigns
+the result continues to work unchanged:
+
+```python
+ctx = SessionContext()
+ctx = ctx.enable_url_table()
+```
+
+Before this change, the two handles could report the same session id while
+holding divergent state; replacing the original handle could also invalidate
+FFI providers bound to it. Treat both handles as views of one session.
+
 ### `sort` and `order_by` now order `NULLS LAST` by default
 
 Calling `.sort(...)` on a DataFrame or using the `.order_by(...)` function now

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

@@ -51,3 +51,19 @@ def test_table_provider_ffi(inner_capsule: bool) -> None:
     result = ctx.read_table(table).collect()
     result = sorted((r.column(0) for r in result), key=lambda a: a[0].as_py())
     assert result == expected
+
+
+def test_table_provider_survives_enable_url_table():
+    ctx = SessionContext()
+    ctx.register_table("t", MyTableProvider(3, 2, 4))
+    session_id = ctx.session_id()
+
+    # Reassigning the derived handle must not orphan FFI providers tied to the
+    # session, and enabling URL tables repeatedly must remain safe.
+    ctx = ctx.enable_url_table()
+    ctx = ctx.enable_url_table()
+
+    assert ctx.session_id() == session_id
+    result = ctx.table("t").collect()
+    assert len(result) == 4
+    assert result[0].num_columns == 3

examples/datafusion-ffi-example/python/tests/_test_table_provider.py

@@ -669,7 +669,7 @@ def enable_url_table(self) -> SessionContext:
         """Control if local files can be queried as tables.
 
         Returns:
-            A new :py:class:`SessionContext` object with url table enabled.
+            A handle on the same session with url table enabled.
         """
         klass = self.__class__
         obj = klass.__new__(klass)

python/datafusion/context.py

@@ -939,6 +939,21 @@ def test_contexts_sharing_a_session_share_the_planner(ctx):
     assert sibling.session_id() == ctx.session_id()
 
 
+def test_enable_url_table_returns_handle_on_shared_session(ctx):
+    sibling = ctx.enable_url_table()
+    session_id = ctx.session_id()
+
+    ctx.sql("SET datafusion.execution.batch_size = 111").collect()
+
+    assert sibling.session_id() == session_id
+    assert ctx.sql("SHOW datafusion.execution.batch_size").collect()[0].column(
+        1
+    ).to_pylist() == ["111"]
+    assert sibling.sql("SHOW datafusion.execution.batch_size").collect()[0].column(
+        1
+    ).to_pylist() == ["111"]
+
+
 class _NamedCodec:
     """Wraps a codec capsule in an object that can name itself.