fix: preserve provided arguments during FFI object construction

[timsaucer]

Which issue does this PR close?

• Closes FFI constructors silently discard arguments when the input is already foreign #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:

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 apache/datafusion-python#1677, which retains the pre-fork SessionContext purely to keep the Weak valid.

🤖 Generated with Claude Code

[timsaucer]

I've marked the 3 line behavior change of the PR below.

[timsaucer]

This is the core of the PR.

[timsaucer]

This is the core of the PR.

[timsaucer]

This is the core of the PR.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 99.09910% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 81.45%. Comparing base (1038d35) to head (7599f8c).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/ffi/src/table_provider.rs	96.15%	0 Missing and 1 partial ⚠️

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

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-Weak failure 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.

[milenkovicm]

thanks @timsaucer

[timsaucer]

Thank you for the review @milenkovicm