Skip to content

enable_url_table returns a handle on a different session, unlike every other with_* derivation #1708

Description

@timsaucer

Describe the bug

SessionContext.enable_url_table is the only derivation in the binding that returns a handle on a different session than the one it was called on, and it is the only one that leaves two live handles reporting the same session_id with divergent state. Every other with_* method returns a handle wrapping the same underlying session.

The root cause is one line. Upstream's SessionContext::enable_url_table consumes self (datafusion/core/src/execution/context/mod.rs:414):

pub fn enable_url_table(self) -> Self {
    let current_catalog_list = Arc::clone(self.state.read().catalog_list());
    let factory = Arc::new(DynamicListTableFactory::new(SessionStore::new()));
    let catalog_list = Arc::new(DynamicFileCatalog::new(
        current_catalog_list,
        Arc::clone(&factory) as Arc<dyn UrlTableFactory>,
    ));
    let session_id = self.session_id.clone();
    let ctx: SessionContext = self
        .into_state_builder()
        .with_session_id(session_id)
        .with_catalog_list(catalog_list)
        .build()
        .into();
    factory.session_store().with_state(ctx.state_weak_ref());
    ctx
}

Taking self by value means the old and new contexts are never meant to coexist, which is exactly why carrying session_id over is correct there. The binding at crates/core/src/context.rs:425 takes &self and clones instead:

ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()),

That manufactures the coexistence upstream's signature prevents, and everything else is downstream of it: the must-not-outlive caveat at crates/core/src/context.rs:428, the "one exception" paragraph at docs/source/contributor-guide/ffi.md:475-477, and the state divergence below.

To Reproduce

from datafusion import SessionContext

a = SessionContext()
b = a.enable_url_table()
print("same session_id:", a.session_id() == b.session_id())

a.sql("SET datafusion.execution.batch_size = 111").collect()
print("a", a.sql("SHOW datafusion.execution.batch_size").collect()[0].column(1).to_pylist())
print("b", b.sql("SHOW datafusion.execution.batch_size").collect()[0].column(1).to_pylist())
same session_id: True
a ['111']
b ['8192']

Two live sessions with independent SessionState, both reporting one id.

The second half of the problem is not reproducible without a built FFI extension, but it is what the existing caveat comments are about: any FFI codec, query planner, or task-context provider handed out by the receiver holds an FFI_TaskContextProvider bound weakly to the receiver's Arc<SessionContext>. The returned context is a different allocation, so ctx = ctx.enable_url_table() drops the last strong reference and those components fail at query time with TaskContextProvider went out of scope over FFI boundary. examples/create-context.py:41 is written as ctx = ctx.enable_url_table() — exactly the pattern the caveat forbids. Harmless in that example because it installs no FFI components, but it is the shape users copy.

Expected behavior

enable_url_table should behave like with_logical_extension_codec, with_physical_extension_codec, and with_python_udf_inlining: return a handle on the same underlying session, so session_id stays unique to a session, state cannot diverge, and FFI components bound to any handle in the lineage stay valid.

The pattern already exists in the same file. set_session_query_planner (crates/core/src/context.rs:1729-1740) mutates through state_ref() rather than deriving a new context, including the load-bearing with_session_id carry-over:

let state_ref = self.ctx.state_ref();
let factory = Arc::new(DynamicListTableFactory::new(SessionStore::new()));
let catalog_list = Arc::new(DynamicFileCatalog::new(
    Arc::clone(state_ref.read().catalog_list()),
    Arc::clone(&factory) as Arc<dyn UrlTableFactory>,
));
{
    let mut guard = state_ref.write();
    *guard = SessionStateBuilder::new_from_existing(guard.clone())
        .with_session_id(guard.session_id().to_string())
        .with_catalog_list(catalog_list)
        .build();
}
factory.session_store().with_state(Arc::downgrade(&state_ref));

The returned handle then shares Arc::clone(&self.ctx) like every sibling.

Two details for whoever picks this up:

  • Needs an idempotence guard. DynamicFileCatalog::new wraps whatever catalog list is current, so calling enable_url_table twice nests wrappers. Upstream has the same wart, but mutating in place makes it easier to hit. with_python_udf_inlining (crates/core/src/context.rs:1571-1588) already establishes the pattern of skipping the state rebuild when the call would change nothing.
  • This is an api change. ctx = ctx.enable_url_table() keeps working, but the original handle also gains url tables afterwards. That is the same trade with_logical_extension_codec already documents — it takes effect on the shared session even if the returned context is discarded — so this is bringing enable_url_table in line rather than introducing a new surprise. Needs a section in docs/source/user-guide/upgrade-guides.md, plus updates to the docstring at python/datafusion/context.py:589, the "one exception" paragraph at docs/source/contributor-guide/ffi.md:475-477, and the shared-derivation list at docs/source/contributor-guide/ffi.md:481-486.

Additional context

Found while reviewing #1679. That PR's _derive_for_extensions currently forks state the same way, so the "one method that mints a second Arc<SessionContext>" comment at crates/core/src/context.rs:428 is temporarily inaccurate. If #1679 lands with that fork removed, enable_url_table is once again the only site and that comment becomes true again.

The two are independent in both directions and should not be chained: #1679 does not need this fixed, and this does not need #1679.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions