timsaucer opened a new pull request, #24723:
URL: https://github.com/apache/datafusion/pull/24723
## Which issue does this PR close?
- Closes #24722.
## Rationale for this change
Three `datafusion-ffi` constructors unwrap an already-foreign input and
return its original handle, dropping the arguments passed alongside without an
error or a warning:
- `FFI_LogicalExtensionCodec::new` — discards `task_ctx_provider`
- `FFI_PhysicalExtensionCodec::new` — discards `task_ctx_provider`
- `FFI_TableProvider::new_with_ffi_codec` — discards `logical_codec`
The consequence is that a consumer which imports a foreign codec can never
rebind it. Re-wrapping with a different provider compiles, runs, and has no
effect, so the handle keeps resolving against whatever session it was first
built with. In `datafusion-python` that shows up as decode callbacks resolving
names against a pre-fork session: a UDF registered after the fork is invisible
to them, and the config they see is a stale snapshot.
There is a second failure mode with the same root cause. The provider is
held as a `Weak`, so a consumer that cannot rebind must keep the original
session alive artificially or the capsule starts failing with
`TaskContextProvider went out of scope over FFI boundary`.
The two sibling constructors that hit the same case already do the opposite
— `FFI_QueryPlanner::new_with_ffi_codecs` and
`FFI_SessionRef::new_with_ffi_codecs` both adopt the supplied codecs on the
unwrap path, and the former documents that guarantee explicitly. This PR makes
the other three consistent with them.
## What changes are included in this PR?
On the already-foreign path, each of the three constructors now clones the
original handle and overwrites the relevant `#[repr(C)]` field before returning
it, matching `FFI_QueryPlanner::new_with_ffi_codecs`:
```rust
if let Some(codec) = (Arc::clone(&codec) as Arc<dyn Any>)
.downcast_ref::<ForeignLogicalExtensionCodec>()
{
let mut codec = codec.0.clone();
codec.task_ctx_provider = task_ctx_provider.into();
return codec;
}
```
The `runtime` argument is a deliberate exception. Unlike the codecs and the
task context provider, `runtime` lives in `private_data`, which belongs to the
library that owns the handle — this side cannot write it without an ABI change.
`FFI_SessionRef::new_with_ffi_codecs` already takes the same position
("retaining its original private data and runtime"). Rather than leave that
silent, all three constructors now document it, alongside the new
adopt-on-unwrap guarantee.
No public signatures change, and no behavior changes on the non-foreign path.
## Are these changes tested?
Yes — five new unit tests, one per behavior, in each affected module's own
test module. All five fail on `main` and pass here.
- `ffi_logical_extension_codec_rebind_adopts_task_ctx_provider`
- `ffi_logical_extension_codec_rebind_releases_original_session` — covers
the dangling-`Weak` failure mode: session A is dropped after the rebind, and
the handle stays usable
- `ffi_physical_extension_codec_rebind_adopts_task_ctx_provider`
- `test_rebind_foreign_table_provider_adopts_logical_codec`
- `test_rebind_foreign_query_planner_adopts_codecs` — a control over the
already-correct sibling, so the two paths stay in agreement
Worth flagging for reviewers, since it is easy to write a test here that
silently proves nothing: `impl From<&FFI_LogicalExtensionCodec> for Arc<dyn
LogicalExtensionCodec>` compares `library_marker_id` first and returns the
original local `Arc` on a match, so within one library the foreign branch is
never reached. Each test overrides `library_marker_id` with
`crate::mock_foreign_marker_id` and asserts the import really did produce a
`Foreign*` wrapper before exercising the rebind.
`cargo test -p datafusion-ffi --all-features` passes (152 tests).
## Are there any user-facing changes?
Yes, a behavior change, though it replaces a silent no-op with the
documented intent.
Callers that pass a `task_ctx_provider` or `logical_codec` to these
constructors alongside an already-foreign input previously had that argument
ignored; it now takes effect. Anything relying on the old handle being returned
untouched would see the change — but since the old path gave no way to observe
or opt into that, it is hard to depend on deliberately.
Downstream, this lets `datafusion-python` drop the workaround in
https://github.com/apache/datafusion-python/pull/1677, which retains the
pre-fork `SessionContext` purely to keep the `Weak` valid.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
