feat: share the session RuntimeEnv across the FFI boundary - #24733
Draft
timsaucer wants to merge 2 commits into
Draft
feat: share the session RuntimeEnv across the FFI boundary#24733timsaucer wants to merge 2 commits into
timsaucer wants to merge 2 commits into
Conversation
A `Session` shared over FFI now carries its `RuntimeEnv`, so a table provider shared over FFI reaches the object stores and memory budget of the session executing it. `ForeignSession::runtime_env` and the `FFI_TaskContext` conversion both built a default `RuntimeEnv`. A provider that registered its object store on the session during `TableProvider::scan` could not reach that store when the resulting plan was executed, failing with "No suitable object store found", and foreign plans ran against an unbounded memory pool regardless of `datafusion.execution.memory_limit`. `RuntimeEnv` is a plain struct whose field list depends on enabled features, so it is not passed as an opaque pointer. Its components cross individually: `ObjectStoreRegistry` and `MemoryPool` are shared as trait objects, while `DiskManager` and `CacheManager` have their configuration copied and each side builds its own. A store used within the library that created it stays on the local fast path: it round trips through the foreign registry but is recovered as the original `Arc<dyn ObjectStore>`, so reads never cross the boundary. Object store error variants are preserved across the boundary rather than flattened to a message, so optimistic concurrency control built on `AlreadyExists` and conditional reads built on `NotModified` and `Precondition` keep working. `ResourcesExhausted` is likewise preserved so spilling operators still spill instead of failing the query. Breaking changes: * `FFI_TaskContext` gained a `runtime_env` field, changing its ABI. * `impl From<Arc<TaskContext>> for FFI_TaskContext` is removed; it could not supply a tokio runtime handle. Use `FFI_TaskContext::new`. Part of apache#19277. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
documentation
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24733 +/- ##
==========================================
+ Coverage 81.44% 81.48% +0.03%
==========================================
Files 1120 1131 +11
Lines 401605 404340 +2735
Branches 401605 404340 +2735
==========================================
+ Hits 327098 329480 +2382
- Misses 55348 55590 +242
- Partials 19159 19270 +111 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
runtime_envandexecution_props; this PR implementsruntime_envonly, so it does not close it.Rationale for this change
A table provider shared over FFI is asked to
scanby a session in another library, and theExecutionPlanit returns is later executed with aTaskContextproduced by that same session. Providers backed by remote storage build their ownObjectStoreand register it on the session duringscan, then read through it at execution time.Both
ForeignSession::runtime_envand theFFI_TaskContextconversion built a defaultRuntimeEnv. Planning and execution happen on opposite sides of the boundary, so the registration was never visible where it was needed and the scan failed with:Registering the store on the consumer's session instead did not help, because the plan executing on the provider's side reconstructed its own default environment.
The same gap meant foreign plans allocated from a fresh
UnboundedMemoryPool, sodatafusion.execution.memory_limitwas silently ignored for them and their allocations were invisible to the session's accounting.What changes are included in this PR?
RuntimeEnvis a plain struct whose field list depends on enabled features (parquet_encryptionadds one), so passing anArc<RuntimeEnv>as an opaque pointer would be undefined behaviour between libraries built with different feature sets. Each component crosses on its own terms instead:ObjectStoreRegistryFFI_ObjectStoreRegistryMemoryPoolFFI_MemoryPoolDiskManagerCacheManagerNew in
datafusion_ffi::execution:FFI_RuntimeEnv/FFI_RuntimeConfigFFI_MemoryPool/ForeignMemoryPool/FFI_TryGrowResultFFI_ObjectStore/FFI_ObjectStoreRegistryand theirForeign*wrappers, covering every requiredObjectStoremethod plus theget_rangesandlist_with_offsetoverridesA few details worth a reviewer's attention:
FFI_ObjectStore::as_localrecovers the originalArc<dyn ObjectStore>, so no data crosses the boundary. Making that work needs a side table keyed on the wrapper's address, becauseObjectStorehas noAnysupertrait and adyn ObjectStorecannot be tested for a concrete type. A one-line upstream change toobject_storewould replace it with a downcast; I plan to propose that separately.AlreadyExistsand conditional reads built onNotModified/Preconditiondepend on the discriminant, and delta-style commit protocols would silently lose their concurrency control otherwise.ResourcesExhaustedis preserved for the same reason: a dozen spilling operators match on it to decide whether to spill rather than fail the query.Bytes::from_owner.GetResultPayload::Fileis not forwarded. A raw file descriptor is not portable, so payloads always cross as a byte stream.Extensionsare dropped on all options structs. They areTypeId-keyed andTypeIdis not stable across separately compiled libraries, so the contents cannot be interpreted on the far side even in principle.Are these changes tested?
Yes. 212 tests pass.
The important one is
test_object_store_crosses_ffi_boundaryindatafusion/ffi/tests/ffi_integration.rs, which is a genuine cross-library test: the table provider is compiled into the cdylib and loaded withlibloading, so the two sides have distinct library markers. It registers an object store on the session duringscanand reads it back duringexecute, and also reports the memory pool limit it observes so the host can assert its limit reached the foreign plan.I verified that test actually catches the bug by reverting each fix independently. Both are load bearing, and each reproduces the reported error:
Unit tests cover error-variant round trips for every
object_store::Errorvariant,PutMode::Createandcopy_if_not_existssurfacingAlreadyExiststhrough the wrapper, memory limits being enforced and reservations released across the boundary, and the local fast path returning the original store.Are there any user-facing changes?
Yes, including breaking changes to public APIs. Documented in the 56.0.0 upgrade guide:
FFI_TaskContextgained aruntime_envfield, changing its ABI. All FFI providers and consumers must be rebuilt together.impl From<Arc<TaskContext>> for FFI_TaskContextis removed. It could not supply a tokio runtime handle, producing a context whose object stores could not be polled from a foreign executor. UseFFI_TaskContext::new.datafusion-ffinow depends onobject_storedirectly and exposes its types, so providers and consumers must agree on theobject_storeversion too.Behavioural changes worth calling out: plans shared over FFI are now bounded by the session's memory limit and may return
ResourcesExhaustedwhere they previously did not, andDiskManageris not shared, so spilling can use up to twicemax_temp_directory_sizeacross the two sides.