diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 37be0bc9d..f1a842cb8 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.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`. 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 diff --git a/Cargo.lock b/Cargo.lock index ec665e285..afc7c53e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1584,6 +1584,7 @@ dependencies = [ "datafusion-ffi", "datafusion-proto", "datafusion-python-util", + "datafusion-session", "datafusion-spark", "datafusion-substrait", "futures", diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index c5f1e0167..e723d1263 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -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 } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index fe75668fb..709f84c2c 100644 --- a/crates/core/src/context.rs +++ b/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::{ @@ -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 { - // Pre-existing caveat, unrelated to query planners: this is the one - // method that mints a second `Arc` 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 . + let state_ref = self.ctx.state_ref(); + let mut guard = state_ref.write(); + if guard + .catalog_list() + .downcast_ref::() + .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, + )); + 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()) diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index bbf6b65e7..1dd950133 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -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`. 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)= diff --git a/docs/source/extension-guide/sessions.md b/docs/source/extension-guide/sessions.md index 21bfb6e35..91dd99cd9 100644 --- a/docs/source/extension-guide/sessions.md +++ b/docs/source/extension-guide/sessions.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. ::: diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 74318c83c..f4fe77ab6 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.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 diff --git a/examples/datafusion-ffi-example/python/tests/_test_table_provider.py b/examples/datafusion-ffi-example/python/tests/_test_table_provider.py index aee16f839..0deed2c9b 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_table_provider.py +++ b/examples/datafusion-ffi-example/python/tests/_test_table_provider.py @@ -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 diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 08dbc4400..84aa76a85 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.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) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 274f2a9b1..496376e63 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_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.