From e04625117f01ae378495639750ce617eb2d9ce00 Mon Sep 17 00:00:00 2001 From: Yifan Chen Date: Sun, 6 Sep 2026 08:21:08 -0700 Subject: [PATCH] fix: enable URL tables on the shared session Preserve context and state identity, initialize the URL factory before publishing it, and avoid nested catalog wrappers under one write lock. Cover SQL configuration, registrations, aliases, and real FFI provider lifetimes; document the API change for #1708. Generated-by: Codex (GPT-6) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 7 +- crates/core/src/context.rs | 33 +++++-- docs/source/contributor-guide/ffi.md | 11 ++- docs/source/user-guide/upgrade-guides.md | 28 ++++++ .../_test_three_library_query_planner.py | 37 ++++++++ python/datafusion/context.py | 20 ++++- python/tests/test_context.py | 86 +++++++++++++++++++ 7 files changed, 206 insertions(+), 16 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 294ebfb3a..cd7c9b22d 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -154,8 +154,9 @@ guards this. Its `WHERE` clause is load-bearing: filter pushdown upgrades the weak handle during logical optimization, before plan serialization could fail first for an unrelated reason. -`SessionContext.enable_url_table` is the one method that mints a second -allocation for a session. Its result must not outlive the receiver. +`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 @@ -185,7 +186,7 @@ pins that; changing it should be deliberate. ## Where the truth is -- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. +- `docs/source/contributor-guide/ffi.md` — the protocol and session ownership. - `docs/source/user-guide/upgrade-guides.md` — every past migration. - `crates/core/src/codec.rs` — the codec chain: the envelope, identity dispatch, and the two unframed cases from Rule 8. diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 84182ff19..a4bfcd4f0 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::{ @@ -423,13 +426,29 @@ 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. 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. + 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::() { + 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, + )); + // 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), }) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d86858a83..4ccee6601 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -417,18 +417,21 @@ codec simply retain the session that built it: a codec handed to a provider is r registered straight back into that same session, which 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. +`SessionContext.enable_url_table` also preserves that allocation. It wraps the +existing catalog list in place, so previously exported weak providers remain valid +while any handle on the session is alive. Repeated calls do not nest catalog wrappers. ### What a derived context shares -`with_logical_extension_codec`, `with_physical_extension_codec`, and +`enable_url_table`, `with_logical_extension_codec`, `with_physical_extension_codec`, and `with_python_udf_inlining` return a new `SessionContext` wrapping the *same* underlying session. Only the Python-side codec settings differ; catalogs, tables, registered functions, and configuration are the one shared session, so a registration on either side is visible to both. +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. + `set_query_planner` does not return anything. The query planner lives in `SessionState`, so it is a property of the session rather than of a handle on it, and installing one is visible to every context sharing that session — including ones a `with_*` call returned diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 257749c3a..e5463e668 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/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}`ffi` hook methods. Users who contribute their own `LogicalExtensionCodec` or `PhysicalExtensionCodec` via FFI must update diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index c6ef2072a..3845ffa74 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -76,6 +76,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. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 644c7b445..181fd2900 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -594,10 +594,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) diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3c95835af..8fc9932dd 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -19,6 +19,7 @@ import gzip import pathlib import shutil +from concurrent.futures import ThreadPoolExecutor import pyarrow as pa import pyarrow.dataset as ds @@ -41,6 +42,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())