fix: preserve provided arguments during FFI object construction - #24723
fix: preserve provided arguments during FFI object construction#24723timsaucer wants to merge 1 commit into
Conversation
FFI_LogicalExtensionCodec::new, FFI_PhysicalExtensionCodec::new, and FFI_TableProvider::new_with_ffi_codec unwrap an already-foreign input and return its original handle, silently dropping the task context provider or logical codec passed alongside. A consumer that imports a foreign codec can therefore never rebind it: re-wrapping with a different provider compiles, runs, and has no effect, so the handle keeps resolving against the session it was first built with. Because that provider is held as a Weak, the same gap forces consumers to keep the original session alive artificially or hit "TaskContextProvider went out of scope over FFI boundary". Clone the original handle and overwrite the relevant repr(C) field before returning it, matching FFI_QueryPlanner::new_with_ffi_codecs and FFI_SessionRef::new_with_ffi_codecs, which already adopt on this path. The runtime argument stays an exception: it lives in private_data owned by the library that holds the handle, so this side cannot write it without an ABI change. Document that on all three constructors rather than leaving it silent. Closes apache#24722 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
timsaucer
left a comment
There was a problem hiding this comment.
I've marked the 3 line behavior change of the PR below.
| let mut codec = codec.0.clone(); | ||
| codec.task_ctx_provider = task_ctx_provider.into(); | ||
| return codec; |
There was a problem hiding this comment.
This is the core of the PR.
| let mut codec = codec.0.clone(); | ||
| codec.task_ctx_provider = task_ctx_provider.into(); | ||
| return codec; |
There was a problem hiding this comment.
This is the core of the PR.
| let mut provider = provider.0.clone(); | ||
| provider.logical_codec = logical_codec; | ||
| return provider; |
There was a problem hiding this comment.
This is the core of the PR.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24723 +/- ##
========================================
Coverage 81.45% 81.45%
========================================
Files 1120 1120
Lines 401289 401397 +108
Branches 401289 401397 +108
========================================
+ Hits 326874 326973 +99
- Misses 55296 55299 +3
- Partials 19119 19125 +6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull request overview
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Makes datafusion-ffi constructors consistent when re-wrapping already-foreign handles by ensuring supplied rebinding arguments (task context provider / logical codec) are adopted rather than silently discarded, addressing #24722.
Changes:
- Update three constructors to clone the foreign handle and overwrite the relevant
#[repr(C)]field(s) before returning. - Add / expand rustdoc to document adopt-on-unwrap behavior and the deliberate exception for
runtime(private data owned by originating library). - Add unit tests covering rebinding behavior and the prior dangling-
Weakfailure mode.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| datafusion/ffi/src/table_provider.rs | Adopt logical_codec when re-wrapping an already-foreign table provider; document runtime exception; add rebinding test. |
| datafusion/ffi/src/query_planner.rs | Add control test asserting the already-correct query planner constructor continues to adopt supplied codecs. |
| datafusion/ffi/src/proto/physical_extension_codec.rs | Adopt task_ctx_provider on already-foreign path; document runtime exception; add rebinding test. |
| datafusion/ffi/src/proto/logical_extension_codec.rs | Adopt task_ctx_provider on already-foreign path; document runtime exception; add rebinding tests including releasing the original session. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Which issue does this PR close?
Rationale for this change
Three
datafusion-fficonstructors unwrap an already-foreign input and return its original handle, dropping the arguments passed alongside without an error or a warning:FFI_LogicalExtensionCodec::new— discardstask_ctx_providerFFI_PhysicalExtensionCodec::new— discardstask_ctx_providerFFI_TableProvider::new_with_ffi_codec— discardslogical_codecThe 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-pythonthat 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 withTaskContextProvider 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_codecsandFFI_SessionRef::new_with_ffi_codecsboth 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, matchingFFI_QueryPlanner::new_with_ffi_codecs:The
runtimeargument is a deliberate exception. Unlike the codecs and the task context provider,runtimelives inprivate_data, which belongs to the library that owns the handle — this side cannot write it without an ABI change.FFI_SessionRef::new_with_ffi_codecsalready 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
mainand pass here.ffi_logical_extension_codec_rebind_adopts_task_ctx_providerffi_logical_extension_codec_rebind_releases_original_session— covers the dangling-Weakfailure mode: session A is dropped after the rebind, and the handle stays usableffi_physical_extension_codec_rebind_adopts_task_ctx_providertest_rebind_foreign_table_provider_adopts_logical_codectest_rebind_foreign_query_planner_adopts_codecs— a control over the already-correct sibling, so the two paths stay in agreementWorth flagging for reviewers, since it is easy to write a test here that silently proves nothing:
impl From<&FFI_LogicalExtensionCodec> for Arc<dyn LogicalExtensionCodec>compareslibrary_marker_idfirst and returns the original localArcon a match, so within one library the foreign branch is never reached. Each test overrideslibrary_marker_idwithcrate::mock_foreign_marker_idand asserts the import really did produce aForeign*wrapper before exercising the rebind.cargo test -p datafusion-ffi --all-featurespasses (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_providerorlogical_codecto 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-pythondrop the workaround in apache/datafusion-python#1677, which retains the pre-forkSessionContextpurely to keep theWeakvalid.🤖 Generated with Claude Code