From cb009e7dd58aba28ea784ec3b55c051b08b1ec88 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 13:09:10 -0400 Subject: [PATCH 1/2] feat: share the session RuntimeEnv across the FFI boundary 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`, 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> for FFI_TaskContext` is removed; it could not supply a tokio runtime handle. Use `FFI_TaskContext::new`. Part of #19277. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 3 + datafusion/ffi/Cargo.toml | 3 + datafusion/ffi/README.md | 34 + datafusion/ffi/src/execution/memory_pool.rs | 602 ++++++++++++++ datafusion/ffi/src/execution/mod.rs | 9 + .../ffi/src/execution/object_store/buffer.rs | 177 +++++ .../ffi/src/execution/object_store/error.rs | 415 ++++++++++ .../ffi/src/execution/object_store/mod.rs | 360 +++++++++ .../src/execution/object_store/multipart.rs | 276 +++++++ .../src/execution/object_store/registry.rs | 394 ++++++++++ .../ffi/src/execution/object_store/store.rs | 734 ++++++++++++++++++ .../ffi/src/execution/object_store/stream.rs | 490 ++++++++++++ .../ffi/src/execution/object_store/types.rs | 716 +++++++++++++++++ datafusion/ffi/src/execution/runtime_env.rs | 492 ++++++++++++ datafusion/ffi/src/execution/task_ctx.rs | 60 +- .../ffi/src/execution/task_ctx_provider.rs | 5 +- datafusion/ffi/src/execution_plan.rs | 4 +- datafusion/ffi/src/session/mod.rs | 30 +- datafusion/ffi/src/tests/mod.rs | 7 + .../ffi/src/tests/object_store_provider.rs | 259 ++++++ datafusion/ffi/tests/ffi_integration.rs | 79 ++ datafusion/ffi/tests/utils/mod.rs | 20 +- .../library-user-guide/upgrading/56.0.0.md | 53 ++ 23 files changed, 5206 insertions(+), 16 deletions(-) create mode 100644 datafusion/ffi/src/execution/memory_pool.rs create mode 100644 datafusion/ffi/src/execution/object_store/buffer.rs create mode 100644 datafusion/ffi/src/execution/object_store/error.rs create mode 100644 datafusion/ffi/src/execution/object_store/mod.rs create mode 100644 datafusion/ffi/src/execution/object_store/multipart.rs create mode 100644 datafusion/ffi/src/execution/object_store/registry.rs create mode 100644 datafusion/ffi/src/execution/object_store/store.rs create mode 100644 datafusion/ffi/src/execution/object_store/stream.rs create mode 100644 datafusion/ffi/src/execution/object_store/types.rs create mode 100644 datafusion/ffi/src/execution/runtime_env.rs create mode 100644 datafusion/ffi/src/tests/object_store_provider.rs diff --git a/Cargo.lock b/Cargo.lock index d19ce6e86d18e..d1c2b1931bc98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2200,6 +2200,7 @@ dependencies = [ "arrow-schema", "async-ffi", "async-trait", + "bytes", "chrono", "datafusion", "datafusion-catalog", @@ -2223,10 +2224,12 @@ dependencies = [ "futures", "libloading", "log", + "object_store", "prost", "semver", "stabby", "tokio", + "url", ] [[package]] diff --git a/datafusion/ffi/Cargo.toml b/datafusion/ffi/Cargo.toml index affcff3dbdcd9..9ccdf7f3752a5 100644 --- a/datafusion/ffi/Cargo.toml +++ b/datafusion/ffi/Cargo.toml @@ -48,6 +48,7 @@ arrow = { workspace = true, features = ["ffi"] } arrow-schema = { workspace = true } async-ffi = { version = "0.5.0" } async-trait = { workspace = true } +bytes = { workspace = true } chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-common = { workspace = true } @@ -69,10 +70,12 @@ datafusion-session = { workspace = true } futures = { workspace = true } libloading = "0.9" log = { workspace = true } +object_store = { workspace = true } prost = { workspace = true } semver = "1.0.28" stabby = "72.1.2" tokio = { workspace = true } +url = { workspace = true } [dev-dependencies] datafusion = { workspace = true, default-features = false, features = ["sql"] } diff --git a/datafusion/ffi/README.md b/datafusion/ffi/README.md index ded54b0a88d09..de361fb4e3ad2 100644 --- a/datafusion/ffi/README.md +++ b/datafusion/ffi/README.md @@ -218,6 +218,40 @@ these methods that your provider remains valid for the lifetime of the calls. The `FFI_TaskContextProvider` is implemented on `SessionContext` and it is easy to implement on any struct that implements `Session`. +## Runtime Environment + +The `RuntimeEnv` of a `Session` crosses the boundary, so a `TableProvider` +shared over FFI can reach the object stores and memory budget of the session +that is executing it. This matters most for providers backed by remote storage: +such a provider typically builds its own `ObjectStore` and registers it during +`TableProvider::scan`, then reads through it when the returned `ExecutionPlan` +is executed. Planning and execution happen on opposite sides of the boundary, +so both must agree on the store. + +`RuntimeEnv` is a plain struct rather than a trait, and its field list depends +on enabled features, so it is not passed as an opaque pointer. Instead each of +its components crosses on its own terms: + +| Component | How it crosses | +| --------------------- | --------------------------------------------- | +| `ObjectStoreRegistry` | Shared, via `FFI_ObjectStoreRegistry` | +| `MemoryPool` | Shared, via `FFI_MemoryPool` | +| `DiskManager` | Configuration copied; each side keeps its own | +| `CacheManager` | Configuration copied; each side keeps its own | + +Because the disk and cache managers are concrete structs they cannot be shared, +only configured alike. Two consequences are worth knowing: spill limits are +enforced per side, so total on-disk usage can reach twice +`max_temp_directory_size`; and file statistics and listings cached by one side +are not reused by the other. + +A store used within the library that created it does not pay for any of this. +`FFI_ObjectStore::as_local` recovers the original `Arc`, so a +provider that registers its own store and then reads through it runs at full +speed even though the store made a round trip through the foreign registry. +Only a store genuinely owned by the _other_ library is driven through the +wrapper. + [apache datafusion]: https://datafusion.apache.org/ [api docs]: http://docs.rs/datafusion-ffi/latest [rust abi]: https://doc.rust-lang.org/reference/abi.html diff --git a/datafusion/ffi/src/execution/memory_pool.rs b/datafusion/ffi/src/execution/memory_pool.rs new file mode 100644 index 0000000000000..0b4a97db5309a --- /dev/null +++ b/datafusion/ffi/src/execution/memory_pool.rs @@ -0,0 +1,602 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`MemoryPool`]. +//! +//! Sharing the pool across the boundary means an [`ExecutionPlan`] shared over +//! FFI allocates from the same budget as the session executing it, so +//! `datafusion.execution.memory_limit` applies to it and its allocations are +//! visible in the session's accounting. +//! +//! # How reservations are tracked +//! +//! [`MemoryPool`] methods take a `&MemoryReservation`, which cannot cross the +//! boundary. What they actually need from it is the consumer identity and the +//! current size, so only the consumer's process-unique +//! [`MemoryConsumer::id`] is sent. The providing side keeps one real +//! [`MemoryReservation`] per foreign consumer and forwards `grow` / `shrink` / +//! `try_grow` onto it, which keeps the real pool's accounting exactly in step +//! with the foreign one. Dropping the entry on `unregister` releases any +//! outstanding bytes, so a foreign library that leaks a reservation cannot +//! permanently consume the host's budget beyond its own lifetime. +//! +//! Consumer ids are only unique within the library that generated them, so each +//! [`FFI_MemoryPool`] clone keeps its own map. A consumer registered through one +//! handle always grows and shrinks through that same handle, so the maps never +//! need to agree. +//! +//! [`ExecutionPlan`]: datafusion_physical_plan::ExecutionPlan + +use std::collections::HashMap; +use std::ffi::c_void; +use std::sync::{Arc, Mutex}; + +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use stabby::string::String as SString; + +/// An FFI-safe [`MemoryLimit`]. +#[repr(C, u8)] +#[derive(Debug, Clone, Copy)] +pub enum FFI_MemoryLimit { + Infinite, + Finite(u64), + Unknown, +} + +impl From for FFI_MemoryLimit { + fn from(limit: MemoryLimit) -> Self { + match limit { + MemoryLimit::Infinite => FFI_MemoryLimit::Infinite, + MemoryLimit::Finite(size) => FFI_MemoryLimit::Finite(size as u64), + MemoryLimit::Unknown => FFI_MemoryLimit::Unknown, + } + } +} + +impl From for MemoryLimit { + fn from(limit: FFI_MemoryLimit) -> Self { + match limit { + FFI_MemoryLimit::Infinite => MemoryLimit::Infinite, + FFI_MemoryLimit::Finite(size) => MemoryLimit::Finite(size as usize), + FFI_MemoryLimit::Unknown => MemoryLimit::Unknown, + } + } +} + +/// The result of a [`MemoryPool::try_grow`] call. +/// +/// [`DataFusionError::ResourcesExhausted`] is carried as its own variant rather +/// than as a generic error, so spilling operators on the far side still +/// recognise a rejected allocation as recoverable and spill to disk instead of +/// failing the query. +#[repr(C, u8)] +#[derive(Debug, Clone)] +pub enum FFI_TryGrowResult { + Ok, + ResourcesExhausted(SString), + Other(SString), +} + +impl From> for FFI_TryGrowResult { + fn from(result: Result<()>) -> Self { + match result { + Ok(()) => FFI_TryGrowResult::Ok, + Err(DataFusionError::ResourcesExhausted(msg)) => { + FFI_TryGrowResult::ResourcesExhausted(msg.as_str().into()) + } + Err(e) => FFI_TryGrowResult::Other(e.to_string().as_str().into()), + } + } +} + +impl From for Result<()> { + fn from(result: FFI_TryGrowResult) -> Self { + match result { + FFI_TryGrowResult::Ok => Ok(()), + FFI_TryGrowResult::ResourcesExhausted(msg) => { + Err(DataFusionError::ResourcesExhausted(msg.to_string())) + } + FFI_TryGrowResult::Other(msg) => datafusion_common::ffi_err!("{msg}"), + } + } +} + +/// A stable struct for sharing [`MemoryPool`] across FFI boundaries. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_MemoryPool { + /// Return the pool name. + pub name: unsafe extern "C" fn(pool: &Self) -> SString, + + /// Register a consumer, identified by its process-unique id. + pub register: unsafe extern "C" fn( + pool: &Self, + consumer_id: u64, + name: SString, + can_spill: bool, + ), + + /// Release a previously registered consumer along with any bytes it still + /// holds. + pub unregister: unsafe extern "C" fn(pool: &Self, consumer_id: u64), + + /// Infallibly grow the consumer's reservation. + pub grow: unsafe extern "C" fn(pool: &Self, consumer_id: u64, additional: u64), + + /// Infallibly shrink the consumer's reservation. + pub shrink: unsafe extern "C" fn(pool: &Self, consumer_id: u64, shrink: u64), + + /// Attempt to grow the consumer's reservation. + pub try_grow: unsafe extern "C" fn( + pool: &Self, + consumer_id: u64, + additional: u64, + ) -> FFI_TryGrowResult, + + /// Return the total number of bytes reserved across the whole pool. + pub reserved: unsafe extern "C" fn(pool: &Self) -> u64, + + /// Return the pool's memory limit. + pub memory_limit: unsafe extern "C" fn(pool: &Self) -> FFI_MemoryLimit, + + /// Used to create a clone on the provider of the pool. This should + /// only need to be called by the receiver of the pool. + pub clone: unsafe extern "C" fn(pool: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the pool. + /// The foreign library should never attempt to access this data. + pub private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. See [`crate::get_library_marker_id`] and + /// the crate's `README.md` for more information. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_MemoryPool {} +unsafe impl Sync for FFI_MemoryPool {} + +struct MemoryPoolPrivateData { + pool: Arc, + /// One real reservation per foreign consumer id. See the module docs. + reservations: Mutex>, +} + +impl MemoryPoolPrivateData { + /// Run `f` against the reservation for `consumer_id`. + /// + /// A missing entry means the foreign side grew a consumer it never + /// registered. That should not happen, since `MemoryConsumer::register` is + /// the only way to obtain a reservation, but dropping the allocation + /// silently would understate usage and defeat the memory limit. Register it + /// late instead and warn. + fn with_reservation( + &self, + consumer_id: u64, + f: impl FnOnce(&MemoryReservation) -> T, + ) -> Option { + let mut reservations = match self.reservations.lock() { + Ok(guard) => guard, + Err(e) => { + log::error!("FFI memory pool reservation map is poisoned: {e}"); + return None; + } + }; + + let reservation = reservations.entry(consumer_id).or_insert_with(|| { + log::warn!( + "Foreign memory consumer {consumer_id} was used before being registered; \ + registering it now so its usage is still accounted for" + ); + MemoryConsumer::new(format!("ffi_consumer_{consumer_id}")) + .register(&self.pool) + }); + + Some(f(reservation)) + } +} + +impl FFI_MemoryPool { + fn private_data(&self) -> &MemoryPoolPrivateData { + unsafe { &*(self.private_data as *const MemoryPoolPrivateData) } + } + + /// Create a new [`FFI_MemoryPool`] from a local pool. + pub fn new(pool: Arc) -> Self { + Self { + name: name_fn_wrapper, + register: register_fn_wrapper, + unregister: unregister_fn_wrapper, + grow: grow_fn_wrapper, + shrink: shrink_fn_wrapper, + try_grow: try_grow_fn_wrapper, + reserved: reserved_fn_wrapper, + memory_limit: memory_limit_fn_wrapper, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + private_data: Box::into_raw(Box::new(MemoryPoolPrivateData { + pool, + reservations: Mutex::new(HashMap::new()), + })) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } + + /// If this pool originated in the current library, return the underlying + /// [`MemoryPool`] directly. + pub fn as_local(&self) -> Option> { + ((self.library_marker_id)() == crate::get_library_marker_id()) + .then(|| Arc::clone(&self.private_data().pool)) + } +} + +unsafe extern "C" fn name_fn_wrapper(pool: &FFI_MemoryPool) -> SString { + pool.private_data().pool.name().into() +} + +unsafe extern "C" fn register_fn_wrapper( + pool: &FFI_MemoryPool, + consumer_id: u64, + name: SString, + can_spill: bool, +) { + let private_data = pool.private_data(); + let reservation = MemoryConsumer::new(name.to_string()) + .with_can_spill(can_spill) + .register(&private_data.pool); + + match private_data.reservations.lock() { + Ok(mut reservations) => { + reservations.insert(consumer_id, reservation); + } + Err(e) => log::error!("FFI memory pool reservation map is poisoned: {e}"), + } +} + +unsafe extern "C" fn unregister_fn_wrapper(pool: &FFI_MemoryPool, consumer_id: u64) { + let private_data = pool.private_data(); + match private_data.reservations.lock() { + // Dropping the reservation frees any outstanding bytes back to the pool + // and unregisters the consumer. + Ok(mut reservations) => { + reservations.remove(&consumer_id); + } + Err(e) => log::error!("FFI memory pool reservation map is poisoned: {e}"), + } +} + +unsafe extern "C" fn grow_fn_wrapper( + pool: &FFI_MemoryPool, + consumer_id: u64, + additional: u64, +) { + pool.private_data() + .with_reservation(consumer_id, |reservation| { + reservation.grow(additional as usize) + }); +} + +unsafe extern "C" fn shrink_fn_wrapper( + pool: &FFI_MemoryPool, + consumer_id: u64, + shrink: u64, +) { + pool.private_data() + .with_reservation(consumer_id, |reservation| { + // `MemoryReservation::shrink` panics if asked to free more than it + // holds. The foreign side tracks its own size and should never do + // that, but a panic across the FFI boundary is undefined behaviour, + // so clamp instead. + let capacity = (shrink as usize).min(reservation.size()); + if capacity > 0 { + reservation.shrink(capacity); + } + }); +} + +unsafe extern "C" fn try_grow_fn_wrapper( + pool: &FFI_MemoryPool, + consumer_id: u64, + additional: u64, +) -> FFI_TryGrowResult { + pool.private_data() + .with_reservation(consumer_id, |reservation| { + FFI_TryGrowResult::from(reservation.try_grow(additional as usize)) + }) + .unwrap_or_else(|| { + FFI_TryGrowResult::Other("FFI memory pool is unavailable".into()) + }) +} + +unsafe extern "C" fn reserved_fn_wrapper(pool: &FFI_MemoryPool) -> u64 { + pool.private_data().pool.reserved() as u64 +} + +unsafe extern "C" fn memory_limit_fn_wrapper(pool: &FFI_MemoryPool) -> FFI_MemoryLimit { + pool.private_data().pool.memory_limit().into() +} + +unsafe extern "C" fn clone_fn_wrapper(pool: &FFI_MemoryPool) -> FFI_MemoryPool { + // A clone gets its own reservation map: consumers registered through one + // handle always grow and shrink through that same handle. + FFI_MemoryPool::new(Arc::clone(&pool.private_data().pool)) +} + +unsafe extern "C" fn release_fn_wrapper(pool: &mut FFI_MemoryPool) { + unsafe { + debug_assert!(!pool.private_data.is_null()); + drop(Box::from_raw( + pool.private_data as *mut MemoryPoolPrivateData, + )); + pool.private_data = std::ptr::null_mut(); + } +} + +impl Clone for FFI_MemoryPool { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +impl Drop for FFI_MemoryPool { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +/// A [`MemoryPool`] backed by a foreign [`FFI_MemoryPool`]. +#[derive(Debug)] +pub struct ForeignMemoryPool { + pool: FFI_MemoryPool, + /// The underlying pool's name, fetched once at construction. + /// + /// [`MemoryPool::name`] returns a borrowed `&str`, but the name arrives + /// owned from across the boundary and so cannot be fetched per call. The + /// name appears in resource-exhaustion messages, so keeping the real one + /// makes those messages name the pool that actually rejected the + /// allocation. + name: String, +} + +unsafe impl Send for ForeignMemoryPool {} +unsafe impl Sync for ForeignMemoryPool {} + +impl From for ForeignMemoryPool { + fn from(pool: FFI_MemoryPool) -> Self { + let name = unsafe { (pool.name)(&pool) }.to_string(); + Self { pool, name } + } +} + +impl From for Arc { + fn from(pool: FFI_MemoryPool) -> Self { + match pool.as_local() { + Some(local) => local, + None => Arc::new(ForeignMemoryPool::from(pool)), + } + } +} + +impl std::fmt::Display for ForeignMemoryPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.name) + } +} + +impl MemoryPool for ForeignMemoryPool { + fn name(&self) -> &str { + &self.name + } + + fn register(&self, consumer: &MemoryConsumer) { + unsafe { + (self.pool.register)( + &self.pool, + consumer.id() as u64, + consumer.name().into(), + consumer.can_spill(), + ) + } + } + + fn unregister(&self, consumer: &MemoryConsumer) { + unsafe { (self.pool.unregister)(&self.pool, consumer.id() as u64) } + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + unsafe { + (self.pool.grow)( + &self.pool, + reservation.consumer().id() as u64, + additional as u64, + ) + } + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + unsafe { + (self.pool.shrink)( + &self.pool, + reservation.consumer().id() as u64, + shrink as u64, + ) + } + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + let result = unsafe { + (self.pool.try_grow)( + &self.pool, + reservation.consumer().id() as u64, + additional as u64, + ) + }; + result.into() + } + + fn reserved(&self) -> usize { + unsafe { (self.pool.reserved)(&self.pool) as usize } + } + + fn memory_limit(&self) -> MemoryLimit { + unsafe { (self.pool.memory_limit)(&self.pool) }.into() + } +} + +#[cfg(test)] +mod tests { + use datafusion_execution::memory_pool::{GreedyMemoryPool, UnboundedMemoryPool}; + + use super::*; + + fn foreign_pool(pool: Arc) -> Arc { + let mut ffi = FFI_MemoryPool::new(pool); + ffi.library_marker_id = crate::mock_foreign_marker_id; + Arc::new(ForeignMemoryPool::from(ffi)) + } + + #[test] + fn grow_and_shrink_are_reflected_in_the_real_pool() { + let real = Arc::new(GreedyMemoryPool::new(1_000)) as Arc; + let foreign = foreign_pool(Arc::clone(&real)); + + let reservation = MemoryConsumer::new("test").register(&foreign); + reservation.grow(400); + assert_eq!(real.reserved(), 400); + assert_eq!(foreign.reserved(), 400); + + reservation.shrink(150); + assert_eq!(real.reserved(), 250); + + drop(reservation); + assert_eq!(real.reserved(), 0); + } + + /// The point of forwarding the pool: a limit configured on the host must + /// actually constrain a foreign consumer. + #[test] + fn memory_limit_is_enforced_across_the_boundary() { + let real = Arc::new(GreedyMemoryPool::new(1_000)) as Arc; + let foreign = foreign_pool(Arc::clone(&real)); + + let reservation = MemoryConsumer::new("test").register(&foreign); + assert!(reservation.try_grow(900).is_ok()); + + let err = reservation.try_grow(200).unwrap_err(); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "expected ResourcesExhausted so operators spill, got {err:?}" + ); + } + + #[test] + fn memory_limit_is_reported() { + let real = Arc::new(GreedyMemoryPool::new(4_096)) as Arc; + let foreign = foreign_pool(real); + assert!(matches!(foreign.memory_limit(), MemoryLimit::Finite(4_096))); + + let unbounded = foreign_pool(Arc::new(UnboundedMemoryPool::default())); + assert!(matches!( + unbounded.memory_limit(), + MemoryLimit::Infinite | MemoryLimit::Unknown + )); + } + + /// Dropping a reservation on the foreign side must release the bytes it + /// held on the host side, otherwise a foreign plan would leak the host's + /// budget. + #[test] + fn unregister_releases_outstanding_bytes() { + let real = Arc::new(GreedyMemoryPool::new(1_000)) as Arc; + let foreign = foreign_pool(Arc::clone(&real)); + + let reservation = MemoryConsumer::new("leaky").register(&foreign); + reservation.grow(500); + assert_eq!(real.reserved(), 500); + + // Drop without shrinking first. + drop(reservation); + assert_eq!(real.reserved(), 0); + } + + #[test] + fn multiple_consumers_are_tracked_independently() { + let real = Arc::new(GreedyMemoryPool::new(1_000)) as Arc; + let foreign = foreign_pool(Arc::clone(&real)); + + let a = MemoryConsumer::new("a").register(&foreign); + let b = MemoryConsumer::new("b").register(&foreign); + + a.grow(300); + b.grow(200); + assert_eq!(real.reserved(), 500); + + drop(a); + assert_eq!(real.reserved(), 200); + drop(b); + assert_eq!(real.reserved(), 0); + } + + #[test] + fn shrink_beyond_size_does_not_panic() { + let real = Arc::new(GreedyMemoryPool::new(1_000)) as Arc; + let foreign = foreign_pool(Arc::clone(&real)); + + let ffi = FFI_MemoryPool::new(Arc::clone(&real)); + // Grow by 10 then ask the wrapper to shrink by far more. A panic here + // would be undefined behaviour across a real FFI boundary. + unsafe { + (ffi.register)(&ffi, 42, "clamped".into(), false); + (ffi.grow)(&ffi, 42, 10); + (ffi.shrink)(&ffi, 42, 10_000); + } + assert_eq!(real.reserved(), 0); + + drop(foreign); + } + + #[test] + fn local_pool_is_unwrapped() { + let original = Arc::new(GreedyMemoryPool::new(1_000)) as Arc; + let ffi = FFI_MemoryPool::new(Arc::clone(&original)); + + let recovered = ffi.as_local().expect("local pool should unwrap"); + assert!(Arc::ptr_eq(&original, &recovered)); + } + + #[test] + fn try_grow_result_round_trip() { + let exhausted: Result<()> = FFI_TryGrowResult::from(Err::<(), _>( + DataFusionError::ResourcesExhausted("over budget".to_string()), + )) + .into(); + match exhausted.unwrap_err() { + DataFusionError::ResourcesExhausted(msg) => assert_eq!(msg, "over budget"), + other => panic!("expected ResourcesExhausted, got {other:?}"), + } + + let ok: Result<()> = + FFI_TryGrowResult::from(Ok::<(), DataFusionError>(())).into(); + assert!(ok.is_ok()); + } +} diff --git a/datafusion/ffi/src/execution/mod.rs b/datafusion/ffi/src/execution/mod.rs index 41107947fff01..3b7a468bf3bc6 100644 --- a/datafusion/ffi/src/execution/mod.rs +++ b/datafusion/ffi/src/execution/mod.rs @@ -15,8 +15,17 @@ // specific language governing permissions and limitations // under the License. +pub mod memory_pool; +pub mod object_store; +pub mod runtime_env; mod task_ctx; pub mod task_ctx_provider; +pub use memory_pool::{FFI_MemoryPool, FFI_TryGrowResult, ForeignMemoryPool}; +pub use object_store::{ + FFI_ObjectStore, FFI_ObjectStoreRegistry, ForeignObjectStore, + ForeignObjectStoreRegistry, +}; +pub use runtime_env::{FFI_RuntimeConfig, FFI_RuntimeEnv}; pub use task_ctx::FFI_TaskContext; pub use task_ctx_provider::FFI_TaskContextProvider; diff --git a/datafusion/ffi/src/execution/object_store/buffer.rs b/datafusion/ffi/src/execution/object_store/buffer.rs new file mode 100644 index 0000000000000..a80a2e65a3467 --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/buffer.rs @@ -0,0 +1,177 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI-safe transfer of [`Bytes`] buffers. + +use std::ffi::c_void; + +use bytes::Bytes; + +/// A stable struct for sharing a byte buffer across FFI boundaries. +/// +/// The buffer is *borrowed* from the producing library for the lifetime of this +/// struct: `private_data` owns a boxed [`Bytes`] and `ptr` points into it. +/// Dropping this struct calls back into the producing library to release that +/// allocation, so the bytes are never freed by the wrong allocator. +/// +/// Because the memory stays valid until release, the consuming side can build a +/// [`Bytes`] that borrows it directly with [`Bytes::from_owner`] rather than +/// copying. This matters on the read path, where every scanned byte of every +/// data file crosses this struct. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_Bytes { + /// Pointer to the start of the buffer. Valid for `len` bytes until + /// `release` is called. Never null, even for an empty buffer. + pub ptr: *const u8, + + /// Length of the buffer in bytes. + pub len: u64, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// buffer. The foreign library should never attempt to access this data. + pub private_data: *mut c_void, +} + +// Safety: the buffer is immutable and owned by the boxed `Bytes` in +// `private_data`, which is itself `Send + Sync`. Access is read-only. +unsafe impl Send for FFI_Bytes {} +unsafe impl Sync for FFI_Bytes {} + +unsafe extern "C" fn release_fn_wrapper(buffer: &mut FFI_Bytes) { + unsafe { + debug_assert!(!buffer.private_data.is_null()); + drop(Box::from_raw(buffer.private_data as *mut Bytes)); + buffer.private_data = std::ptr::null_mut(); + buffer.ptr = std::ptr::NonNull::dangling().as_ptr(); + buffer.len = 0; + } +} + +impl From for FFI_Bytes { + fn from(bytes: Bytes) -> Self { + let len = bytes.len() as u64; + let boxed = Box::new(bytes); + // Take the pointer from the boxed value so it remains valid after the + // move into the box. + let ptr = if boxed.is_empty() { + // `Bytes::as_ptr` on an empty buffer may be dangling but is still + // required to be non-null for `slice::from_raw_parts`. + std::ptr::NonNull::dangling().as_ptr() + } else { + boxed.as_ptr() + }; + + Self { + ptr, + len, + release: release_fn_wrapper, + private_data: Box::into_raw(boxed) as *mut c_void, + } + } +} + +impl AsRef<[u8]> for FFI_Bytes { + fn as_ref(&self) -> &[u8] { + // Safety: `ptr` is non-null and valid for `len` bytes until `release`, + // which only runs in `Drop`. + unsafe { std::slice::from_raw_parts(self.ptr, self.len as usize) } + } +} + +impl From for Bytes { + fn from(buffer: FFI_Bytes) -> Self { + let release_ptr = buffer.release as *const (); + if std::ptr::eq(release_ptr, release_fn_wrapper as *const ()) { + // Same library: unwrap the original `Bytes` and keep its refcount + // rather than adding an owner indirection. + let mut buffer = std::mem::ManuallyDrop::new(buffer); + let bytes = unsafe { *Box::from_raw(buffer.private_data as *mut Bytes) }; + buffer.private_data = std::ptr::null_mut(); + return bytes; + } + + // Foreign library: borrow the buffer in place. The `FFI_Bytes` is kept + // alive by the resulting `Bytes` and released when its refcount drops. + Bytes::from_owner(buffer) + } +} + +impl Drop for FFI_Bytes { + fn drop(&mut self) { + if !self.private_data.is_null() { + unsafe { (self.release)(self) } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn bytes_round_trip() { + let original = Bytes::from_static(b"hello world"); + let ffi = FFI_Bytes::from(original.clone()); + assert_eq!(ffi.len, 11); + assert_eq!(ffi.as_ref(), b"hello world"); + + let restored: Bytes = ffi.into(); + assert_eq!(restored, original); + } + + #[test] + fn empty_bytes_round_trip() { + let ffi = FFI_Bytes::from(Bytes::new()); + assert_eq!(ffi.len, 0); + assert!(!ffi.ptr.is_null()); + assert_eq!(ffi.as_ref(), b""); + + let restored: Bytes = ffi.into(); + assert!(restored.is_empty()); + } + + /// Simulate the foreign path, where the release function pointer belongs to + /// another library and the buffer must be borrowed rather than unwrapped. + #[test] + fn foreign_bytes_borrow_round_trip() { + unsafe extern "C" fn foreign_release(buffer: &mut FFI_Bytes) { + unsafe { + drop(Box::from_raw(buffer.private_data as *mut Bytes)); + buffer.private_data = std::ptr::null_mut(); + } + } + + let original = Bytes::from(vec![1u8, 2, 3, 4, 5]); + let mut ffi = FFI_Bytes::from(original.clone()); + ffi.release = foreign_release; + + let restored: Bytes = ffi.into(); + assert_eq!(restored, original); + // Dropping the borrowed `Bytes` must run the foreign release. + drop(restored); + } + + #[test] + fn dropping_without_conversion_releases() { + let ffi = FFI_Bytes::from(Bytes::from(vec![0u8; 64])); + drop(ffi); + } +} diff --git a/datafusion/ffi/src/execution/object_store/error.rs b/datafusion/ffi/src/execution/object_store/error.rs new file mode 100644 index 0000000000000..2e50f3193aa9e --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/error.rs @@ -0,0 +1,415 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI-safe representation of [`object_store::Error`]. +//! +//! The *variant* of an object store error is preserved across the boundary, +//! not just its message, so code that branches on the variant behaves the same +//! either side. That matters for callers implementing optimistic concurrency +//! control on top of [`ObjectStore::rename_if_not_exists`] or +//! [`PutMode::Create`], which detect a conflicting write via +//! [`object_store::Error::AlreadyExists`], and for conditional reads, which +//! rely on [`object_store::Error::NotModified`] and +//! [`object_store::Error::Precondition`]. +//! +//! The `source` of an error is a `Box` and cannot cross +//! the boundary, so it is rendered to a string. The chain of causes is +//! therefore flattened into the message but never dropped entirely. +//! +//! [`ObjectStore::rename_if_not_exists`]: object_store::ObjectStore::rename_if_not_exists +//! [`PutMode::Create`]: object_store::PutMode::Create + +use object_store::Error as ObjectStoreError; +use object_store::path::Error as PathError; +use stabby::string::String as SString; + +/// The variant of an [`object_store::Error`], preserved across the FFI +/// boundary. +/// +/// Values are explicitly assigned so that adding a variant cannot silently +/// renumber the existing ones. +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[expect(non_camel_case_types)] +pub enum FFI_ObjectStoreErrorKind { + Generic = 0, + NotFound = 1, + InvalidPath = 2, + JoinError = 3, + NotSupported = 4, + AlreadyExists = 5, + Precondition = 6, + NotModified = 7, + NotImplemented = 8, + PermissionDenied = 9, + Unauthenticated = 10, + UnknownConfigurationKey = 11, +} + +/// A stable struct for sharing [`object_store::Error`] across FFI boundaries. +/// +/// `path` carries the location for the variants that have one. For +/// [`FFI_ObjectStoreErrorKind::UnknownConfigurationKey`] it carries the +/// configuration key and for [`FFI_ObjectStoreErrorKind::NotImplemented`] it +/// carries the operation, since those variants use those fields in place of a +/// path. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_ObjectStoreError { + /// Which [`object_store::Error`] variant this is. + pub kind: FFI_ObjectStoreErrorKind, + + /// The path, configuration key, or operation associated with the error. + /// Empty when the variant has none. + pub path: SString, + + /// The rendered source error, or the store / implementer name for the + /// variants that carry one instead of a source. + pub message: SString, +} + +impl From<&ObjectStoreError> for FFI_ObjectStoreError { + fn from(err: &ObjectStoreError) -> Self { + use FFI_ObjectStoreErrorKind as Kind; + + let (kind, path, message) = match err { + ObjectStoreError::Generic { store, source } => { + (Kind::Generic, (*store).to_owned(), source.to_string()) + } + ObjectStoreError::NotFound { path, source } => { + (Kind::NotFound, path.clone(), source.to_string()) + } + ObjectStoreError::InvalidPath { source } => { + (Kind::InvalidPath, String::new(), source.to_string()) + } + ObjectStoreError::NotSupported { source } => { + (Kind::NotSupported, String::new(), source.to_string()) + } + ObjectStoreError::AlreadyExists { path, source } => { + (Kind::AlreadyExists, path.clone(), source.to_string()) + } + ObjectStoreError::Precondition { path, source } => { + (Kind::Precondition, path.clone(), source.to_string()) + } + ObjectStoreError::NotModified { path, source } => { + (Kind::NotModified, path.clone(), source.to_string()) + } + ObjectStoreError::NotImplemented { + operation, + implementer, + } => (Kind::NotImplemented, operation.clone(), implementer.clone()), + ObjectStoreError::PermissionDenied { path, source } => { + (Kind::PermissionDenied, path.clone(), source.to_string()) + } + ObjectStoreError::Unauthenticated { path, source } => { + (Kind::Unauthenticated, path.clone(), source.to_string()) + } + ObjectStoreError::UnknownConfigurationKey { store, key } => ( + Kind::UnknownConfigurationKey, + key.clone(), + (*store).to_owned(), + ), + // `JoinError` is `#[cfg(feature = "tokio")]` and any variant added + // to `object_store::Error` in a future release lands here. Preserve + // the rendered message under the fallback variant. + other => (Kind::Generic, String::new(), other.to_string()), + }; + + Self { + kind, + path: path.as_str().into(), + message: message.as_str().into(), + } + } +} + +impl From for FFI_ObjectStoreError { + fn from(err: ObjectStoreError) -> Self { + Self::from(&err) + } +} + +impl From for ObjectStoreError { + fn from(err: FFI_ObjectStoreError) -> Self { + use FFI_ObjectStoreErrorKind as Kind; + + let path = err.path.to_string(); + let message = err.message.to_string(); + // `source` is a boxed error, so the original chain is rebuilt as a + // single opaque string error carrying the rendered chain. + let source = || -> Box { + message.clone().into() + }; + + match err.kind { + Kind::Generic => ObjectStoreError::Generic { + // `store` is a `&'static str`, so the original value cannot be + // rebuilt. Fold it into the message instead of leaking it. + store: "ForeignObjectStore", + source: if path.is_empty() { + source() + } else { + format!("{path}: {message}").into() + }, + }, + Kind::NotFound => ObjectStoreError::NotFound { + path, + source: source(), + }, + Kind::InvalidPath => ObjectStoreError::InvalidPath { + source: PathError::InvalidPath { + path: message.into(), + }, + }, + // `tokio::task::JoinError` cannot be constructed outside tokio, so + // a join failure surfaces as a generic error with its message. + Kind::JoinError | Kind::NotSupported => { + ObjectStoreError::NotSupported { source: source() } + } + Kind::AlreadyExists => ObjectStoreError::AlreadyExists { + path, + source: source(), + }, + Kind::Precondition => ObjectStoreError::Precondition { + path, + source: source(), + }, + Kind::NotModified => ObjectStoreError::NotModified { + path, + source: source(), + }, + Kind::NotImplemented => ObjectStoreError::NotImplemented { + operation: path, + implementer: message, + }, + Kind::PermissionDenied => ObjectStoreError::PermissionDenied { + path, + source: source(), + }, + Kind::Unauthenticated => ObjectStoreError::Unauthenticated { + path, + source: source(), + }, + Kind::UnknownConfigurationKey => ObjectStoreError::UnknownConfigurationKey { + store: "ForeignObjectStore", + key: path, + }, + } + } +} + +/// An FFI-safe result carrying an [`FFI_ObjectStoreError`]. +/// +/// This mirrors [`crate::util::FFI_Result`] but preserves the object store +/// error variant rather than reducing it to a message. +#[repr(C, u8)] +#[derive(Debug, Clone)] +pub enum FFI_ObjectStoreResult { + Ok(T), + Err(FFI_ObjectStoreError), +} + +impl From> for FFI_ObjectStoreResult { + fn from(res: Result) -> Self { + match res { + Ok(v) => FFI_ObjectStoreResult::Ok(v), + Err(e) => FFI_ObjectStoreResult::Err(e.into()), + } + } +} + +impl From> for Result { + fn from(res: FFI_ObjectStoreResult) -> Self { + match res { + FFI_ObjectStoreResult::Ok(v) => Ok(v), + FFI_ObjectStoreResult::Err(e) => Err(e.into()), + } + } +} + +impl FFI_ObjectStoreResult { + pub fn map U>(self, f: F) -> FFI_ObjectStoreResult { + match self { + FFI_ObjectStoreResult::Ok(v) => FFI_ObjectStoreResult::Ok(f(v)), + FFI_ObjectStoreResult::Err(e) => FFI_ObjectStoreResult::Err(e), + } + } +} + +/// Convert a [`Result`] into an +/// [`FFI_ObjectStoreResult`], returning early on error. Mirrors +/// [`crate::sresult_return`], which cannot be used here because it discards the +/// error variant. +#[macro_export] +macro_rules! os_result_return { + ( $x:expr ) => { + match $x { + Ok(v) => v, + Err(e) => { + return $crate::execution::object_store::FFI_ObjectStoreResult::Err( + e.into(), + ); + } + } + }; +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Every variant must survive the round trip, because callers branch on + /// the variant to implement conditional writes. + #[test] + fn error_kind_round_trip() { + use FFI_ObjectStoreErrorKind as Kind; + + let cases: Vec<(ObjectStoreError, FFI_ObjectStoreErrorKind)> = vec![ + ( + ObjectStoreError::NotFound { + path: "a/b".into(), + source: "missing".into(), + }, + FFI_ObjectStoreErrorKind::NotFound, + ), + ( + ObjectStoreError::AlreadyExists { + path: "a/b".into(), + source: "exists".into(), + }, + FFI_ObjectStoreErrorKind::AlreadyExists, + ), + ( + ObjectStoreError::Precondition { + path: "a/b".into(), + source: "etag".into(), + }, + FFI_ObjectStoreErrorKind::Precondition, + ), + ( + ObjectStoreError::NotModified { + path: "a/b".into(), + source: "same".into(), + }, + FFI_ObjectStoreErrorKind::NotModified, + ), + ( + ObjectStoreError::PermissionDenied { + path: "a/b".into(), + source: "denied".into(), + }, + FFI_ObjectStoreErrorKind::PermissionDenied, + ), + ( + ObjectStoreError::Unauthenticated { + path: "a/b".into(), + source: "no creds".into(), + }, + FFI_ObjectStoreErrorKind::Unauthenticated, + ), + ( + ObjectStoreError::NotSupported { + source: "nope".into(), + }, + FFI_ObjectStoreErrorKind::NotSupported, + ), + ( + ObjectStoreError::NotImplemented { + operation: "put_multipart_opts".into(), + implementer: "TestStore".into(), + }, + FFI_ObjectStoreErrorKind::NotImplemented, + ), + ( + ObjectStoreError::UnknownConfigurationKey { + store: "S3", + key: "bogus".into(), + }, + FFI_ObjectStoreErrorKind::UnknownConfigurationKey, + ), + ( + ObjectStoreError::Generic { + store: "S3", + source: "boom".into(), + }, + FFI_ObjectStoreErrorKind::Generic, + ), + ]; + + for (original, expected_kind) in cases { + let ffi = FFI_ObjectStoreError::from(&original); + assert_eq!(ffi.kind, expected_kind, "for {original}"); + + let restored = ObjectStoreError::from(ffi.clone()); + let restored_ffi = FFI_ObjectStoreError::from(&restored); + assert_eq!( + restored_ffi.kind, expected_kind, + "kind changed on rebuild for {original}" + ); + + match expected_kind { + // `Generic` carries a `store: &'static str` rather than a path. + // A `&'static str` cannot be rebuilt on the far side, so the + // originating store name is folded into the message instead. + Kind::Generic => assert!( + restored.to_string().contains("S3"), + "originating store name should survive in the message, got {restored}" + ), + // Every other variant carries a path (or, for + // `UnknownConfigurationKey` and `NotImplemented`, the field that + // takes its place), and callers surface it to users. + _ => { + assert_eq!(restored_ffi.path, ffi.path, "path changed for {original}") + } + } + } + } + + #[test] + fn error_preserves_path_and_message() { + let original = ObjectStoreError::AlreadyExists { + path: "_delta_log/00000000000000000001.json".into(), + source: "conditional put failed".into(), + }; + + let restored = ObjectStoreError::from(FFI_ObjectStoreError::from(&original)); + + let ObjectStoreError::AlreadyExists { path, source } = &restored else { + panic!("expected AlreadyExists, got {restored:?}"); + }; + assert_eq!(path, "_delta_log/00000000000000000001.json"); + assert_eq!(source.to_string(), "conditional put failed"); + } + + #[test] + fn result_round_trip() { + let ok: FFI_ObjectStoreResult = Ok(7).into(); + let ok: Result = ok.into(); + assert_eq!(ok.unwrap(), 7); + + let err: FFI_ObjectStoreResult = Err(ObjectStoreError::NotFound { + path: "x".into(), + source: "gone".into(), + }) + .into(); + let err: Result = err.into(); + assert!(matches!( + err.unwrap_err(), + ObjectStoreError::NotFound { .. } + )); + } +} diff --git a/datafusion/ffi/src/execution/object_store/mod.rs b/datafusion/ffi/src/execution/object_store/mod.rs new file mode 100644 index 0000000000000..3905bd8f7c89c --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/mod.rs @@ -0,0 +1,360 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`object_store::ObjectStore`] and +//! [`datafusion_execution::object_store::ObjectStoreRegistry`]. +//! +//! # Sharing the registry +//! +//! A table provider shared over FFI is asked to `scan` by a session in another +//! library, and the [`ExecutionPlan`] it returns is later executed with a +//! [`TaskContext`] produced by that same session. Both sides need to agree on +//! which object store serves a given URL. Sharing the registry is what makes +//! that agreement possible: a store registered by the provider during planning +//! is visible to the session at execution time, and a store registered on the +//! session by the host is visible to the provider. +//! +//! # The local fast path +//! +//! A provider that creates its own store, registers it, and then reads through +//! it during execution never actually moves data across the boundary. The store +//! makes a round trip through the foreign registry but +//! [`FFI_ObjectStore::as_local`] recovers the original `Arc`, +//! so reads run at full speed. Only a store genuinely owned by the *other* +//! library is driven through the wrapper. +//! +//! [`ExecutionPlan`]: datafusion_physical_plan::ExecutionPlan +//! [`TaskContext`]: datafusion_execution::TaskContext + +mod buffer; +mod error; +mod multipart; +mod registry; +mod store; +mod stream; +mod types; + +pub use buffer::FFI_Bytes; +pub use error::{FFI_ObjectStoreError, FFI_ObjectStoreErrorKind, FFI_ObjectStoreResult}; +pub use multipart::{FFI_MultipartUpload, ForeignMultipartUpload}; +pub use registry::{FFI_ObjectStoreRegistry, ForeignObjectStoreRegistry}; +pub use store::{FFI_ByteRange, FFI_GetResult, FFI_ObjectStore, ForeignObjectStore}; +pub use stream::{FFI_BytesStream, FFI_ObjectMetaStream, FFI_PathStream}; +pub use types::{ + FFI_Attribute, FFI_CopyOptions, FFI_GetOptions, FFI_GetRange, FFI_ListResult, + FFI_ObjectMeta, FFI_PutMode, FFI_PutMultipartOptions, FFI_PutOptions, FFI_PutResult, + FFI_RenameOptions, FFI_Timestamp, +}; + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use bytes::Bytes; + use futures::StreamExt; + use object_store::memory::InMemory; + use object_store::path::Path; + use object_store::{ + Error as ObjectStoreError, ObjectStore, ObjectStoreExt, PutMode, PutOptions, + PutPayload, + }; + + use super::*; + + /// Wrap a store and force the foreign code path so the FFI functions are + /// exercised rather than the local unwrap. + fn foreign_store(store: Arc) -> Arc { + let mut ffi = FFI_ObjectStore::new(store, None); + ffi.library_marker_id = crate::mock_foreign_marker_id; + Arc::new(ForeignObjectStore::from(ffi)) + } + + #[tokio::test] + async fn put_and_get_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("a/b.parquet"); + + store + .put(&path, PutPayload::from_static(b"hello ffi")) + .await?; + + let result = store.get(&path).await?; + assert_eq!(result.bytes().await?, Bytes::from_static(b"hello ffi")); + + Ok(()) + } + + #[tokio::test] + async fn get_ranges_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("data.bin"); + store + .put(&path, PutPayload::from_static(b"0123456789")) + .await?; + + let ranges = store.get_ranges(&path, &[0..3, 5..8]).await?; + assert_eq!(ranges.len(), 2); + assert_eq!(ranges[0], Bytes::from_static(b"012")); + assert_eq!(ranges[1], Bytes::from_static(b"567")); + + Ok(()) + } + + #[tokio::test] + async fn get_range_option_is_honored() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("data.bin"); + store + .put(&path, PutPayload::from_static(b"0123456789")) + .await?; + + let bytes = store.get_range(&path, 2..6).await?; + assert_eq!(bytes, Bytes::from_static(b"2345")); + + Ok(()) + } + + #[tokio::test] + async fn head_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("a/b.parquet"); + store.put(&path, PutPayload::from_static(b"12345")).await?; + + let meta = store.head(&path).await?; + assert_eq!(meta.location, path); + assert_eq!(meta.size, 5); + + Ok(()) + } + + #[tokio::test] + async fn list_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + for name in ["p/1.parquet", "p/2.parquet", "q/3.parquet"] { + store + .put(&Path::from(name), PutPayload::from_static(b"x")) + .await?; + } + + let mut listed: Vec = store + .list(Some(&Path::from("p"))) + .map(|m| m.unwrap().location.to_string()) + .collect() + .await; + listed.sort(); + + assert_eq!(listed, vec!["p/1.parquet", "p/2.parquet"]); + + Ok(()) + } + + #[tokio::test] + async fn list_with_delimiter_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + for name in ["p/1.parquet", "p/sub/2.parquet"] { + store + .put(&Path::from(name), PutPayload::from_static(b"x")) + .await?; + } + + let result = store.list_with_delimiter(Some(&Path::from("p"))).await?; + assert_eq!(result.objects.len(), 1); + assert_eq!(result.objects[0].location, Path::from("p/1.parquet")); + assert_eq!(result.common_prefixes, vec![Path::from("p/sub")]); + + Ok(()) + } + + #[tokio::test] + async fn delete_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("gone.parquet"); + store.put(&path, PutPayload::from_static(b"x")).await?; + store.delete(&path).await?; + + assert!(matches!( + store.head(&path).await.unwrap_err(), + ObjectStoreError::NotFound { .. } + )); + + Ok(()) + } + + #[tokio::test] + async fn copy_and_rename_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let src = Path::from("src.parquet"); + store.put(&src, PutPayload::from_static(b"payload")).await?; + + let copied = Path::from("copied.parquet"); + store.copy(&src, &copied).await?; + assert_eq!(store.get(&copied).await?.bytes().await?, "payload"); + + let renamed = Path::from("renamed.parquet"); + store.rename(&copied, &renamed).await?; + assert!(matches!( + store.head(&copied).await.unwrap_err(), + ObjectStoreError::NotFound { .. } + )); + assert_eq!(store.get(&renamed).await?.bytes().await?, "payload"); + + Ok(()) + } + + /// `copy_if_not_exists` is an extension method built on `copy_opts` with + /// `CopyMode::Create`. Its atomicity depends on the mode surviving the + /// boundary and on `AlreadyExists` being reported as that variant. + #[tokio::test] + async fn copy_if_not_exists_preserves_already_exists() -> datafusion_common::Result<()> + { + let store = foreign_store(Arc::new(InMemory::new())); + let src = Path::from("src.parquet"); + let dst = Path::from("dst.parquet"); + store.put(&src, PutPayload::from_static(b"a")).await?; + store.put(&dst, PutPayload::from_static(b"b")).await?; + + let err = store.copy_if_not_exists(&src, &dst).await.unwrap_err(); + assert!( + matches!(err, ObjectStoreError::AlreadyExists { .. }), + "expected AlreadyExists, got {err:?}" + ); + + Ok(()) + } + + /// `PutMode::Create` underpins optimistic concurrency control in commit + /// protocols such as Delta's. A conflicting write must surface as + /// `AlreadyExists`, not as a generic error. + #[tokio::test] + async fn put_mode_create_preserves_already_exists() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("_delta_log/00000000000000000001.json"); + + let opts = PutOptions { + mode: PutMode::Create, + ..Default::default() + }; + store + .put_opts(&path, PutPayload::from_static(b"first"), opts.clone()) + .await?; + + let err = store + .put_opts(&path, PutPayload::from_static(b"second"), opts) + .await + .unwrap_err(); + assert!( + matches!(err, ObjectStoreError::AlreadyExists { .. }), + "expected AlreadyExists, got {err:?}" + ); + + Ok(()) + } + + #[tokio::test] + async fn not_found_preserves_variant() { + let store = foreign_store(Arc::new(InMemory::new())); + let err = store.get(&Path::from("missing")).await.unwrap_err(); + assert!( + matches!(err, ObjectStoreError::NotFound { .. }), + "expected NotFound, got {err:?}" + ); + } + + #[tokio::test] + async fn multipart_upload_round_trip() -> datafusion_common::Result<()> { + let store = foreign_store(Arc::new(InMemory::new())); + let path = Path::from("multipart.bin"); + + let mut upload = store.put_multipart(&path).await?; + upload.put_part(PutPayload::from(vec![1u8; 8])).await?; + upload.put_part(PutPayload::from(vec![2u8; 8])).await?; + upload.complete().await?; + + let bytes = store.get(&path).await?.bytes().await?; + assert_eq!(bytes.len(), 16); + assert_eq!(&bytes[..8], &[1u8; 8]); + assert_eq!(&bytes[8..], &[2u8; 8]); + + Ok(()) + } + + /// A store owned by this library must be handed back unchanged rather than + /// wrapped, so that reads never cross the boundary. + #[test] + fn local_store_is_unwrapped() { + let original = Arc::new(InMemory::new()) as Arc; + let ffi = FFI_ObjectStore::new(Arc::clone(&original), None); + + let recovered = ffi.as_local().expect("local store should unwrap"); + assert!(Arc::ptr_eq(&original, &recovered)); + } + + /// A store that has crossed into another library and is then sent back + /// arrives as the original handle rather than a second wrapper, keeping + /// reads on the local fast path. + /// + /// The check is on store identity rather than `library_marker_id`, because + /// `clone` regenerates the marker from whichever library runs it. That is + /// correct across a real boundary, where the clone executes in the owning + /// library, but means a mocked marker does not survive a clone. + #[test] + fn foreign_store_round_trips_back_to_original() { + let original = Arc::new(InMemory::new()) as Arc; + + // Cross into a "foreign" library. + let mut ffi = FFI_ObjectStore::new(Arc::clone(&original), None); + ffi.library_marker_id = crate::mock_foreign_marker_id; + let wrapped = Arc::::from(ffi); + assert!( + !Arc::ptr_eq(&wrapped, &original), + "crossing the boundary should produce a wrapper" + ); + + // Send it back the other way. + let returned = FFI_ObjectStore::new(Arc::clone(&wrapped), None); + let recovered = returned + .as_local() + .expect("a store sent home should be recognised as local"); + + assert!( + Arc::ptr_eq(&recovered, &original), + "sending a store home should recover the original store, not wrap the wrapper" + ); + } + + /// The re-export must not survive the wrapper: once the `ForeignObjectStore` + /// is dropped its side-table entry has to go, otherwise a later store + /// allocated at the same address would be mistaken for it. + #[test] + fn dropped_wrapper_is_forgotten() { + let original = Arc::new(InMemory::new()) as Arc; + let mut ffi = FFI_ObjectStore::new(Arc::clone(&original), None); + ffi.library_marker_id = crate::mock_foreign_marker_id; + + let wrapped = Arc::::from(ffi); + let key = Arc::as_ptr(&wrapped) as *const () as usize; + drop(wrapped); + + let unrelated = Arc::new(InMemory::new()) as Arc; + let handle = FFI_ObjectStore::new(Arc::clone(&unrelated), None); + let recovered = handle.as_local().expect("local store should unwrap"); + assert!( + Arc::ptr_eq(&recovered, &unrelated), + "a stale side-table entry at address {key:#x} leaked into a new store" + ); + } +} diff --git a/datafusion/ffi/src/execution/object_store/multipart.rs b/datafusion/ffi/src/execution/object_store/multipart.rs new file mode 100644 index 0000000000000..d8a9e5e5a3029 --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/multipart.rs @@ -0,0 +1,276 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`MultipartUpload`]. + +use std::ffi::c_void; + +use async_ffi::{FfiFuture, FutureExt}; +use async_trait::async_trait; +use object_store::{MultipartUpload, PutResult, UploadPart}; +use stabby::vec::Vec as SVec; +use tokio::runtime::Handle; + +use super::buffer::FFI_Bytes; +use super::error::FFI_ObjectStoreResult; +use super::types::{FFI_PutResult, put_payload_from_ffi, put_payload_to_ffi}; + +/// A stable struct for sharing [`MultipartUpload`] across FFI boundaries. +/// +/// # Safety +/// +/// [`MultipartUpload::complete`] and [`MultipartUpload::abort`] take +/// `&mut self`, so the futures they return borrow the upload. The function +/// pointers below erase that lifetime. Callers must keep this struct alive +/// until the returned future has been driven to completion or dropped. +/// [`ForeignMultipartUpload`] upholds this because the `async_trait` methods +/// borrow `self` for the duration of the returned future. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_MultipartUpload { + /// Upload the next part. The returned future does not borrow this struct. + pub put_part: unsafe extern "C" fn( + upload: &mut Self, + payload: SVec, + ) -> FfiFuture>, + + /// Complete the upload. The returned future borrows this struct. + pub complete: unsafe extern "C" fn( + upload: &mut Self, + ) + -> FfiFuture>, + + /// Abort the upload. The returned future borrows this struct. + pub abort: + unsafe extern "C" fn(upload: &mut Self) -> FfiFuture>, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// upload. The foreign library should never attempt to access this data. + pub private_data: *mut c_void, +} + +// Safety: the inner `Box` is `Send` and is only reached +// through `&mut Self`, so access is exclusive. +unsafe impl Send for FFI_MultipartUpload {} + +struct MultipartUploadPrivateData { + upload: Box, + runtime: Option, +} + +impl FFI_MultipartUpload { + pub fn new(upload: Box, runtime: Option) -> Self { + Self { + put_part: put_part_fn_wrapper, + complete: complete_fn_wrapper, + abort: abort_fn_wrapper, + release: release_fn_wrapper, + private_data: Box::into_raw(Box::new(MultipartUploadPrivateData { + upload, + runtime, + })) as *mut c_void, + } + } + + unsafe fn private_data(&mut self) -> &mut MultipartUploadPrivateData { + unsafe { &mut *(self.private_data as *mut MultipartUploadPrivateData) } + } +} + +unsafe extern "C" fn put_part_fn_wrapper( + upload: &mut FFI_MultipartUpload, + payload: SVec, +) -> FfiFuture> { + unsafe { + let private_data = upload.private_data(); + let _guard = private_data.runtime.as_ref().map(|rt| rt.enter()); + // `put_part` returns a `'static` future, so nothing borrowed from + // `upload` escapes into the returned future. + let part: UploadPart = + private_data.upload.put_part(put_payload_from_ffi(payload)); + async move { FFI_ObjectStoreResult::from(part.await) }.into_ffi() + } +} + +unsafe extern "C" fn complete_fn_wrapper( + upload: &mut FFI_MultipartUpload, +) -> FfiFuture> { + unsafe { + let private_data = upload.private_data as *mut MultipartUploadPrivateData; + // See the safety note on `FFI_MultipartUpload`: the caller keeps the + // struct alive until this future resolves. + let private_data = private_data as usize; + // A tokio `EnterGuard` is not `Send` and so cannot be held across an + // await. As elsewhere in this crate, runtime context is established per + // poll by the stream wrappers rather than around a whole future. + async move { + let private_data = &mut *(private_data as *mut MultipartUploadPrivateData); + FFI_ObjectStoreResult::from( + private_data + .upload + .complete() + .await + .map(|r| FFI_PutResult::from(&r)), + ) + } + .into_ffi() + } +} + +unsafe extern "C" fn abort_fn_wrapper( + upload: &mut FFI_MultipartUpload, +) -> FfiFuture> { + unsafe { + let private_data = upload.private_data as usize; + async move { + let private_data = &mut *(private_data as *mut MultipartUploadPrivateData); + FFI_ObjectStoreResult::from(private_data.upload.abort().await) + } + .into_ffi() + } +} + +unsafe extern "C" fn release_fn_wrapper(upload: &mut FFI_MultipartUpload) { + unsafe { + debug_assert!(!upload.private_data.is_null()); + drop(Box::from_raw( + upload.private_data as *mut MultipartUploadPrivateData, + )); + upload.private_data = std::ptr::null_mut(); + } +} + +impl Drop for FFI_MultipartUpload { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +/// A [`MultipartUpload`] backed by a foreign [`FFI_MultipartUpload`]. +#[derive(Debug)] +pub struct ForeignMultipartUpload { + upload: FFI_MultipartUpload, +} + +unsafe impl Send for ForeignMultipartUpload {} + +impl From for ForeignMultipartUpload { + fn from(upload: FFI_MultipartUpload) -> Self { + Self { upload } + } +} + +#[async_trait] +impl MultipartUpload for ForeignMultipartUpload { + fn put_part(&mut self, data: object_store::PutPayload) -> UploadPart { + let future = unsafe { + (self.upload.put_part)(&mut self.upload, put_payload_to_ffi(&data)) + }; + Box::pin(async move { future.await.into() }) + } + + async fn complete(&mut self) -> object_store::Result { + let future = unsafe { (self.upload.complete)(&mut self.upload) }; + Result::::from(future.await).map(PutResult::from) + } + + async fn abort(&mut self) -> object_store::Result<()> { + let future = unsafe { (self.upload.abort)(&mut self.upload) }; + future.await.into() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use object_store::PutPayload; + + use super::*; + + #[derive(Debug)] + struct TestUpload { + parts: Arc, + aborted: Arc, + } + + #[async_trait] + impl MultipartUpload for TestUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + self.parts + .fetch_add(data.content_length(), Ordering::SeqCst); + Box::pin(async { Ok(()) }) + } + + async fn complete(&mut self) -> object_store::Result { + Ok(PutResult { + e_tag: Some("final".to_string()), + version: Some("v2".to_string()), + }) + } + + async fn abort(&mut self) -> object_store::Result<()> { + self.aborted.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + #[tokio::test] + async fn multipart_upload_round_trip() -> datafusion_common::Result<()> { + let parts = Arc::new(AtomicUsize::new(0)); + let aborted = Arc::new(AtomicUsize::new(0)); + let upload = TestUpload { + parts: Arc::clone(&parts), + aborted: Arc::clone(&aborted), + }; + + let ffi = FFI_MultipartUpload::new(Box::new(upload), None); + let mut foreign = ForeignMultipartUpload::from(ffi); + + foreign.put_part(PutPayload::from(vec![0u8; 16])).await?; + foreign.put_part(PutPayload::from(vec![0u8; 32])).await?; + assert_eq!(parts.load(Ordering::SeqCst), 48); + + let result = foreign.complete().await?; + assert_eq!(result.e_tag.as_deref(), Some("final")); + assert_eq!(result.version.as_deref(), Some("v2")); + + Ok(()) + } + + #[tokio::test] + async fn multipart_abort_round_trip() -> datafusion_common::Result<()> { + let parts = Arc::new(AtomicUsize::new(0)); + let aborted = Arc::new(AtomicUsize::new(0)); + let upload = TestUpload { + parts: Arc::clone(&parts), + aborted: Arc::clone(&aborted), + }; + + let ffi = FFI_MultipartUpload::new(Box::new(upload), None); + let mut foreign = ForeignMultipartUpload::from(ffi); + + foreign.abort().await?; + assert_eq!(aborted.load(Ordering::SeqCst), 1); + + Ok(()) + } +} diff --git a/datafusion/ffi/src/execution/object_store/registry.rs b/datafusion/ffi/src/execution/object_store/registry.rs new file mode 100644 index 0000000000000..057bab95c0315 --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/registry.rs @@ -0,0 +1,394 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`ObjectStoreRegistry`]. +//! +//! Sharing the registry is what makes object stores visible across the +//! boundary in both directions: a store registered by a table provider during +//! planning is found by the session at execution time, and a store registered +//! on the session by the host is found by the provider. + +use std::ffi::c_void; +use std::sync::Arc; + +use datafusion_common::{Result, ffi_datafusion_err}; +use datafusion_execution::object_store::ObjectStoreRegistry; +use object_store::ObjectStore; +use stabby::string::String as SString; +use tokio::runtime::Handle; +use url::Url; + +use crate::util::{FFI_Option, FFI_Result}; +use crate::{df_result, sresult, sresult_return}; + +use super::store::FFI_ObjectStore; + +/// A stable struct for sharing [`ObjectStoreRegistry`] across FFI boundaries. +/// +/// All three methods are synchronous, so unlike most of this crate no futures +/// are involved; a registry lookup never performs I/O. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_ObjectStoreRegistry { + /// Register a store, returning the store it replaced if any. + pub register_store: unsafe extern "C" fn( + registry: &Self, + url: SString, + store: FFI_ObjectStore, + ) -> FFI_Option, + + /// Deregister the store registered for `url`. + pub deregister_store: unsafe extern "C" fn( + registry: &Self, + url: SString, + ) -> FFI_Result, + + /// Get a suitable store for `url`. + pub get_store: unsafe extern "C" fn( + registry: &Self, + url: SString, + ) -> FFI_Result, + + /// Used to create a clone on the provider of the registry. This should + /// only need to be called by the receiver of the registry. + pub clone: unsafe extern "C" fn(registry: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// registry. The foreign library should never attempt to access this data. + pub private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. See [`crate::get_library_marker_id`] and + /// the crate's `README.md` for more information. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_ObjectStoreRegistry {} +unsafe impl Sync for FFI_ObjectStoreRegistry {} + +struct RegistryPrivateData { + registry: Arc, + runtime: Option, +} + +impl FFI_ObjectStoreRegistry { + fn private_data(&self) -> &RegistryPrivateData { + unsafe { &*(self.private_data as *const RegistryPrivateData) } + } + + fn inner(&self) -> &Arc { + &self.private_data().registry + } + + fn runtime(&self) -> &Option { + &self.private_data().runtime + } + + /// Create a new [`FFI_ObjectStoreRegistry`] from a local registry. + /// + /// `runtime` is the tokio runtime handle of the providing library, attached + /// to any store handed out by this registry so that stores which spawn + /// tasks work when driven by a foreign executor. + pub fn new(registry: Arc, runtime: Option) -> Self { + Self { + register_store: register_store_fn_wrapper, + deregister_store: deregister_store_fn_wrapper, + get_store: get_store_fn_wrapper, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + private_data: Box::into_raw(Box::new(RegistryPrivateData { + registry, + runtime, + })) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } + + /// If this registry originated in the current library, return the + /// underlying [`ObjectStoreRegistry`] directly. + pub fn as_local(&self) -> Option> { + ((self.library_marker_id)() == crate::get_library_marker_id()) + .then(|| Arc::clone(self.inner())) + } +} + +fn parse_url(url: &SString) -> Result { + Url::parse(url.as_str()) + .map_err(|e| ffi_datafusion_err!("Invalid object store URL '{url}': {e}")) +} + +unsafe extern "C" fn register_store_fn_wrapper( + registry: &FFI_ObjectStoreRegistry, + url: SString, + store: FFI_ObjectStore, +) -> FFI_Option { + let Ok(url) = parse_url(&url) else { + // `register_store` cannot report an error. An unparseable URL can only + // come from a caller that built one by hand, and silently dropping the + // registration would be worse than a log line. + log::warn!("Ignoring object store registration for unparseable URL '{url}'"); + return FFI_Option::None; + }; + + let runtime = registry.runtime().clone(); + let store = Arc::::from(store); + + match registry.inner().register_store(&url, store) { + Some(previous) => FFI_Option::Some(FFI_ObjectStore::new(previous, runtime)), + None => FFI_Option::None, + } +} + +unsafe extern "C" fn deregister_store_fn_wrapper( + registry: &FFI_ObjectStoreRegistry, + url: SString, +) -> FFI_Result { + let url = sresult_return!(parse_url(&url)); + let runtime = registry.runtime().clone(); + let store = sresult_return!(registry.inner().deregister_store(&url)); + + FFI_Result::Ok(FFI_ObjectStore::new(store, runtime)) +} + +unsafe extern "C" fn get_store_fn_wrapper( + registry: &FFI_ObjectStoreRegistry, + url: SString, +) -> FFI_Result { + let url = sresult_return!(parse_url(&url)); + let runtime = registry.runtime().clone(); + + sresult!( + registry + .inner() + .get_store(&url) + .map(|store| FFI_ObjectStore::new(store, runtime)) + ) +} + +unsafe extern "C" fn clone_fn_wrapper( + registry: &FFI_ObjectStoreRegistry, +) -> FFI_ObjectStoreRegistry { + let private_data = registry.private_data(); + FFI_ObjectStoreRegistry::new( + Arc::clone(&private_data.registry), + private_data.runtime.clone(), + ) +} + +unsafe extern "C" fn release_fn_wrapper(registry: &mut FFI_ObjectStoreRegistry) { + unsafe { + debug_assert!(!registry.private_data.is_null()); + drop(Box::from_raw( + registry.private_data as *mut RegistryPrivateData, + )); + registry.private_data = std::ptr::null_mut(); + } +} + +impl Clone for FFI_ObjectStoreRegistry { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +impl Drop for FFI_ObjectStoreRegistry { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +/// An [`ObjectStoreRegistry`] backed by a foreign [`FFI_ObjectStoreRegistry`]. +/// +/// Unlike [`super::ForeignObjectStore`], no attempt is made to unwrap a +/// registry that is being sent back toward its owning library. A registry +/// crosses the boundary once per session and its methods perform no I/O, so an +/// extra layer of indirection costs nothing measurable; the stores it returns +/// are still unwrapped to their local form by +/// [`FFI_ObjectStore::as_local`]. +#[derive(Debug)] +pub struct ForeignObjectStoreRegistry { + registry: FFI_ObjectStoreRegistry, +} + +unsafe impl Send for ForeignObjectStoreRegistry {} +unsafe impl Sync for ForeignObjectStoreRegistry {} + +impl From for ForeignObjectStoreRegistry { + fn from(registry: FFI_ObjectStoreRegistry) -> Self { + Self { registry } + } +} + +impl From for Arc { + fn from(registry: FFI_ObjectStoreRegistry) -> Self { + match registry.as_local() { + Some(local) => local, + None => Arc::new(ForeignObjectStoreRegistry::from(registry)), + } + } +} + +impl ObjectStoreRegistry for ForeignObjectStoreRegistry { + fn register_store( + &self, + url: &Url, + store: Arc, + ) -> Option> { + // Attach the runtime this registration is happening on. A store that + // spawns tasks or uses timers needs it once the other side starts + // driving it from its own executor. + let runtime = Handle::try_current().ok(); + let previous = unsafe { + (self.registry.register_store)( + &self.registry, + url.as_str().into(), + FFI_ObjectStore::new(store, runtime), + ) + }; + + previous.into_option().map(Arc::::from) + } + + fn deregister_store(&self, url: &Url) -> Result> { + let store = unsafe { + (self.registry.deregister_store)(&self.registry, url.as_str().into()) + }; + df_result!(store).map(Arc::::from) + } + + fn get_store(&self, url: &Url) -> Result> { + let store = + unsafe { (self.registry.get_store)(&self.registry, url.as_str().into()) }; + df_result!(store).map(Arc::::from) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion_execution::object_store::DefaultObjectStoreRegistry; + use object_store::memory::InMemory; + + use super::*; + + fn foreign_registry( + registry: Arc, + ) -> ForeignObjectStoreRegistry { + let mut ffi = FFI_ObjectStoreRegistry::new(registry, None); + // Force the foreign path so the test exercises the FFI functions rather + // than the local unwrap. + ffi.library_marker_id = crate::mock_foreign_marker_id; + ForeignObjectStoreRegistry::from(ffi) + } + + #[test] + fn register_then_get_round_trip() -> Result<()> { + let local = Arc::new(DefaultObjectStoreRegistry::new()); + let foreign = + foreign_registry(Arc::clone(&local) as Arc); + + let url = Url::parse("s3://bucket").unwrap(); + let store = Arc::new(InMemory::new()) as Arc; + assert!(foreign.register_store(&url, store).is_none()); + + // Visible through the foreign handle... + assert!(foreign.get_store(&url).is_ok()); + // ...and in the underlying local registry. + assert!(local.get_store(&url).is_ok()); + + Ok(()) + } + + /// A store registered on the local registry must be reachable through the + /// foreign handle, so a host that registers a store on the session before + /// handing it across the boundary is honoured on the far side. + #[test] + fn store_registered_locally_is_visible_to_foreign() -> Result<()> { + let local = Arc::new(DefaultObjectStoreRegistry::new()); + let url = Url::parse("s3://host-registered").unwrap(); + local.register_store(&url, Arc::new(InMemory::new())); + + let foreign = + foreign_registry(Arc::clone(&local) as Arc); + assert!(foreign.get_store(&url).is_ok()); + + Ok(()) + } + + #[test] + fn missing_store_reports_error() { + let local = Arc::new(DefaultObjectStoreRegistry::new()); + let foreign = foreign_registry(local as Arc); + + let url = Url::parse("s3://never-registered").unwrap(); + assert!(foreign.get_store(&url).is_err()); + } + + #[test] + fn register_returns_replaced_store() -> Result<()> { + let local = Arc::new(DefaultObjectStoreRegistry::new()); + let foreign = foreign_registry(local as Arc); + + let url = Url::parse("s3://bucket").unwrap(); + assert!( + foreign + .register_store(&url, Arc::new(InMemory::new())) + .is_none() + ); + assert!( + foreign + .register_store(&url, Arc::new(InMemory::new())) + .is_some() + ); + + Ok(()) + } + + #[test] + fn deregister_round_trip() -> Result<()> { + let local = Arc::new(DefaultObjectStoreRegistry::new()); + let foreign = foreign_registry(local as Arc); + + let url = Url::parse("s3://bucket").unwrap(); + foreign.register_store(&url, Arc::new(InMemory::new())); + + assert!(foreign.deregister_store(&url).is_ok()); + assert!(foreign.get_store(&url).is_err()); + + Ok(()) + } + + #[test] + fn local_registry_is_unwrapped() { + let local = Arc::new(DefaultObjectStoreRegistry::new()); + let ffi = FFI_ObjectStoreRegistry::new( + Arc::clone(&local) as Arc, + None, + ); + assert!(ffi.as_local().is_some()); + + let registry = Arc::::from(ffi); + // The same registry instance, not a wrapper. + let url = Url::parse("s3://bucket").unwrap(); + registry.register_store(&url, Arc::new(InMemory::new())); + assert!(local.get_store(&url).is_ok()); + } +} diff --git a/datafusion/ffi/src/execution/object_store/store.rs b/datafusion/ffi/src/execution/object_store/store.rs new file mode 100644 index 0000000000000..9ea2433d23149 --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/store.rs @@ -0,0 +1,734 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`ObjectStore`]. +//! +//! # Which methods cross the boundary +//! +//! Every required method of [`ObjectStore`] is forwarded, plus the optional +//! [`ObjectStore::get_ranges`] and [`ObjectStore::list_with_offset`] overrides. +//! `get_ranges` matters because the Parquet reader issues one call per row +//! group to fetch column chunks; leaving it to the default implementation would +//! turn a single batched request into one `get_opts` round trip per range. +//! +//! The `ObjectStoreExt` methods (`head`, `copy_if_not_exists`, +//! `rename_if_not_exists`) are deliberately *not* forwarded. They are extension +//! methods defined in terms of the core trait, so they compose correctly on top +//! of the forwarded methods: `copy_if_not_exists` becomes `copy_opts` with +//! [`object_store::CopyMode::Create`], which crosses the boundary intact and so +//! keeps its atomicity guarantee. + +use std::ffi::c_void; +use std::ops::Range; +use std::sync::Arc; + +use async_ffi::{FfiFuture, FutureExt}; +use async_trait::async_trait; +use bytes::Bytes; +use futures::stream::BoxStream; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + RenameOptions, +}; +use stabby::string::String as SString; +use stabby::vec::Vec as SVec; +use tokio::runtime::Handle; + +use crate::util::FFI_Option; + +use super::buffer::FFI_Bytes; +use super::error::FFI_ObjectStoreResult; +use super::multipart::{FFI_MultipartUpload, ForeignMultipartUpload}; +use super::stream::{FFI_BytesStream, FFI_ObjectMetaStream, FFI_PathStream}; +use super::types::{ + FFI_Attribute, FFI_CopyOptions, FFI_GetOptions, FFI_ListResult, FFI_ObjectMeta, + FFI_PutMultipartOptions, FFI_PutOptions, FFI_PutResult, FFI_RenameOptions, + attributes_from_ffi, attributes_to_ffi, put_payload_from_ffi, put_payload_to_ffi, +}; + +/// An FFI-safe byte range, used by [`ObjectStore::get_ranges`]. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct FFI_ByteRange { + pub start: u64, + pub end: u64, +} + +/// An FFI-safe [`GetResult`]. +/// +/// [`GetResultPayload`] has a `File` variant that lets a local filesystem store +/// hand back an open file descriptor. A raw file descriptor is not portable +/// across the boundary, so the payload is always converted to a byte stream +/// with [`GetResult::into_stream`]. A local filesystem store reached through +/// FFI therefore loses the file fast path, which is why +/// [`FFI_ObjectStore::as_local`] exists: a store used within its own library +/// never goes through this struct at all. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_GetResult { + pub payload: FFI_BytesStream, + pub meta: FFI_ObjectMeta, + pub range_start: u64, + pub range_end: u64, + pub attributes: SVec, +} + +impl FFI_GetResult { + fn new(result: GetResult, runtime: Option) -> Self { + let meta = FFI_ObjectMeta::from(&result.meta); + let range_start = result.range.start; + let range_end = result.range.end; + let attributes = attributes_to_ffi(&result.attributes); + + Self { + payload: FFI_BytesStream::new(result.into_stream(), runtime), + meta, + range_start, + range_end, + attributes, + } + } +} + +impl From for GetResult { + fn from(result: FFI_GetResult) -> Self { + GetResult { + meta: ObjectMeta::from(result.meta), + range: result.range_start..result.range_end, + attributes: attributes_from_ffi(result.attributes), + payload: GetResultPayload::Stream(Box::pin(result.payload)), + } + } +} + +/// A stable struct for sharing [`ObjectStore`] across FFI boundaries. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_ObjectStore { + pub put_opts: unsafe extern "C" fn( + store: &Self, + location: SString, + payload: SVec, + options: FFI_PutOptions, + ) + -> FfiFuture>, + + pub put_multipart_opts: unsafe extern "C" fn( + store: &Self, + location: SString, + options: FFI_PutMultipartOptions, + ) -> FfiFuture< + FFI_ObjectStoreResult, + >, + + pub get_opts: unsafe extern "C" fn( + store: &Self, + location: SString, + options: FFI_GetOptions, + ) + -> FfiFuture>, + + pub get_ranges: unsafe extern "C" fn( + store: &Self, + location: SString, + ranges: SVec, + ) -> FfiFuture< + FFI_ObjectStoreResult>, + >, + + /// Delete a stream of locations, yielding those successfully deleted. + /// + /// This is the required trait method; the simpler `delete` is an + /// `ObjectStoreExt` extension defined in terms of it. + pub delete_stream: + unsafe extern "C" fn(store: &Self, locations: FFI_PathStream) -> FFI_PathStream, + + pub list: unsafe extern "C" fn( + store: &Self, + prefix: FFI_Option, + ) -> FFI_ObjectMetaStream, + + pub list_with_offset: unsafe extern "C" fn( + store: &Self, + prefix: FFI_Option, + offset: SString, + ) -> FFI_ObjectMetaStream, + + pub list_with_delimiter: + unsafe extern "C" fn( + store: &Self, + prefix: FFI_Option, + ) -> FfiFuture>, + + pub copy_opts: unsafe extern "C" fn( + store: &Self, + from: SString, + to: SString, + options: FFI_CopyOptions, + ) -> FfiFuture>, + + pub rename_opts: unsafe extern "C" fn( + store: &Self, + from: SString, + to: SString, + options: FFI_RenameOptions, + ) -> FfiFuture>, + + /// Backs the [`std::fmt::Display`] bound on [`ObjectStore`]. + pub display: unsafe extern "C" fn(store: &Self) -> SString, + + /// Used to create a clone on the provider of the store. This should + /// only need to be called by the receiver of the store. + pub clone: unsafe extern "C" fn(store: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the store. + /// The foreign library should never attempt to access this data. + pub private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. See [`crate::get_library_marker_id`] and + /// the crate's `README.md` for more information. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_ObjectStore {} +unsafe impl Sync for FFI_ObjectStore {} + +struct ObjectStorePrivateData { + store: Arc, + runtime: Option, +} + +impl FFI_ObjectStore { + fn private_data(&self) -> &ObjectStorePrivateData { + unsafe { &*(self.private_data as *const ObjectStorePrivateData) } + } + + fn inner(&self) -> &Arc { + &self.private_data().store + } + + fn runtime(&self) -> &Option { + &self.private_data().runtime + } + + /// Create a new [`FFI_ObjectStore`] from a local store. + /// + /// `runtime` is the tokio runtime handle of the providing library. Stores + /// that spawn tasks or use timers require it when the future is driven by a + /// foreign executor. + pub fn new(store: Arc, runtime: Option) -> Self { + // Re-export rather than double wrap when this store is itself a foreign + // store being handed back toward its owning library. See + // `foreign_store_handle` for why this is not a downcast. + if let Some(handle) = foreign_store_handle(&store) { + return handle; + } + + Self::new_unchecked(store, runtime) + } + + /// Wrap `store` without checking whether it is a foreign store being sent + /// home. Used by `clone`, where the handle is already known to be the right + /// one and re-running the check would take the `FOREIGN_STORES` lock a + /// second time on the same thread. + fn new_unchecked(store: Arc, runtime: Option) -> Self { + Self { + put_opts: put_opts_fn_wrapper, + put_multipart_opts: put_multipart_opts_fn_wrapper, + get_opts: get_opts_fn_wrapper, + get_ranges: get_ranges_fn_wrapper, + delete_stream: delete_stream_fn_wrapper, + list: list_fn_wrapper, + list_with_offset: list_with_offset_fn_wrapper, + list_with_delimiter: list_with_delimiter_fn_wrapper, + copy_opts: copy_opts_fn_wrapper, + rename_opts: rename_opts_fn_wrapper, + display: display_fn_wrapper, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + private_data: Box::into_raw(Box::new(ObjectStorePrivateData { + store, + runtime, + })) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } + + /// If this store originated in the current library, return the underlying + /// [`ObjectStore`] directly. + /// + /// This is the path that matters for a table provider that creates its own + /// store, registers it with the session, and then reads through it during + /// execution: the store makes a round trip through the registry but is + /// unwrapped back to the original `Arc` at execution time, so no data + /// actually crosses the FFI boundary. + pub fn as_local(&self) -> Option> { + ((self.library_marker_id)() == crate::get_library_marker_id()) + .then(|| Arc::clone(self.inner())) + } +} + +fn path_from(location: &SString) -> Path { + Path::from(location.to_string()) +} + +fn prefix_from(prefix: FFI_Option) -> Option { + prefix.into_option().map(|p| path_from(&p)) +} + +unsafe extern "C" fn put_opts_fn_wrapper( + store: &FFI_ObjectStore, + location: SString, + payload: SVec, + options: FFI_PutOptions, +) -> FfiFuture> { + let inner = Arc::clone(store.inner()); + async move { + let result = inner + .put_opts( + &path_from(&location), + put_payload_from_ffi(payload), + options.into(), + ) + .await; + FFI_ObjectStoreResult::from(result.map(|r| FFI_PutResult::from(&r))) + } + .into_ffi() +} + +unsafe extern "C" fn put_multipart_opts_fn_wrapper( + store: &FFI_ObjectStore, + location: SString, + options: FFI_PutMultipartOptions, +) -> FfiFuture> { + let inner = Arc::clone(store.inner()); + let runtime = store.runtime().clone(); + async move { + let result = inner + .put_multipart_opts(&path_from(&location), options.into()) + .await; + FFI_ObjectStoreResult::from( + result.map(|upload| FFI_MultipartUpload::new(upload, runtime)), + ) + } + .into_ffi() +} + +unsafe extern "C" fn get_opts_fn_wrapper( + store: &FFI_ObjectStore, + location: SString, + options: FFI_GetOptions, +) -> FfiFuture> { + let inner = Arc::clone(store.inner()); + let runtime = store.runtime().clone(); + async move { + let result = inner.get_opts(&path_from(&location), options.into()).await; + FFI_ObjectStoreResult::from( + result.map(|result| FFI_GetResult::new(result, runtime)), + ) + } + .into_ffi() +} + +unsafe extern "C" fn get_ranges_fn_wrapper( + store: &FFI_ObjectStore, + location: SString, + ranges: SVec, +) -> FfiFuture>> { + let inner = Arc::clone(store.inner()); + async move { + let ranges: Vec> = + ranges.into_iter().map(|r| r.start..r.end).collect(); + let result = inner.get_ranges(&path_from(&location), &ranges).await; + FFI_ObjectStoreResult::from( + result.map(|chunks| chunks.into_iter().map(FFI_Bytes::from).collect()), + ) + } + .into_ffi() +} + +unsafe extern "C" fn delete_stream_fn_wrapper( + store: &FFI_ObjectStore, + locations: FFI_PathStream, +) -> FFI_PathStream { + let deleted = store.inner().delete_stream(Box::pin(locations)); + FFI_PathStream::new(deleted, store.runtime().clone()) +} + +unsafe extern "C" fn list_fn_wrapper( + store: &FFI_ObjectStore, + prefix: FFI_Option, +) -> FFI_ObjectMetaStream { + let stream = store.inner().list(prefix_from(prefix).as_ref()); + FFI_ObjectMetaStream::new(stream, store.runtime().clone()) +} + +unsafe extern "C" fn list_with_offset_fn_wrapper( + store: &FFI_ObjectStore, + prefix: FFI_Option, + offset: SString, +) -> FFI_ObjectMetaStream { + let stream = store + .inner() + .list_with_offset(prefix_from(prefix).as_ref(), &path_from(&offset)); + FFI_ObjectMetaStream::new(stream, store.runtime().clone()) +} + +unsafe extern "C" fn list_with_delimiter_fn_wrapper( + store: &FFI_ObjectStore, + prefix: FFI_Option, +) -> FfiFuture> { + let inner = Arc::clone(store.inner()); + async move { + let result = inner + .list_with_delimiter(prefix_from(prefix).as_ref()) + .await; + FFI_ObjectStoreResult::from(result.map(|r| FFI_ListResult::from(&r))) + } + .into_ffi() +} + +unsafe extern "C" fn copy_opts_fn_wrapper( + store: &FFI_ObjectStore, + from: SString, + to: SString, + options: FFI_CopyOptions, +) -> FfiFuture> { + let inner = Arc::clone(store.inner()); + async move { + FFI_ObjectStoreResult::from( + inner + .copy_opts(&path_from(&from), &path_from(&to), options.into()) + .await, + ) + } + .into_ffi() +} + +unsafe extern "C" fn rename_opts_fn_wrapper( + store: &FFI_ObjectStore, + from: SString, + to: SString, + options: FFI_RenameOptions, +) -> FfiFuture> { + let inner = Arc::clone(store.inner()); + async move { + FFI_ObjectStoreResult::from( + inner + .rename_opts(&path_from(&from), &path_from(&to), options.into()) + .await, + ) + } + .into_ffi() +} + +unsafe extern "C" fn display_fn_wrapper(store: &FFI_ObjectStore) -> SString { + store.inner().to_string().as_str().into() +} + +unsafe extern "C" fn clone_fn_wrapper(store: &FFI_ObjectStore) -> FFI_ObjectStore { + let private_data = store.private_data(); + FFI_ObjectStore::new_unchecked( + Arc::clone(&private_data.store), + private_data.runtime.clone(), + ) +} + +unsafe extern "C" fn release_fn_wrapper(store: &mut FFI_ObjectStore) { + unsafe { + debug_assert!(!store.private_data.is_null()); + drop(Box::from_raw( + store.private_data as *mut ObjectStorePrivateData, + )); + store.private_data = std::ptr::null_mut(); + } +} + +impl Clone for FFI_ObjectStore { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +impl Drop for FFI_ObjectStore { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +/// Tracks the [`FFI_ObjectStore`] handle behind every live +/// [`ForeignObjectStore`], keyed by the address of the `ForeignObjectStore` +/// inside its `Arc`. +/// +/// A store that crosses from library B into library A and is later handed back +/// toward B must arrive as B's original handle. Looking the handle up here lets +/// [`FFI_ObjectStore::new`] re-export it unchanged, so +/// [`FFI_ObjectStore::as_local`] succeeds in B and reads stay in B rather than +/// making a B -> A -> B round trip through the wrapper. +/// +/// A downcast would express this more directly, but [`ObjectStore`] has no +/// `Any` supertrait and no `as_any` method, so a `dyn ObjectStore` cannot be +/// tested for a concrete type. +/// +/// Two invariants keep this sound: +/// +/// * Entries are removed in `Drop for ForeignObjectStore`, which runs before +/// the allocation is freed, so one address never maps to two live entries. +/// * Handles are held behind an `Arc` so a lookup can clone the `Arc` under the +/// lock and run [`FFI_ObjectStore::clone`] after releasing it. That clone +/// calls back into the owning library and re-enters +/// [`FFI_ObjectStore::new`]; running it under the lock would deadlock when +/// the owning library is this one, since [`std::sync::Mutex`] is not +/// reentrant. +static FOREIGN_STORES: std::sync::OnceLock< + std::sync::Mutex>>, +> = std::sync::OnceLock::new(); + +fn foreign_stores() +-> &'static std::sync::Mutex>> { + FOREIGN_STORES.get_or_init(Default::default) +} + +/// If `store` is a [`ForeignObjectStore`], return a clone of the original +/// [`FFI_ObjectStore`] handle it wraps. +fn foreign_store_handle(store: &Arc) -> Option { + let key = Arc::as_ptr(store) as *const () as usize; + + let handle = { + let stores = foreign_stores().lock().ok()?; + Arc::clone(stores.get(&key)?) + // Lock released here, before the cross-library clone below. + }; + + Some(handle.as_ref().clone()) +} + +/// An [`ObjectStore`] backed by a foreign [`FFI_ObjectStore`]. +#[derive(Debug)] +pub struct ForeignObjectStore { + pub(crate) store: FFI_ObjectStore, +} + +unsafe impl Send for ForeignObjectStore {} +unsafe impl Sync for ForeignObjectStore {} + +impl From for ForeignObjectStore { + fn from(store: FFI_ObjectStore) -> Self { + Self { store } + } +} + +impl Drop for ForeignObjectStore { + fn drop(&mut self) { + let key = std::ptr::from_ref::(self) as *const () as usize; + let removed = foreign_stores() + .lock() + .ok() + .and_then(|mut stores| stores.remove(&key)); + // Drop the handle after the lock is released: releasing it calls back + // into the owning library, which must not happen under our lock. + drop(removed); + } +} + +impl std::fmt::Display for ForeignObjectStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let name = unsafe { (self.store.display)(&self.store) }; + write!(f, "{}", name.as_str()) + } +} + +/// Convert an [`FFI_ObjectStore`] into an [`ObjectStore`], unwrapping it when +/// it originated in this library. +impl From for Arc { + fn from(store: FFI_ObjectStore) -> Self { + if let Some(local) = store.as_local() { + return local; + } + + // Keep the original handle recoverable so that sending this store back + // toward its owning library re-exports it rather than adding a second + // wrapper. See `FOREIGN_STORES`. + let handle = Arc::new(store.clone()); + let wrapper = Arc::new(ForeignObjectStore::from(store)); + let key = Arc::as_ptr(&wrapper) as *const () as usize; + if let Ok(mut stores) = foreign_stores().lock() { + stores.insert(key, handle); + } + + wrapper + } +} + +fn path_to(path: &Path) -> SString { + path.as_ref().into() +} + +fn prefix_to(prefix: Option<&Path>) -> FFI_Option { + match prefix { + Some(p) => FFI_Option::Some(path_to(p)), + None => FFI_Option::None, + } +} + +#[async_trait] +impl ObjectStore for ForeignObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + let future = unsafe { + (self.store.put_opts)( + &self.store, + path_to(location), + put_payload_to_ffi(&payload), + FFI_PutOptions::from(&opts), + ) + }; + Result::::from(future.await).map(PutResult::from) + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> object_store::Result> { + let future = unsafe { + (self.store.put_multipart_opts)( + &self.store, + path_to(location), + FFI_PutMultipartOptions::from(&opts), + ) + }; + let upload = Result::::from(future.await)?; + Ok(Box::new(ForeignMultipartUpload::from(upload))) + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + let future = unsafe { + (self.store.get_opts)( + &self.store, + path_to(location), + FFI_GetOptions::from(&options), + ) + }; + Result::::from(future.await).map(GetResult::from) + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> object_store::Result> { + let ffi_ranges: SVec = ranges + .iter() + .map(|r| FFI_ByteRange { + start: r.start, + end: r.end, + }) + .collect(); + let future = unsafe { + (self.store.get_ranges)(&self.store, path_to(location), ffi_ranges) + }; + let chunks = Result::, _>::from(future.await)?; + Ok(chunks.into_iter().map(Bytes::from).collect()) + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + let deleted = unsafe { + (self.store.delete_stream)(&self.store, FFI_PathStream::new(locations, None)) + }; + Box::pin(deleted) + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'static, object_store::Result> { + let stream = unsafe { (self.store.list)(&self.store, prefix_to(prefix)) }; + Box::pin(stream) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, object_store::Result> { + let stream = unsafe { + (self.store.list_with_offset)(&self.store, prefix_to(prefix), path_to(offset)) + }; + Box::pin(stream) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> object_store::Result { + let future = + unsafe { (self.store.list_with_delimiter)(&self.store, prefix_to(prefix)) }; + Result::::from(future.await).map(ListResult::from) + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> object_store::Result<()> { + let future = unsafe { + (self.store.copy_opts)( + &self.store, + path_to(from), + path_to(to), + FFI_CopyOptions::from(&options), + ) + }; + future.await.into() + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + options: RenameOptions, + ) -> object_store::Result<()> { + let future = unsafe { + (self.store.rename_opts)( + &self.store, + path_to(from), + path_to(to), + FFI_RenameOptions::from(&options), + ) + }; + future.await.into() + } +} diff --git a/datafusion/ffi/src/execution/object_store/stream.rs b/datafusion/ffi/src/execution/object_store/stream.rs new file mode 100644 index 0000000000000..c6eb452de4bae --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/stream.rs @@ -0,0 +1,490 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI-safe wrappers for the streams returned by +//! [`object_store::ObjectStore`]. +//! +//! Three streams cross the boundary: the byte stream behind +//! [`object_store::GetResultPayload::Stream`], the [`ObjectMeta`] stream +//! returned by [`object_store::ObjectStore::list`], and the [`Path`] stream +//! used by [`object_store::ObjectStore::delete_stream`] in both directions. +//! All follow the same `poll_next` pattern used by +//! [`crate::record_batch_stream`], entering the producing library's tokio +//! runtime for the duration of the poll so that stores which spawn tasks or use +//! timers work when driven by a foreign executor. + +use std::ffi::c_void; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use async_ffi::{ContextExt, FfiContext, FfiPoll}; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use object_store::path::Path; +use object_store::{Error as ObjectStoreError, ObjectMeta}; +use stabby::string::String as SString; +use tokio::runtime::Handle; + +use crate::util::FFI_Option; + +use super::buffer::FFI_Bytes; +use super::error::FFI_ObjectStoreResult; +use super::types::FFI_ObjectMeta; + +type BytesStream = futures::stream::BoxStream<'static, object_store::Result>; +type ObjectMetaStream = + futures::stream::BoxStream<'static, object_store::Result>; +type PathStream = futures::stream::BoxStream<'static, object_store::Result>; + +/// Panic message used when a foreign `poll_next` unwinds. +const POLL_PANIC: &str = "Panic occurred while polling a foreign object store stream"; + +// ----------------------------------------------------------------------------- +// Byte stream +// ----------------------------------------------------------------------------- + +/// A stable struct for sharing a stream of [`Bytes`] across FFI boundaries. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_BytesStream { + /// Mirrors [`Stream::poll_next`] in an FFI safe manner. + pub poll_next: unsafe extern "C" fn( + stream: &Self, + cx: &mut FfiContext, + ) -> FfiPoll< + FFI_Option>, + >, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// stream. The foreign library should never attempt to access this data. + pub private_data: *mut c_void, +} + +// Safety: the inner stream is a `BoxStream` which is `Send`, and access is +// serialized by the caller polling it. +unsafe impl Send for FFI_BytesStream {} + +struct BytesStreamPrivateData { + stream: BytesStream, + runtime: Option, +} + +unsafe extern "C" fn bytes_poll_next_fn_wrapper( + stream: &FFI_BytesStream, + cx: &mut FfiContext, +) -> FfiPoll>> { + unsafe { + let private_data = stream.private_data as *mut BytesStreamPrivateData; + let _guard = (*private_data).runtime.as_ref().map(|rt| rt.enter()); + let stream = &mut (*private_data).stream; + + cx.with_context(|std_cx| { + stream.poll_next_unpin(std_cx).map(|item| match item { + Some(Ok(bytes)) => { + FFI_Option::Some(FFI_ObjectStoreResult::Ok(FFI_Bytes::from(bytes))) + } + Some(Err(e)) => FFI_Option::Some(FFI_ObjectStoreResult::Err(e.into())), + None => FFI_Option::None, + }) + }) + .into() + } +} + +unsafe extern "C" fn bytes_release_fn_wrapper(stream: &mut FFI_BytesStream) { + unsafe { + debug_assert!(!stream.private_data.is_null()); + drop(Box::from_raw( + stream.private_data as *mut BytesStreamPrivateData, + )); + stream.private_data = std::ptr::null_mut(); + } +} + +impl FFI_BytesStream { + pub fn new(stream: BytesStream, runtime: Option) -> Self { + Self { + poll_next: bytes_poll_next_fn_wrapper, + release: bytes_release_fn_wrapper, + private_data: Box::into_raw(Box::new(BytesStreamPrivateData { + stream, + runtime, + })) as *mut c_void, + } + } +} + +impl Drop for FFI_BytesStream { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl Stream for FFI_BytesStream { + type Item = object_store::Result; + + // `Stream::Item` is `object_store::Result`, so the error type is fixed by + // the trait and cannot be boxed to shrink it. + #[expect(clippy::result_large_err)] + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let poll_result = + unsafe { cx.with_ffi_context(|ffi_cx| (self.poll_next)(&self, ffi_cx)) }; + + match poll_result { + FfiPoll::Ready(item) => Poll::Ready(item.into_option().map(|result| { + Result::::from(result).map(Bytes::from) + })), + FfiPoll::Pending => Poll::Pending, + FfiPoll::Panicked => Poll::Ready(Some(Err(ObjectStoreError::Generic { + store: "ForeignObjectStore", + source: POLL_PANIC.into(), + }))), + } + } +} + +// ----------------------------------------------------------------------------- +// ObjectMeta stream +// ----------------------------------------------------------------------------- + +/// A stable struct for sharing a stream of [`ObjectMeta`] across FFI +/// boundaries. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_ObjectMetaStream { + /// Mirrors [`Stream::poll_next`] in an FFI safe manner. + pub poll_next: unsafe extern "C" fn( + stream: &Self, + cx: &mut FfiContext, + ) -> FfiPoll< + FFI_Option>, + >, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// stream. The foreign library should never attempt to access this data. + pub private_data: *mut c_void, +} + +// Safety: see `FFI_BytesStream`. +unsafe impl Send for FFI_ObjectMetaStream {} + +struct ObjectMetaStreamPrivateData { + stream: ObjectMetaStream, + runtime: Option, +} + +unsafe extern "C" fn meta_poll_next_fn_wrapper( + stream: &FFI_ObjectMetaStream, + cx: &mut FfiContext, +) -> FfiPoll>> { + unsafe { + let private_data = stream.private_data as *mut ObjectMetaStreamPrivateData; + let _guard = (*private_data).runtime.as_ref().map(|rt| rt.enter()); + let stream = &mut (*private_data).stream; + + cx.with_context(|std_cx| { + stream.poll_next_unpin(std_cx).map(|item| match item { + Some(Ok(meta)) => FFI_Option::Some(FFI_ObjectStoreResult::Ok( + FFI_ObjectMeta::from(&meta), + )), + Some(Err(e)) => FFI_Option::Some(FFI_ObjectStoreResult::Err(e.into())), + None => FFI_Option::None, + }) + }) + .into() + } +} + +unsafe extern "C" fn meta_release_fn_wrapper(stream: &mut FFI_ObjectMetaStream) { + unsafe { + debug_assert!(!stream.private_data.is_null()); + drop(Box::from_raw( + stream.private_data as *mut ObjectMetaStreamPrivateData, + )); + stream.private_data = std::ptr::null_mut(); + } +} + +impl FFI_ObjectMetaStream { + pub fn new(stream: ObjectMetaStream, runtime: Option) -> Self { + Self { + poll_next: meta_poll_next_fn_wrapper, + release: meta_release_fn_wrapper, + private_data: Box::into_raw(Box::new(ObjectMetaStreamPrivateData { + stream, + runtime, + })) as *mut c_void, + } + } +} + +impl Drop for FFI_ObjectMetaStream { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl Stream for FFI_ObjectMetaStream { + type Item = object_store::Result; + + // `Stream::Item` is `object_store::Result`, so the error type is fixed by + // the trait and cannot be boxed to shrink it. + #[expect(clippy::result_large_err)] + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let poll_result = + unsafe { cx.with_ffi_context(|ffi_cx| (self.poll_next)(&self, ffi_cx)) }; + + match poll_result { + FfiPoll::Ready(item) => Poll::Ready(item.into_option().map(|result| { + Result::::from(result) + .map(ObjectMeta::from) + })), + FfiPoll::Pending => Poll::Pending, + FfiPoll::Panicked => Poll::Ready(Some(Err(ObjectStoreError::Generic { + store: "ForeignObjectStore", + source: POLL_PANIC.into(), + }))), + } + } +} + +// ----------------------------------------------------------------------------- +// Path stream +// ----------------------------------------------------------------------------- + +/// A stable struct for sharing a stream of [`Path`] across FFI boundaries. +/// +/// [`object_store::ObjectStore::delete_stream`] both consumes and produces a +/// path stream, so this struct crosses the boundary in both directions. Each +/// side wraps its own stream, so the same struct serves both roles. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_PathStream { + /// Mirrors [`Stream::poll_next`] in an FFI safe manner. + pub poll_next: unsafe extern "C" fn( + stream: &Self, + cx: &mut FfiContext, + ) -> FfiPoll< + FFI_Option>, + >, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// stream. The foreign library should never attempt to access this data. + pub private_data: *mut c_void, +} + +// Safety: see `FFI_BytesStream`. +unsafe impl Send for FFI_PathStream {} + +struct PathStreamPrivateData { + stream: PathStream, + runtime: Option, +} + +unsafe extern "C" fn path_poll_next_fn_wrapper( + stream: &FFI_PathStream, + cx: &mut FfiContext, +) -> FfiPoll>> { + unsafe { + let private_data = stream.private_data as *mut PathStreamPrivateData; + let _guard = (*private_data).runtime.as_ref().map(|rt| rt.enter()); + let stream = &mut (*private_data).stream; + + cx.with_context(|std_cx| { + stream.poll_next_unpin(std_cx).map(|item| match item { + Some(Ok(path)) => { + FFI_Option::Some(FFI_ObjectStoreResult::Ok(path.as_ref().into())) + } + Some(Err(e)) => FFI_Option::Some(FFI_ObjectStoreResult::Err(e.into())), + None => FFI_Option::None, + }) + }) + .into() + } +} + +unsafe extern "C" fn path_release_fn_wrapper(stream: &mut FFI_PathStream) { + unsafe { + debug_assert!(!stream.private_data.is_null()); + drop(Box::from_raw( + stream.private_data as *mut PathStreamPrivateData, + )); + stream.private_data = std::ptr::null_mut(); + } +} + +impl FFI_PathStream { + pub fn new(stream: PathStream, runtime: Option) -> Self { + Self { + poll_next: path_poll_next_fn_wrapper, + release: path_release_fn_wrapper, + private_data: Box::into_raw(Box::new(PathStreamPrivateData { + stream, + runtime, + })) as *mut c_void, + } + } +} + +impl Drop for FFI_PathStream { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl Stream for FFI_PathStream { + type Item = object_store::Result; + + // `Stream::Item` is `object_store::Result`, so the error type is fixed by + // the trait and cannot be boxed to shrink it. + #[expect(clippy::result_large_err)] + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let poll_result = + unsafe { cx.with_ffi_context(|ffi_cx| (self.poll_next)(&self, ffi_cx)) }; + + match poll_result { + FfiPoll::Ready(item) => Poll::Ready(item.into_option().map(|result| { + Result::::from(result) + .map(|p| Path::from(p.to_string())) + })), + FfiPoll::Pending => Poll::Pending, + FfiPoll::Panicked => Poll::Ready(Some(Err(ObjectStoreError::Generic { + store: "ForeignObjectStore", + source: POLL_PANIC.into(), + }))), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn bytes_stream_round_trip() { + let chunks = vec![ + Ok(Bytes::from_static(b"abc")), + Ok(Bytes::from_static(b"defg")), + ]; + let stream = futures::stream::iter(chunks).boxed(); + + let ffi = FFI_BytesStream::new(stream, None); + let collected: Vec<_> = ffi.collect().await; + + assert_eq!(collected.len(), 2); + assert_eq!(collected[0].as_ref().unwrap(), &Bytes::from_static(b"abc")); + assert_eq!(collected[1].as_ref().unwrap(), &Bytes::from_static(b"defg")); + } + + #[tokio::test] + async fn bytes_stream_propagates_error_variant() { + let chunks = vec![ + Ok(Bytes::from_static(b"ok")), + Err(ObjectStoreError::NotFound { + path: "gone".into(), + source: "missing".into(), + }), + ]; + let ffi = FFI_BytesStream::new(futures::stream::iter(chunks).boxed(), None); + let collected: Vec<_> = ffi.collect().await; + + assert!(collected[0].is_ok()); + assert!(matches!( + collected[1].as_ref().unwrap_err(), + ObjectStoreError::NotFound { .. } + )); + } + + #[tokio::test] + async fn object_meta_stream_round_trip() { + let metas = vec![ + Ok(ObjectMeta { + location: Path::from("a/1.parquet"), + last_modified: chrono::DateTime::from_timestamp(5, 0).unwrap(), + size: 10, + e_tag: Some("e1".to_string()), + version: None, + }), + Ok(ObjectMeta { + location: Path::from("a/2.parquet"), + last_modified: chrono::DateTime::from_timestamp(6, 0).unwrap(), + size: 20, + e_tag: None, + version: None, + }), + ]; + let expected: Vec = + metas.iter().map(|m| m.as_ref().unwrap().clone()).collect(); + + let ffi = FFI_ObjectMetaStream::new(futures::stream::iter(metas).boxed(), None); + let collected: Vec<_> = ffi.collect().await; + + assert_eq!(collected.len(), 2); + for (actual, expected) in collected.into_iter().zip(expected) { + assert_eq!(actual.unwrap(), expected); + } + } + + #[tokio::test] + async fn empty_stream_terminates() { + let ffi = FFI_BytesStream::new(futures::stream::empty().boxed(), None); + let collected: Vec<_> = ffi.collect().await; + assert!(collected.is_empty()); + } + + #[tokio::test] + async fn path_stream_round_trip() { + let paths = vec![ + Ok(Path::from("a/1.parquet")), + Ok(Path::from("b/c/2.parquet")), + ]; + let ffi = FFI_PathStream::new(futures::stream::iter(paths).boxed(), None); + let collected: Vec<_> = ffi.collect().await; + + assert_eq!(collected.len(), 2); + assert_eq!(collected[0].as_ref().unwrap(), &Path::from("a/1.parquet")); + assert_eq!(collected[1].as_ref().unwrap(), &Path::from("b/c/2.parquet")); + } + + #[tokio::test] + async fn path_stream_propagates_error_variant() { + let paths = vec![ + Ok(Path::from("ok")), + Err(ObjectStoreError::NotFound { + path: "gone".into(), + source: "missing".into(), + }), + ]; + let ffi = FFI_PathStream::new(futures::stream::iter(paths).boxed(), None); + let collected: Vec<_> = ffi.collect().await; + + assert!(collected[0].is_ok()); + assert!(matches!( + collected[1].as_ref().unwrap_err(), + ObjectStoreError::NotFound { .. } + )); + } +} diff --git a/datafusion/ffi/src/execution/object_store/types.rs b/datafusion/ffi/src/execution/object_store/types.rs new file mode 100644 index 0000000000000..0a3708617d273 --- /dev/null +++ b/datafusion/ffi/src/execution/object_store/types.rs @@ -0,0 +1,716 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI-safe representations of the plain-data types in the +//! [`object_store::ObjectStore`] API. +//! +//! # Extensions are not carried across the boundary +//! +//! [`GetOptions`], [`PutOptions`], [`PutMultipartOptions`], [`CopyOptions`], +//! and [`RenameOptions`] each carry an `extensions: ::http::Extensions` field. +//! `Extensions` is a heterogeneous `TypeId`-keyed map of arbitrary values; +//! `TypeId` is not stable across separately compiled libraries, so the contents +//! cannot be interpreted on the far side even in principle. Extensions are +//! therefore dropped in both directions. The `object_store` crate documents +//! that its own backends ignore extensions entirely, so this only affects +//! third-party stores that use them for out-of-band context such as tracing +//! spans. + +use chrono::{DateTime, Utc}; +use object_store::path::Path; +use object_store::{ + Attribute, AttributeValue, Attributes, CopyMode, CopyOptions, GetOptions, GetRange, + ListResult, ObjectMeta, PutMode, PutMultipartOptions, PutOptions, PutPayload, + PutResult, RenameOptions, RenameTargetMode, TagSet, UpdateVersion, +}; +use stabby::string::String as SString; +use stabby::vec::Vec as SVec; + +use crate::util::FFI_Option; + +/// An FFI-safe [`DateTime`]. +/// +/// Sent as a split timestamp rather than a nanosecond count so that timestamps +/// outside the roughly 1677-2262 range representable in `i64` nanoseconds +/// survive the round trip. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FFI_Timestamp { + pub secs: i64, + pub nanos: u32, +} + +impl From> for FFI_Timestamp { + fn from(value: DateTime) -> Self { + Self { + secs: value.timestamp(), + nanos: value.timestamp_subsec_nanos(), + } + } +} + +impl From for DateTime { + fn from(value: FFI_Timestamp) -> Self { + DateTime::from_timestamp(value.secs, value.nanos).unwrap_or_default() + } +} + +fn opt_string(value: &Option) -> FFI_Option { + match value { + Some(v) => FFI_Option::Some(v.as_str().into()), + None => FFI_Option::None, + } +} + +fn from_opt_string(value: FFI_Option) -> Option { + value.into_option().map(|v| v.to_string()) +} + +/// An FFI-safe [`ObjectMeta`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_ObjectMeta { + pub location: SString, + pub last_modified: FFI_Timestamp, + pub size: u64, + pub e_tag: FFI_Option, + pub version: FFI_Option, +} + +impl From<&ObjectMeta> for FFI_ObjectMeta { + fn from(meta: &ObjectMeta) -> Self { + Self { + location: meta.location.as_ref().into(), + last_modified: meta.last_modified.into(), + size: meta.size, + e_tag: opt_string(&meta.e_tag), + version: opt_string(&meta.version), + } + } +} + +impl From for FFI_ObjectMeta { + fn from(meta: ObjectMeta) -> Self { + Self::from(&meta) + } +} + +impl From for ObjectMeta { + fn from(meta: FFI_ObjectMeta) -> Self { + ObjectMeta { + // `Path::from` normalizes, and `location` was produced by + // `Path::as_ref` on an already-normalized path, so this round trips. + location: Path::from(meta.location.to_string()), + last_modified: meta.last_modified.into(), + size: meta.size, + e_tag: from_opt_string(meta.e_tag), + version: from_opt_string(meta.version), + } + } +} + +/// An FFI-safe [`GetRange`]. +#[repr(C, u8)] +#[derive(Debug, Clone, Copy)] +pub enum FFI_GetRange { + Bounded { start: u64, end: u64 }, + Offset(u64), + Suffix(u64), +} + +impl From<&GetRange> for FFI_GetRange { + fn from(range: &GetRange) -> Self { + match range { + GetRange::Bounded(r) => FFI_GetRange::Bounded { + start: r.start, + end: r.end, + }, + GetRange::Offset(o) => FFI_GetRange::Offset(*o), + GetRange::Suffix(s) => FFI_GetRange::Suffix(*s), + } + } +} + +impl From for GetRange { + fn from(range: FFI_GetRange) -> Self { + match range { + FFI_GetRange::Bounded { start, end } => GetRange::Bounded(start..end), + FFI_GetRange::Offset(o) => GetRange::Offset(o), + FFI_GetRange::Suffix(s) => GetRange::Suffix(s), + } + } +} + +/// An FFI-safe [`Attribute`]. +/// +/// [`Attribute::Metadata`] carries a user-defined key, which travels in the +/// accompanying string; the other variants ignore it. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_Attribute { + pub kind: u8, + pub metadata_key: SString, + pub value: SString, +} + +const ATTR_CONTENT_DISPOSITION: u8 = 0; +const ATTR_CONTENT_ENCODING: u8 = 1; +const ATTR_CONTENT_LANGUAGE: u8 = 2; +const ATTR_CONTENT_TYPE: u8 = 3; +const ATTR_CACHE_CONTROL: u8 = 4; +const ATTR_STORAGE_CLASS: u8 = 5; +const ATTR_METADATA: u8 = 6; + +pub(crate) fn attributes_to_ffi(attributes: &Attributes) -> SVec { + attributes + .iter() + .map(|(key, value)| { + let (kind, metadata_key) = match key { + Attribute::ContentDisposition => (ATTR_CONTENT_DISPOSITION, ""), + Attribute::ContentEncoding => (ATTR_CONTENT_ENCODING, ""), + Attribute::ContentLanguage => (ATTR_CONTENT_LANGUAGE, ""), + Attribute::ContentType => (ATTR_CONTENT_TYPE, ""), + Attribute::CacheControl => (ATTR_CACHE_CONTROL, ""), + Attribute::StorageClass => (ATTR_STORAGE_CLASS, ""), + Attribute::Metadata(k) => (ATTR_METADATA, k.as_ref()), + // `Attribute` is `#[non_exhaustive]`; an unknown variant is + // dropped rather than mistranslated. + _ => (u8::MAX, ""), + }; + FFI_Attribute { + kind, + metadata_key: metadata_key.into(), + value: value.as_ref().into(), + } + }) + .filter(|attr| attr.kind != u8::MAX) + .collect() +} + +pub(crate) fn attributes_from_ffi(attributes: SVec) -> Attributes { + attributes + .into_iter() + .filter_map(|attr| { + let key = match attr.kind { + ATTR_CONTENT_DISPOSITION => Attribute::ContentDisposition, + ATTR_CONTENT_ENCODING => Attribute::ContentEncoding, + ATTR_CONTENT_LANGUAGE => Attribute::ContentLanguage, + ATTR_CONTENT_TYPE => Attribute::ContentType, + ATTR_CACHE_CONTROL => Attribute::CacheControl, + ATTR_STORAGE_CLASS => Attribute::StorageClass, + ATTR_METADATA => { + Attribute::Metadata(attr.metadata_key.to_string().into()) + } + _ => return None, + }; + Some((key, AttributeValue::from(attr.value.to_string()))) + }) + .collect() +} + +/// Encode a [`TagSet`] as its URL-encoded wire form. +/// +/// [`TagSet`] exposes [`TagSet::encoded`] but can only be built with +/// [`TagSet::push`], so the encoded form is parsed back into pairs and re-pushed +/// on the far side. `push` applies the same encoding, so this round trips. +pub(crate) fn tags_to_ffi(tags: &TagSet) -> SString { + tags.encoded().into() +} + +pub(crate) fn tags_from_ffi(encoded: &SString) -> TagSet { + let encoded = encoded.to_string(); + let mut tags = TagSet::default(); + for (key, value) in url::form_urlencoded::parse(encoded.as_bytes()) { + tags.push(key.as_ref(), value.as_ref()); + } + tags +} + +/// An FFI-safe [`GetOptions`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_GetOptions { + pub if_match: FFI_Option, + pub if_none_match: FFI_Option, + pub if_modified_since: FFI_Option, + pub if_unmodified_since: FFI_Option, + pub range: FFI_Option, + pub version: FFI_Option, + pub head: bool, +} + +impl From<&GetOptions> for FFI_GetOptions { + fn from(options: &GetOptions) -> Self { + Self { + if_match: opt_string(&options.if_match), + if_none_match: opt_string(&options.if_none_match), + if_modified_since: options.if_modified_since.map(FFI_Timestamp::from).into(), + if_unmodified_since: options + .if_unmodified_since + .map(FFI_Timestamp::from) + .into(), + range: options.range.as_ref().map(FFI_GetRange::from).into(), + version: opt_string(&options.version), + head: options.head, + } + } +} + +impl From for GetOptions { + fn from(options: FFI_GetOptions) -> Self { + GetOptions { + if_match: from_opt_string(options.if_match), + if_none_match: from_opt_string(options.if_none_match), + if_modified_since: options.if_modified_since.into_option().map(Into::into), + if_unmodified_since: options + .if_unmodified_since + .into_option() + .map(Into::into), + range: options.range.into_option().map(Into::into), + version: from_opt_string(options.version), + head: options.head, + extensions: Default::default(), + } + } +} + +/// An FFI-safe [`PutMode`]. +#[repr(C, u8)] +#[derive(Debug, Clone)] +pub enum FFI_PutMode { + Overwrite, + Create, + Update { + e_tag: FFI_Option, + version: FFI_Option, + }, +} + +impl From<&PutMode> for FFI_PutMode { + fn from(mode: &PutMode) -> Self { + match mode { + PutMode::Overwrite => FFI_PutMode::Overwrite, + PutMode::Create => FFI_PutMode::Create, + PutMode::Update(v) => FFI_PutMode::Update { + e_tag: opt_string(&v.e_tag), + version: opt_string(&v.version), + }, + } + } +} + +impl From for PutMode { + fn from(mode: FFI_PutMode) -> Self { + match mode { + FFI_PutMode::Overwrite => PutMode::Overwrite, + FFI_PutMode::Create => PutMode::Create, + FFI_PutMode::Update { e_tag, version } => PutMode::Update(UpdateVersion { + e_tag: from_opt_string(e_tag), + version: from_opt_string(version), + }), + } + } +} + +/// An FFI-safe [`PutOptions`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_PutOptions { + pub mode: FFI_PutMode, + pub tags: SString, + pub attributes: SVec, +} + +impl From<&PutOptions> for FFI_PutOptions { + fn from(options: &PutOptions) -> Self { + Self { + mode: (&options.mode).into(), + tags: tags_to_ffi(&options.tags), + attributes: attributes_to_ffi(&options.attributes), + } + } +} + +impl From for PutOptions { + fn from(options: FFI_PutOptions) -> Self { + PutOptions { + mode: options.mode.into(), + tags: tags_from_ffi(&options.tags), + attributes: attributes_from_ffi(options.attributes), + extensions: Default::default(), + } + } +} + +/// An FFI-safe [`PutMultipartOptions`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_PutMultipartOptions { + pub tags: SString, + pub attributes: SVec, +} + +impl From<&PutMultipartOptions> for FFI_PutMultipartOptions { + fn from(options: &PutMultipartOptions) -> Self { + Self { + tags: tags_to_ffi(&options.tags), + attributes: attributes_to_ffi(&options.attributes), + } + } +} + +impl From for PutMultipartOptions { + fn from(options: FFI_PutMultipartOptions) -> Self { + PutMultipartOptions { + tags: tags_from_ffi(&options.tags), + attributes: attributes_from_ffi(options.attributes), + extensions: Default::default(), + } + } +} + +/// An FFI-safe [`PutResult`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_PutResult { + pub e_tag: FFI_Option, + pub version: FFI_Option, +} + +impl From<&PutResult> for FFI_PutResult { + fn from(result: &PutResult) -> Self { + Self { + e_tag: opt_string(&result.e_tag), + version: opt_string(&result.version), + } + } +} + +impl From for PutResult { + fn from(result: FFI_PutResult) -> Self { + PutResult { + e_tag: from_opt_string(result.e_tag), + version: from_opt_string(result.version), + } + } +} + +/// An FFI-safe [`CopyOptions`]. +/// +/// [`CopyMode::Create`] is what makes `copy_if_not_exists` atomic, so the mode +/// must survive the boundary intact. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct FFI_CopyOptions { + /// `true` for [`CopyMode::Create`], `false` for [`CopyMode::Overwrite`]. + pub create: bool, +} + +impl From<&CopyOptions> for FFI_CopyOptions { + fn from(options: &CopyOptions) -> Self { + Self { + create: matches!(options.mode, CopyMode::Create), + } + } +} + +impl From for CopyOptions { + fn from(options: FFI_CopyOptions) -> Self { + CopyOptions { + mode: if options.create { + CopyMode::Create + } else { + CopyMode::Overwrite + }, + extensions: Default::default(), + } + } +} + +/// An FFI-safe [`RenameOptions`]. +/// +/// [`RenameTargetMode::Create`] is what makes `rename_if_not_exists` atomic, +/// which delta-style commit protocols rely on for concurrency control. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct FFI_RenameOptions { + /// `true` for [`RenameTargetMode::Create`], `false` for + /// [`RenameTargetMode::Overwrite`]. + pub create: bool, +} + +impl From<&RenameOptions> for FFI_RenameOptions { + fn from(options: &RenameOptions) -> Self { + Self { + create: matches!(options.target_mode, RenameTargetMode::Create), + } + } +} + +impl From for RenameOptions { + fn from(options: FFI_RenameOptions) -> Self { + RenameOptions { + target_mode: if options.create { + RenameTargetMode::Create + } else { + RenameTargetMode::Overwrite + }, + extensions: Default::default(), + } + } +} + +/// An FFI-safe [`ListResult`]. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_ListResult { + pub common_prefixes: SVec, + pub objects: SVec, +} + +impl From<&ListResult> for FFI_ListResult { + fn from(result: &ListResult) -> Self { + Self { + common_prefixes: result + .common_prefixes + .iter() + .map(|p| p.as_ref().into()) + .collect(), + objects: result.objects.iter().map(FFI_ObjectMeta::from).collect(), + } + } +} + +impl From for ListResult { + fn from(result: FFI_ListResult) -> Self { + ListResult { + common_prefixes: result + .common_prefixes + .into_iter() + .map(|p| Path::from(p.to_string())) + .collect(), + objects: result.objects.into_iter().map(ObjectMeta::from).collect(), + } + } +} + +/// Convert a [`PutPayload`] into its FFI-safe wire form. +/// +/// The payload is a sequence of [`bytes::Bytes`] blocks; each block crosses as +/// its own buffer so the block structure is preserved and no data is copied. +pub(crate) fn put_payload_to_ffi(payload: &PutPayload) -> SVec { + payload + .iter() + .map(|b| super::FFI_Bytes::from(b.clone())) + .collect() +} + +pub(crate) fn put_payload_from_ffi(payload: SVec) -> PutPayload { + payload + .into_iter() + .map(bytes::Bytes::from) + .collect::() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn object_meta_round_trip() { + let original = ObjectMeta { + location: Path::from("a/b/c.parquet"), + last_modified: DateTime::from_timestamp(1_700_000_000, 123_456_789).unwrap(), + size: 4096, + e_tag: Some("\"abc123\"".to_string()), + version: Some("v7".to_string()), + }; + + let restored = ObjectMeta::from(FFI_ObjectMeta::from(&original)); + assert_eq!(restored, original); + } + + #[test] + fn object_meta_round_trip_without_optionals() { + let original = ObjectMeta { + location: Path::from("root.json"), + last_modified: DateTime::from_timestamp(0, 0).unwrap(), + size: 0, + e_tag: None, + version: None, + }; + + let restored = ObjectMeta::from(FFI_ObjectMeta::from(&original)); + assert_eq!(restored, original); + } + + #[test] + fn get_range_round_trip() { + for original in [ + GetRange::Bounded(10..200), + GetRange::Offset(42), + GetRange::Suffix(8), + ] { + let restored = GetRange::from(FFI_GetRange::from(&original)); + assert_eq!(format!("{restored:?}"), format!("{original:?}")); + } + } + + #[test] + fn get_options_round_trip() { + let original = GetOptions { + if_match: Some("\"e1\"".to_string()), + if_none_match: Some("\"e2\"".to_string()), + if_modified_since: Some(DateTime::from_timestamp(1_600_000_000, 0).unwrap()), + if_unmodified_since: Some( + DateTime::from_timestamp(1_600_000_001, 500).unwrap(), + ), + range: Some(GetRange::Bounded(5..25)), + version: Some("v1".to_string()), + head: true, + extensions: Default::default(), + }; + + let restored = GetOptions::from(FFI_GetOptions::from(&original)); + assert_eq!(restored.if_match, original.if_match); + assert_eq!(restored.if_none_match, original.if_none_match); + assert_eq!(restored.if_modified_since, original.if_modified_since); + assert_eq!(restored.if_unmodified_since, original.if_unmodified_since); + assert_eq!( + format!("{:?}", restored.range), + format!("{:?}", original.range) + ); + assert_eq!(restored.version, original.version); + assert_eq!(restored.head, original.head); + } + + #[test] + fn put_mode_round_trip() { + let cases = [ + PutMode::Overwrite, + PutMode::Create, + PutMode::Update(UpdateVersion { + e_tag: Some("\"tag\"".to_string()), + version: None, + }), + ]; + for original in cases { + let restored = PutMode::from(FFI_PutMode::from(&original)); + assert_eq!(restored, original); + } + } + + #[test] + fn attributes_round_trip() { + let original: Attributes = [ + (Attribute::ContentType, "application/parquet"), + (Attribute::CacheControl, "no-cache"), + (Attribute::StorageClass, "GLACIER"), + (Attribute::Metadata("custom".into()), "value"), + ] + .into_iter() + .collect(); + + let restored = attributes_from_ffi(attributes_to_ffi(&original)); + assert_eq!(restored, original); + } + + #[test] + fn tags_round_trip() { + let mut original = TagSet::default(); + original.push("test/foo", "value sdlks"); + original.push("foo", " sdf _ /+./sd"); + + let restored = tags_from_ffi(&tags_to_ffi(&original)); + assert_eq!(restored.encoded(), original.encoded()); + } + + #[test] + fn empty_tags_round_trip() { + let restored = tags_from_ffi(&tags_to_ffi(&TagSet::default())); + assert_eq!(restored.encoded(), ""); + } + + #[test] + fn copy_and_rename_options_round_trip() { + for create in [true, false] { + let original = CopyOptions { + mode: if create { + CopyMode::Create + } else { + CopyMode::Overwrite + }, + extensions: Default::default(), + }; + let restored = CopyOptions::from(FFI_CopyOptions::from(&original)); + assert_eq!(restored.mode, original.mode); + + let original = RenameOptions { + target_mode: if create { + RenameTargetMode::Create + } else { + RenameTargetMode::Overwrite + }, + extensions: Default::default(), + }; + let restored = RenameOptions::from(FFI_RenameOptions::from(&original)); + assert_eq!(restored.target_mode, original.target_mode); + } + } + + #[test] + fn list_result_round_trip() { + let original = ListResult { + common_prefixes: vec![Path::from("a"), Path::from("b/c")], + objects: vec![ObjectMeta { + location: Path::from("a/1.parquet"), + last_modified: DateTime::from_timestamp(10, 0).unwrap(), + size: 12, + e_tag: None, + version: None, + }], + }; + + let restored = ListResult::from(FFI_ListResult::from(&original)); + assert_eq!(restored.common_prefixes, original.common_prefixes); + assert_eq!(restored.objects, original.objects); + } + + #[test] + fn put_payload_round_trip_preserves_blocks() { + let original: PutPayload = vec![ + bytes::Bytes::from_static(b"first"), + bytes::Bytes::from_static(b"second"), + ] + .into_iter() + .collect(); + + let restored = put_payload_from_ffi(put_payload_to_ffi(&original)); + assert_eq!(restored.content_length(), original.content_length()); + let restored_blocks: Vec<_> = restored.iter().cloned().collect(); + let original_blocks: Vec<_> = original.iter().cloned().collect(); + assert_eq!(restored_blocks, original_blocks); + } + + #[test] + fn timestamp_round_trip_beyond_nanosecond_range() { + // Year 2500, outside the i64-nanosecond representable range. + let original = DateTime::from_timestamp(16_725_225_600, 0).unwrap(); + let restored: DateTime = FFI_Timestamp::from(original).into(); + assert_eq!(restored, original); + } +} diff --git a/datafusion/ffi/src/execution/runtime_env.rs b/datafusion/ffi/src/execution/runtime_env.rs new file mode 100644 index 0000000000000..4702f83d26d6a --- /dev/null +++ b/datafusion/ffi/src/execution/runtime_env.rs @@ -0,0 +1,492 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`RuntimeEnv`]. +//! +//! # How the environment crosses +//! +//! Each component is carried across on its own terms: +//! +//! | Component | How it crosses | +//! |---|---| +//! | [`ObjectStoreRegistry`] | Shared: it is a trait object, wrapped by [`FFI_ObjectStoreRegistry`] | +//! | [`MemoryPool`] | Shared: it is a trait object, wrapped by [`FFI_MemoryPool`] | +//! | [`DiskManager`] | Configuration copied; each side keeps its own instance | +//! | [`CacheManager`] | Configuration copied; each side keeps its own instance | +//! +//! The first two are the ones that matter for correctness. Sharing the registry +//! lets a table provider and its session agree on object stores, and sharing +//! the memory pool makes the session's memory limit apply to a foreign plan. +//! +//! [`RuntimeEnv`] is a plain struct rather than a trait, so it is tempting to +//! pass an `Arc` across as an opaque pointer instead. That is +//! unsound: the struct is not `repr(C)`, and its field list depends on enabled +//! features (`parquet_encryption` adds one). Two libraries built against the +//! same DataFusion version but different feature sets disagree about the +//! layout, so reading through such a pointer is undefined behaviour rather than +//! merely version-fragile. +//! +//! # What the configuration copy does and does not give you +//! +//! [`DiskManager`] and [`CacheManager`] are concrete structs and cannot be +//! wrapped as trait objects, so only their configuration is copied and each +//! side builds its own. Two consequences are worth knowing: +//! +//! * Spill limits are enforced *per side*. Both disk managers honour the same +//! `max_temp_directory_size`, but they count independently, so total on-disk +//! usage can reach twice the configured limit. +//! * Caches are not shared, so file statistics and listings fetched by one side +//! are not reused by the other. +//! +//! The host's temp directory *paths* are deliberately not propagated. +//! [`DiskManager::temp_dir_paths`] returns directories the host has already +//! created and owns, and it deletes them when it is dropped. Pointing a second +//! disk manager at them would risk writing into a directory that has been +//! removed underneath it, so the foreign side gets its own OS temporary +//! directory instead. Whether spilling is enabled at all *is* propagated. +//! +//! # The local fast path +//! +//! When both sides are the same library, [`FFI_RuntimeEnv::as_local`] returns +//! the original `Arc` and none of the above applies: the two sides +//! share one runtime environment exactly, including its disk manager and +//! caches. +//! +//! [`DiskManager`]: datafusion_execution::disk_manager::DiskManager +//! [`DiskManager::temp_dir_paths`]: datafusion_execution::disk_manager::DiskManager::temp_dir_paths +//! [`CacheManager`]: datafusion_execution::cache::cache_manager::CacheManager +//! [`ObjectStoreRegistry`]: datafusion_execution::object_store::ObjectStoreRegistry +//! [`MemoryPool`]: datafusion_execution::memory_pool::MemoryPool + +use std::ffi::c_void; +use std::sync::Arc; +use std::time::Duration; + +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::memory_pool::MemoryPool; +use datafusion_execution::object_store::ObjectStoreRegistry; +use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + +use crate::util::FFI_Option; + +use super::memory_pool::FFI_MemoryPool; +use super::object_store::FFI_ObjectStoreRegistry; + +/// The parts of [`RuntimeEnv`] that are copied rather than shared. +/// +/// See the module documentation for why these are copied and what that costs. +#[repr(C)] +#[derive(Debug, Clone)] +pub struct FFI_RuntimeConfig { + /// Whether the providing side permits spilling to disk at all, so a + /// session that has disabled spilling keeps it disabled across the + /// boundary. + pub tmp_files_enabled: bool, + + /// Maximum temporary directory size in bytes. + pub max_temp_directory_size: u64, + + /// Maximum number of spill files opened by one merge pass. Zero means + /// unlimited. + pub max_spill_merge_fan_in: u64, + + pub metadata_cache_limit: u64, + pub list_files_cache_limit: u64, + pub list_files_cache_ttl_secs: FFI_Option, + pub file_statistics_cache_limit: u64, +} + +impl From<&RuntimeEnv> for FFI_RuntimeConfig { + fn from(runtime_env: &RuntimeEnv) -> Self { + Self { + tmp_files_enabled: runtime_env.disk_manager.tmp_files_enabled(), + max_temp_directory_size: runtime_env.disk_manager.max_temp_directory_size(), + max_spill_merge_fan_in: runtime_env.disk_manager.max_spill_merge_fan_in() + as u64, + metadata_cache_limit: runtime_env.cache_manager.get_metadata_cache_limit() + as u64, + list_files_cache_limit: runtime_env.cache_manager.get_list_files_cache_limit() + as u64, + list_files_cache_ttl_secs: runtime_env + .cache_manager + .get_list_files_cache_ttl() + .map(|ttl| ttl.as_secs()) + .into(), + file_statistics_cache_limit: runtime_env + .cache_manager + .get_file_statistic_cache_limit() + as u64, + } + } +} + +impl FFI_RuntimeConfig { + /// Apply this configuration to a [`RuntimeEnvBuilder`]. + fn apply(self, builder: RuntimeEnvBuilder) -> RuntimeEnvBuilder { + let disk_manager = DiskManagerBuilder::default() + .with_mode(if self.tmp_files_enabled { + DiskManagerMode::OsTmpDirectory + } else { + DiskManagerMode::Disabled + }) + .with_max_temp_directory_size(self.max_temp_directory_size) + .with_max_spill_merge_fan_in(self.max_spill_merge_fan_in as usize); + + builder + .with_disk_manager_builder(disk_manager) + .with_metadata_cache_limit(self.metadata_cache_limit as usize) + .with_object_list_cache_limit(self.list_files_cache_limit as usize) + .with_object_list_cache_ttl( + self.list_files_cache_ttl_secs + .into_option() + .map(Duration::from_secs), + ) + .with_file_statistics_cache_limit(self.file_statistics_cache_limit as usize) + } +} + +/// A stable struct for sharing [`RuntimeEnv`] across FFI boundaries. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_RuntimeEnv { + /// Return the object store registry. Shared, not copied. + pub object_store_registry: + unsafe extern "C" fn(runtime_env: &Self) -> FFI_ObjectStoreRegistry, + + /// Return the memory pool. Shared, not copied. + pub memory_pool: unsafe extern "C" fn(runtime_env: &Self) -> FFI_MemoryPool, + + /// Return the configuration of the components that are copied. + pub config: unsafe extern "C" fn(runtime_env: &Self) -> FFI_RuntimeConfig, + + /// Used to create a clone on the provider of the environment. This should + /// only need to be called by the receiver of the environment. + pub clone: unsafe extern "C" fn(runtime_env: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + pub release: unsafe extern "C" fn(arg: &mut Self), + + /// Internal data. This is only to be accessed by the provider of the + /// environment. The foreign library should never attempt to access this + /// data. + pub private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. See [`crate::get_library_marker_id`] and + /// the crate's `README.md` for more information. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_RuntimeEnv {} +unsafe impl Sync for FFI_RuntimeEnv {} + +struct RuntimeEnvPrivateData { + runtime_env: Arc, + runtime: Option, +} + +impl FFI_RuntimeEnv { + fn private_data(&self) -> &RuntimeEnvPrivateData { + unsafe { &*(self.private_data as *const RuntimeEnvPrivateData) } + } + + fn inner(&self) -> &Arc { + &self.private_data().runtime_env + } + + /// Create a new [`FFI_RuntimeEnv`] from a local runtime environment. + /// + /// `runtime` is the tokio runtime handle of the providing library, attached + /// to object stores handed out by the registry so that stores which spawn + /// tasks work when driven by a foreign executor. + pub fn new( + runtime_env: Arc, + runtime: Option, + ) -> Self { + Self { + object_store_registry: object_store_registry_fn_wrapper, + memory_pool: memory_pool_fn_wrapper, + config: config_fn_wrapper, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + private_data: Box::into_raw(Box::new(RuntimeEnvPrivateData { + runtime_env, + runtime, + })) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } + + /// If this environment originated in the current library, return the + /// underlying [`RuntimeEnv`] directly. + /// + /// This is an exact identity match, so the two sides share the disk manager + /// and caches as well as the registry and memory pool. + pub fn as_local(&self) -> Option> { + ((self.library_marker_id)() == crate::get_library_marker_id()) + .then(|| Arc::clone(self.inner())) + } +} + +unsafe extern "C" fn object_store_registry_fn_wrapper( + runtime_env: &FFI_RuntimeEnv, +) -> FFI_ObjectStoreRegistry { + let private_data = runtime_env.private_data(); + FFI_ObjectStoreRegistry::new( + Arc::clone(&private_data.runtime_env.object_store_registry), + private_data.runtime.clone(), + ) +} + +unsafe extern "C" fn memory_pool_fn_wrapper( + runtime_env: &FFI_RuntimeEnv, +) -> FFI_MemoryPool { + FFI_MemoryPool::new(Arc::clone(&runtime_env.inner().memory_pool)) +} + +unsafe extern "C" fn config_fn_wrapper( + runtime_env: &FFI_RuntimeEnv, +) -> FFI_RuntimeConfig { + FFI_RuntimeConfig::from(runtime_env.inner().as_ref()) +} + +unsafe extern "C" fn clone_fn_wrapper(runtime_env: &FFI_RuntimeEnv) -> FFI_RuntimeEnv { + let private_data = runtime_env.private_data(); + FFI_RuntimeEnv::new( + Arc::clone(&private_data.runtime_env), + private_data.runtime.clone(), + ) +} + +unsafe extern "C" fn release_fn_wrapper(runtime_env: &mut FFI_RuntimeEnv) { + unsafe { + debug_assert!(!runtime_env.private_data.is_null()); + drop(Box::from_raw( + runtime_env.private_data as *mut RuntimeEnvPrivateData, + )); + runtime_env.private_data = std::ptr::null_mut(); + } +} + +impl Clone for FFI_RuntimeEnv { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +impl Drop for FFI_RuntimeEnv { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl TryFrom<&FFI_RuntimeEnv> for Arc { + type Error = DataFusionError; + + fn try_from(runtime_env: &FFI_RuntimeEnv) -> Result { + if let Some(local) = runtime_env.as_local() { + return Ok(local); + } + + let registry: Arc = + unsafe { (runtime_env.object_store_registry)(runtime_env) }.into(); + let memory_pool: Arc = + unsafe { (runtime_env.memory_pool)(runtime_env) }.into(); + let config = unsafe { (runtime_env.config)(runtime_env) }; + + config + .apply( + RuntimeEnvBuilder::new() + .with_object_store_registry(registry) + .with_memory_pool(memory_pool), + ) + .build_arc() + } +} + +#[cfg(test)] +mod tests { + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryLimit, + }; + use datafusion_execution::object_store::ObjectStoreUrl; + use object_store::ObjectStore; + use object_store::memory::InMemory; + use url::Url; + + use super::*; + + /// `RuntimeEnv::object_store` takes an `ObjectStoreUrl` rather than a raw + /// `Url`. + fn store_url(url: &Url) -> ObjectStoreUrl { + ObjectStoreUrl::parse(url.as_str()).expect("valid object store url") + } + + fn foreign_runtime_env(runtime_env: Arc) -> Arc { + let mut ffi = FFI_RuntimeEnv::new(runtime_env, None); + ffi.library_marker_id = crate::mock_foreign_marker_id; + Arc::::try_from(&ffi).expect("build foreign runtime env") + } + + /// The bug this whole module exists to fix: a store registered on one side + /// must be visible from the other. + #[test] + fn object_store_registered_on_host_is_visible() -> Result<()> { + let host = RuntimeEnvBuilder::new().build_arc()?; + let url = Url::parse("s3://bucket").unwrap(); + host.register_object_store(&url, Arc::new(InMemory::new())); + + let foreign = foreign_runtime_env(Arc::clone(&host)); + assert!(foreign.object_store(store_url(&url)).is_ok()); + + Ok(()) + } + + /// A store registered by a table provider on the session it is handed + /// during planning must be visible to the host at execution time. + #[test] + fn object_store_registered_on_foreign_is_visible_to_host() -> Result<()> { + let host = RuntimeEnvBuilder::new().build_arc()?; + let foreign = foreign_runtime_env(Arc::clone(&host)); + + let url = Url::parse("s3://provider-registered").unwrap(); + foreign.register_object_store(&url, Arc::new(InMemory::new())); + + assert!( + host.object_store(store_url(&url)).is_ok(), + "store registered through the foreign runtime env should reach the host" + ); + + Ok(()) + } + + /// A store that a provider registers and then reads back must come back as + /// its own local store, so reads do not cross the boundary. + #[test] + fn round_tripped_store_is_recovered_locally() -> Result<()> { + let host = RuntimeEnvBuilder::new().build_arc()?; + let foreign = foreign_runtime_env(Arc::clone(&host)); + + let url = Url::parse("s3://mine").unwrap(); + let original = Arc::new(InMemory::new()) as Arc; + foreign.register_object_store(&url, Arc::clone(&original)); + + let recovered = foreign.object_store(store_url(&url))?; + assert!( + Arc::ptr_eq(&original, &recovered), + "a provider's own store should not be wrapped when read back" + ); + + Ok(()) + } + + #[test] + fn memory_limit_crosses_the_boundary() -> Result<()> { + let host = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1_000))) + .build_arc()?; + let foreign = foreign_runtime_env(Arc::clone(&host)); + + assert!(matches!( + foreign.memory_pool.memory_limit(), + MemoryLimit::Finite(1_000) + )); + + let reservation = MemoryConsumer::new("foreign").register(&foreign.memory_pool); + reservation.grow(600); + assert_eq!( + host.memory_pool.reserved(), + 600, + "foreign allocations must count against the host pool" + ); + + Ok(()) + } + + #[test] + fn disk_manager_config_is_copied() -> Result<()> { + let host = RuntimeEnvBuilder::new() + .with_max_temp_directory_size(4_096) + .with_max_spill_merge_fan_in(7) + .build_arc()?; + let foreign = foreign_runtime_env(host); + + assert_eq!(foreign.disk_manager.max_temp_directory_size(), 4_096); + assert_eq!(foreign.disk_manager.max_spill_merge_fan_in(), 7); + assert!(foreign.disk_manager.tmp_files_enabled()); + + Ok(()) + } + + /// A host that disabled spilling must not have a foreign plan spill behind + /// its back. + #[test] + fn disabled_spilling_is_propagated() -> Result<()> { + let host = RuntimeEnvBuilder::new() + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + assert!(!host.disk_manager.tmp_files_enabled()); + + let foreign = foreign_runtime_env(host); + assert!( + !foreign.disk_manager.tmp_files_enabled(), + "spilling must stay disabled across the boundary" + ); + + Ok(()) + } + + #[test] + fn cache_config_is_copied() -> Result<()> { + let host = RuntimeEnvBuilder::new() + .with_metadata_cache_limit(1_234) + .with_object_list_cache_limit(2_345) + .with_object_list_cache_ttl(Some(Duration::from_mins(1))) + .with_file_statistics_cache_limit(3_456) + .build_arc()?; + let foreign = foreign_runtime_env(host); + + assert_eq!(foreign.cache_manager.get_metadata_cache_limit(), 1_234); + assert_eq!(foreign.cache_manager.get_list_files_cache_limit(), 2_345); + assert_eq!( + foreign.cache_manager.get_list_files_cache_ttl(), + Some(Duration::from_mins(1)) + ); + assert_eq!( + foreign.cache_manager.get_file_statistic_cache_limit(), + 3_456 + ); + + Ok(()) + } + + /// Within one library the runtime environment is shared exactly, so the + /// config copy never runs. + #[test] + fn local_runtime_env_is_shared_exactly() -> Result<()> { + let original = RuntimeEnvBuilder::new().build_arc()?; + let ffi = FFI_RuntimeEnv::new(Arc::clone(&original), None); + + let recovered = Arc::::try_from(&ffi)?; + assert!(Arc::ptr_eq(&original, &recovered)); + + Ok(()) + } +} diff --git a/datafusion/ffi/src/execution/task_ctx.rs b/datafusion/ffi/src/execution/task_ctx.rs index 0a48a5fe5af36..d476eec53dbf8 100644 --- a/datafusion/ffi/src/execution/task_ctx.rs +++ b/datafusion/ffi/src/execution/task_ctx.rs @@ -25,10 +25,12 @@ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::{ AggregateUDF, AggregateUDFImpl, ScalarUDF, ScalarUDFImpl, WindowUDF, WindowUDFImpl, }; +use tokio::runtime::Handle; use stabby::string::String as SString; use stabby::vec::Vec as SVec; +use crate::execution::runtime_env::FFI_RuntimeEnv; use crate::session::config::FFI_SessionConfig; use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; @@ -58,6 +60,15 @@ pub struct FFI_TaskContext { /// Returns a vec of name-function pairs for window functions. pub window_functions: unsafe extern "C" fn(&Self) -> SVec<(SString, FFI_WindowUDF)>, + /// Returns the runtime environment. + /// + /// A plan executing on the far side of the boundary reaches the object + /// stores and the memory budget of the executing session through this, + /// so a store registered on the session during planning is available at + /// execution time and allocations count against the session's memory + /// limit. + pub runtime_env: unsafe extern "C" fn(&Self) -> FFI_RuntimeEnv, + /// Release the memory of the private data when it is no longer being used. pub release: unsafe extern "C" fn(arg: &mut Self), @@ -73,6 +84,9 @@ pub struct FFI_TaskContext { struct TaskContextPrivateData { ctx: Arc, + /// Tokio runtime handle of the providing library, attached to object stores + /// handed out by this context's runtime environment. + runtime: Option, } impl FFI_TaskContext { @@ -150,6 +164,16 @@ unsafe extern "C" fn window_functions_fn_wrapper( } } +unsafe extern "C" fn runtime_env_fn_wrapper(ctx: &FFI_TaskContext) -> FFI_RuntimeEnv { + unsafe { + let private_data = ctx.private_data as *const TaskContextPrivateData; + FFI_RuntimeEnv::new( + Arc::clone(&(*private_data).ctx.runtime_env()), + (*private_data).runtime.clone(), + ) + } +} + unsafe extern "C" fn release_fn_wrapper(ctx: &mut FFI_TaskContext) { unsafe { let private_data = Box::from_raw(ctx.private_data as *mut TaskContextPrivateData); @@ -163,9 +187,21 @@ impl Drop for FFI_TaskContext { } } -impl From> for FFI_TaskContext { - fn from(ctx: Arc) -> Self { - let private_data = Box::new(TaskContextPrivateData { ctx }); +impl FFI_TaskContext { + /// Create a new [`FFI_TaskContext`] from a local task context. + /// + /// `runtime` is the tokio runtime handle of the library creating this + /// context. It is attached to the object stores reached through this + /// context's runtime environment and entered while they are polled, so + /// that stores which spawn tasks or use timers work when driven by a + /// foreign executor. + /// + /// Pass `None` only when the context will not be used to reach an object + /// store, such as when it serves purely as a + /// [`FunctionRegistry`](datafusion_expr::registry::FunctionRegistry) + /// while encoding or decoding a plan. + pub fn new(ctx: Arc, runtime: Option) -> Self { + let private_data = Box::new(TaskContextPrivateData { ctx, runtime }); FFI_TaskContext { session_id: session_id_fn_wrapper, @@ -174,6 +210,7 @@ impl From> for FFI_TaskContext { scalar_functions: scalar_functions_fn_wrapper, aggregate_functions: aggregate_functions_fn_wrapper, window_functions: window_functions_fn_wrapper, + runtime_env: runtime_env_fn_wrapper, release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, @@ -228,7 +265,18 @@ impl From for Arc { }) .collect(); - let runtime = Arc::new(RuntimeEnv::default()); + // The providing side's runtime environment carries the registered + // object stores and the memory pool this context should execute + // against. + let ffi_runtime_env = (ffi_ctx.runtime_env)(&ffi_ctx); + let runtime_env = >::try_from(&ffi_runtime_env) + .unwrap_or_else(|e| { + log::warn!( + "Unable to reconstruct the runtime environment across \ + the FFI boundary, falling back to a default: {e}" + ); + Arc::new(RuntimeEnv::default()) + }); Arc::new(TaskContext::new( task_id, @@ -238,7 +286,7 @@ impl From for Arc { HashMap::new(), aggregate_functions, window_functions, - runtime, + runtime_env, )) } } @@ -258,7 +306,7 @@ mod tests { fn ffi_task_ctx_round_trip() -> Result<()> { let session_ctx = SessionContext::new(); let original = session_ctx.task_ctx(); - let mut ffi_task_ctx = FFI_TaskContext::from(Arc::clone(&original)); + let mut ffi_task_ctx = FFI_TaskContext::new(Arc::clone(&original), None); ffi_task_ctx.library_marker_id = crate::mock_foreign_marker_id; let foreign_task_ctx: Arc = ffi_task_ctx.into(); diff --git a/datafusion/ffi/src/execution/task_ctx_provider.rs b/datafusion/ffi/src/execution/task_ctx_provider.rs index b8fa68a16ae1f..09af169fdc3b6 100644 --- a/datafusion/ffi/src/execution/task_ctx_provider.rs +++ b/datafusion/ffi/src/execution/task_ctx_provider.rs @@ -79,7 +79,10 @@ unsafe extern "C" fn task_ctx_fn_wrapper( sresult!( ctx_provider .inner() - .map(FFI_TaskContext::from) + // This context is used as a `FunctionRegistry` while + // encoding and decoding plans, not to reach object stores, so + // it needs no runtime handle. + .map(|ctx| FFI_TaskContext::new(ctx, None)) .ok_or_else(|| { ffi_datafusion_err!( "TaskContextProvider went out of scope over FFI boundary." diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index d7ee5dace30cc..5a5b8f07639f5 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -491,7 +491,9 @@ impl ExecutionPlan for ForeignExecutionPlan { partition: usize, context: Arc, ) -> Result { - let context = FFI_TaskContext::from(context); + // Attach the runtime this plan is being executed on. The foreign + // plan enters it when polling object stores owned by this library. + let context = FFI_TaskContext::new(context, Handle::try_current().ok()); unsafe { df_result!((self.plan.execute)(&self.plan, partition, context)) .map(|stream| Pin::new(Box::new(stream)) as SendableRecordBatchStream) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index f215a6ba5a568..15339862ac4f9 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -73,7 +73,7 @@ use tokio::runtime::Handle; use crate::arrow_wrappers::WrappedSchema; use crate::catalog_provider_list::FFI_CatalogProviderList; -use crate::execution::FFI_TaskContext; +use crate::execution::{FFI_RuntimeEnv, FFI_TaskContext}; use crate::execution_plan::FFI_ExecutionPlan; use crate::physical_expr::FFI_PhysicalExpr; use crate::physical_optimizer::FFI_PhysicalOptimizerRule; @@ -141,6 +141,13 @@ pub(crate) struct FFI_SessionRef { task_ctx: unsafe extern "C" fn(&Self) -> FFI_TaskContext, + /// Returns the session's runtime environment. + /// + /// This is the session's own environment, so a table provider that + /// registers an object store on it during planning finds that store again + /// when the resulting plan is executed. + runtime_env: unsafe extern "C" fn(&Self) -> FFI_RuntimeEnv, + physical_optimizers: unsafe extern "C" fn(&Self) -> SVec, logical_codec: FFI_LogicalExtensionCodec, @@ -369,7 +376,15 @@ unsafe extern "C" fn default_table_options_fn_wrapper( } unsafe extern "C" fn task_ctx_fn_wrapper(session: &FFI_SessionRef) -> FFI_TaskContext { - session.inner().task_ctx().into() + FFI_TaskContext::new(session.inner().task_ctx(), unsafe { + session.runtime().clone() + }) +} + +unsafe extern "C" fn runtime_env_fn_wrapper(session: &FFI_SessionRef) -> FFI_RuntimeEnv { + FFI_RuntimeEnv::new(Arc::clone(session.inner().runtime_env()), unsafe { + session.runtime().clone() + }) } unsafe extern "C" fn physical_optimizers_fn_wrapper( @@ -415,6 +430,7 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_SessionRef) -> FFI_SessionR table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, + runtime_env: runtime_env_fn_wrapper, physical_optimizers: physical_optimizers_fn_wrapper, logical_codec: provider.logical_codec.clone(), physical_codec: provider.physical_codec.clone(), @@ -507,6 +523,7 @@ impl FFI_SessionRef { table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, + runtime_env: runtime_env_fn_wrapper, physical_optimizers: physical_optimizers_fn_wrapper, logical_codec, physical_codec, @@ -618,7 +635,14 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { aggregate_functions, window_functions, extension_types: Arc::new(MemoryExtensionTypeRegistry::default()), - runtime_env: Default::default(), + // `Session::runtime_env` returns a borrow, so this has to be + // built once here rather than lazily per call. That also makes + // the identity stable: a table provider that registers an + // object store through this environment finds it again later, + // which would not hold if each call rebuilt the environment. + runtime_env: >::try_from(&(session.runtime_env)( + session, + ))?, props: Default::default(), query_planner: OnceLock::new(), physical_optimizers: OnceLock::new(), diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index fbc3e83ba49fc..1d7fb3fa84919 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -59,6 +59,7 @@ use crate::util::FFI_Option; mod async_provider; pub mod catalog; pub mod config; +pub mod object_store_provider; mod physical_optimizer; mod query_planner; mod sync_provider; @@ -139,6 +140,11 @@ pub struct ForeignLibraryModule { /// Create an aggregate UDAF using first_value pub create_first_value_udaf: extern "C" fn() -> FFI_AggregateUDF, + + /// Construct a table provider that registers its own object store on the + /// session during planning and reads it back at execution time. + pub create_object_store_table: + extern "C" fn(codec: FFI_LogicalExtensionCodec) -> FFI_TableProvider, } pub fn create_test_schema() -> Arc { @@ -370,5 +376,6 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_query_planner: query_planner::create_query_planner, version: super::version, create_first_value_udaf: create_ffi_first_value_func, + create_object_store_table: object_store_provider::create_object_store_table, } } diff --git a/datafusion/ffi/src/tests/object_store_provider.rs b/datafusion/ffi/src/tests/object_store_provider.rs new file mode 100644 index 0000000000000..107088e943b50 --- /dev/null +++ b/datafusion/ffi/src/tests/object_store_provider.rs @@ -0,0 +1,259 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A table provider that reproduces the pattern used by table providers backed +//! by remote storage, such as `delta-rs`. +//! +//! The provider builds its *own* object store, registers it on the session it +//! is handed during planning, and returns a plan that looks that store up again +//! at execution time. Planning and execution happen on opposite sides of the +//! FFI boundary, so this exercises the session's runtime environment crossing +//! that boundary intact. +//! +//! The scan also reports the memory pool limit it observes at execution time, +//! letting the same test check that the host's memory limit reaches a foreign +//! plan. + +use std::sync::Arc; + +use arrow::array::{Int32Array, RecordBatch, UInt64Array}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use async_trait::async_trait; +use datafusion_catalog::{Session, TableProvider}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, exec_datafusion_err, exec_err}; +use datafusion_execution::memory_pool::MemoryLimit; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::{Expr, TableType}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion_physical_plan::Partitioning; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::execution_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, +}; +use datafusion_physical_plan::stream::RecordBatchStreamAdapter; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, +}; +use object_store::memory::InMemory; +use object_store::path::Path; +use object_store::{ObjectStoreExt, PutPayload}; +use url::Url; + +use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::table_provider::FFI_TableProvider; + +/// The URL the provider registers its store under. +pub const OBJECT_STORE_URL: &str = "ffitest://ffi-object-store"; + +/// The object within that store holding the scan's data. +const DATA_PATH: &str = "data.bin"; + +/// The values the scan produces, encoded little-endian into the object above. +pub const EXPECTED_VALUES: [i32; 5] = [10, 20, 30, 40, 50]; + +/// Reported by the scan when the memory pool it sees has no finite limit, +/// letting a test distinguish the session's pool from an unbounded default. +pub const UNLIMITED_MEMORY: u64 = u64::MAX; + +pub fn object_store_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("mem_limit", DataType::UInt64, false), + ])) +} + +#[derive(Debug)] +struct ObjectStoreTableProvider; + +#[async_trait] +impl TableProvider for ObjectStoreTableProvider { + fn schema(&self) -> SchemaRef { + object_store_schema() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + session: &dyn Session, + _projection: Option<&[usize]>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + // Build a store owned by *this* library and populate it. + let store = InMemory::new(); + let payload: Vec = EXPECTED_VALUES + .iter() + .flat_map(|v| v.to_le_bytes()) + .collect(); + store + .put(&Path::from(DATA_PATH), PutPayload::from(payload)) + .await + .map_err(|e| exec_datafusion_err!("Unable to seed the test store: {e}"))?; + + // Register it on the session handed to us during planning. The plan + // returned below looks it up again at execution time, by which point + // the session lives on the other side of the boundary. + let url = Url::parse(OBJECT_STORE_URL) + .map_err(|e| exec_datafusion_err!("Invalid test store URL: {e}"))?; + session + .runtime_env() + .register_object_store(&url, Arc::new(store)); + + Ok(Arc::new(ObjectStoreScanExec::new())) + } +} + +#[derive(Debug)] +struct ObjectStoreScanExec { + props: Arc, +} + +impl ObjectStoreScanExec { + fn new() -> Self { + Self { + props: Arc::new(PlanProperties::new( + EquivalenceProperties::new(object_store_schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } +} + +impl DisplayAs for ObjectStoreScanExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "ObjectStoreScanExec") + } +} + +impl ExecutionPlan for ObjectStoreScanExec { + fn name(&self) -> &'static str { + "ObjectStoreScanExec" + } + + fn properties(&self) -> &Arc { + &self.props + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn replace_children( + self: Arc, + _children: Vec>, + _options: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition != 0 { + return exec_err!("ObjectStoreScanExec only has one partition"); + } + + let schema = object_store_schema(); + let stream_schema = Arc::clone(&schema); + + let stream = futures::stream::once(async move { + let runtime_env = context.runtime_env(); + + // The store registered during planning must be reachable here. + let url = ObjectStoreUrl::parse(OBJECT_STORE_URL)?; + let store = runtime_env.object_store(url)?; + + let bytes = store + .get(&Path::from(DATA_PATH)) + .await + .map_err(|e| exec_datafusion_err!("Unable to read the test store: {e}"))? + .bytes() + .await + .map_err(|e| { + exec_datafusion_err!("Unable to collect the test store bytes: {e}") + })?; + + let values: Vec = bytes + .chunks_exact(4) + .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + + let memory_limit = match runtime_env.memory_pool.memory_limit() { + MemoryLimit::Finite(limit) => limit as u64, + MemoryLimit::Infinite | MemoryLimit::Unknown => UNLIMITED_MEMORY, + }; + + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(values)), + Arc::new(UInt64Array::from(vec![ + memory_limit; + EXPECTED_VALUES.len() + ])), + ], + ) + .map_err(Into::into) + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + stream_schema, + stream, + ))) + } +} + +pub(crate) extern "C" fn create_object_store_table( + codec: FFI_LogicalExtensionCodec, +) -> FFI_TableProvider { + FFI_TableProvider::new_with_ffi_codec( + Arc::new(ObjectStoreTableProvider), + true, + None, + codec, + ) +} diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index cba334e8fae57..9cd05e5ecb667 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -137,6 +137,85 @@ mod tests { Ok(()) } + /// A table provider that builds its own object store, registers it on the + /// session during planning, and reads it back during execution. + /// + /// Planning happens in the loaded module and execution is driven from this + /// executable, so the session's `RuntimeEnv` has to cross the FFI boundary + /// intact for the store to be found at execution time. + /// + /// The same scan reports the memory pool limit it observes, checking that + /// `datafusion.execution.memory_limit` reaches a foreign plan. + #[tokio::test] + async fn test_object_store_crosses_ffi_boundary() -> Result<()> { + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::prelude::{SessionConfig, SessionContext}; + use datafusion_execution::memory_pool::GreedyMemoryPool; + use datafusion_ffi::tests::object_store_provider::{ + EXPECTED_VALUES, UNLIMITED_MEMORY, + }; + use std::sync::Arc; + + const MEMORY_LIMIT: usize = 64 * 1024 * 1024; + + let module = get_module()?; + + // Build a context with a distinctive memory limit so the value observed + // inside the module identifies which pool it actually reached. + let runtime_env = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(MEMORY_LIMIT))) + .build_arc()?; + let ctx = Arc::new(SessionContext::new_with_config_rt( + SessionConfig::new(), + runtime_env, + )); + let codec = super::utils::codec_for(&ctx); + + let ffi_provider = (module.create_object_store_table)(codec); + let foreign: Arc = (&ffi_provider).into(); + + ctx.register_table("remote_table", foreign)?; + let batches = ctx.table("remote_table").await?.collect().await?; + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, + EXPECTED_VALUES.len(), + "scan should read every value back out of the registered store" + ); + + let values: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .expect("column a should be Int32") + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, EXPECTED_VALUES.to_vec()); + + let observed_limit = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .expect("column mem_limit should be UInt64") + .value(0); + assert_ne!( + observed_limit, UNLIMITED_MEMORY, + "the foreign plan saw an unbounded pool, so the host's memory limit \ + did not cross the boundary" + ); + assert_eq!( + observed_limit, MEMORY_LIMIT as u64, + "the foreign plan should see the host's configured memory limit" + ); + + Ok(()) + } + #[tokio::test] async fn test_table_provider_factory() -> Result<()> { let table_provider_module = get_module()?; diff --git a/datafusion/ffi/tests/utils/mod.rs b/datafusion/ffi/tests/utils/mod.rs index acf59de7f3464..452f10960a908 100644 --- a/datafusion/ffi/tests/utils/mod.rs +++ b/datafusion/ffi/tests/utils/mod.rs @@ -31,13 +31,23 @@ use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; #[cfg_attr(not(feature = "integration-tests"), expect(dead_code))] pub fn ctx_and_codec() -> (Arc, FFI_LogicalExtensionCodec) { let ctx = Arc::new(SessionContext::default()); - let task_ctx_provider = Arc::clone(&ctx) as Arc; + let codec = codec_for(&ctx); + + (ctx, codec) +} + +/// Build an FFI logical extension codec bound to an existing context. +/// +/// Use this instead of [`ctx_and_codec`] when the test needs a context +/// configured a particular way, for example with a specific memory pool. +#[cfg_attr(not(feature = "integration-tests"), expect(dead_code))] +pub fn codec_for(ctx: &Arc) -> FFI_LogicalExtensionCodec { + let task_ctx_provider = Arc::clone(ctx) as Arc; let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); - let codec = FFI_LogicalExtensionCodec::new( + + FFI_LogicalExtensionCodec::new( Arc::new(DefaultLogicalExtensionCodec {}), None, task_ctx_provider, - ); - - (ctx, codec) + ) } diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 1a6b983025e39..53aef93408323 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -70,3 +70,56 @@ Use the walk instead: StatisticsContext::new_with_registry(registry) .compute_extended(plan, &StatisticsArgs::new())?; // or .compute(...) for core Statistics ``` + +### FFI table providers now share the session's `RuntimeEnv` + +A `Session` shared over FFI now carries its `RuntimeEnv` across the boundary. +Object stores registered on the session during `TableProvider::scan` are found +at execution time, and the session's memory limit applies to plans shared over +FFI. Previously a scan that relied on a store registered during planning failed +with `No suitable object store found`. + +**Who is affected:** + +- Anyone building or consuming `datafusion-ffi` libraries. The ABI of + `FFI_TaskContext` changed. +- FFI table providers backed by remote storage. +- FFI plans that allocate large amounts of memory or spill to disk. + +**Migration guide:** + +Rebuild every FFI provider and consumer against DataFusion 56, and do not +exchange these structs with libraries built against an older major version. +`datafusion-ffi` now depends on `object_store` directly and exposes its types in +its public API, so providers and consumers must agree on the `object_store` +version as well as the DataFusion version. + +`impl From> for FFI_TaskContext` has been removed because it +could not supply a tokio runtime handle. Call `FFI_TaskContext::new` instead: + +```diff +- let ffi_ctx = FFI_TaskContext::from(task_ctx); ++ let ffi_ctx = FFI_TaskContext::new(task_ctx, Handle::try_current().ok()); +``` + +Pass `None` as the handle only when the context will not be used to reach an +object store, such as when it serves purely as a `FunctionRegistry` while +encoding or decoding a plan. If you construct `FFI_TaskContext` with a struct +literal, populate the new `runtime_env` field. + +If you worked around object stores not crossing the boundary, that workaround +can be removed. Registering the store on the session from inside +`TableProvider::scan` now works. + +If an FFI plan starts failing with `DataFusionError::ResourcesExhausted`, it is +being charged against the session's memory pool for the first time. Raise +`datafusion.execution.memory_limit` or reduce the plan's memory use. + +`DiskManager` and `CacheManager` are not shared; each side builds its own from +copied configuration. If you sized `datafusion.runtime.max_temp_directory_size` +against total disk, note that both sides now enforce it independently, so +spilling can use up to twice that amount. + +See the `datafusion_ffi::execution::runtime_env` module documentation for +exactly what is shared and what is copied, and +[issue #19277](https://github.com/apache/datafusion/issues/19277) for details. From f509f500641ad1e81f3b21ad6f8007fae3d8dd10 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 14:43:13 -0400 Subject: [PATCH 2/2] spell check --- datafusion/ffi/src/execution/object_store/registry.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/ffi/src/execution/object_store/registry.rs b/datafusion/ffi/src/execution/object_store/registry.rs index 057bab95c0315..28742216ead82 100644 --- a/datafusion/ffi/src/execution/object_store/registry.rs +++ b/datafusion/ffi/src/execution/object_store/registry.rs @@ -140,10 +140,10 @@ unsafe extern "C" fn register_store_fn_wrapper( store: FFI_ObjectStore, ) -> FFI_Option { let Ok(url) = parse_url(&url) else { - // `register_store` cannot report an error. An unparseable URL can only + // `register_store` cannot report an error. An unparsable URL can only // come from a caller that built one by hand, and silently dropping the // registration would be worse than a log line. - log::warn!("Ignoring object store registration for unparseable URL '{url}'"); + log::warn!("Ignoring object store registration for unparsable URL '{url}'"); return FFI_Option::None; };