You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by pa...@apache.org on 2026/09/25 20:59:28 UTC
(sedona-db) branch main updated: feat(c/sedona-extension): implement async FFI execution (#1333)
This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new ae903aa79 feat(c/sedona-extension): implement async FFI execution (#1333)
ae903aa79 is described below
commit ae903aa795db8317bb5eb865be2da4d34c795252
Author: Dewey Dunnington <de...@dunnington.ca>
AuthorDate: Fri Sep 25 20:59:21 2026 +0000
feat(c/sedona-extension): implement async FFI execution (#1333)
---
Cargo.lock | 1 +
c/sedona-extension/src/execution_plan.rs | 196 +++-
.../src/export_sendable_record_batch_stream.rs | 597 ++++++++++++
c/sedona-extension/src/extension.rs | 98 ++
.../src/import_sendable_record_batch_stream.rs | 1027 ++++++++++++++++++++
c/sedona-extension/src/lib.rs | 2 +
c/sedona-extension/src/sedona_extension.h | 66 +-
c/sedona-extension/src/table_provider.rs | 43 +-
c/sedona-extension/tests/async_stream_interop.rs | 398 ++++++++
.../tests/async_stream_regressions.rs | 546 +++++++++++
python/sedonadb/Cargo.toml | 1 +
python/sedonadb/src/context.rs | 12 +-
python/sedonadb/src/import_from.rs | 7 +-
python/sedonadb/tests/test_dataframe_ffi.py | 14 +-
rust/sedona-common/src/option.rs | 21 +
15 files changed, 3010 insertions(+), 19 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 95a542799..cd1b7512c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -6435,6 +6435,7 @@ dependencies = [
"pyo3",
"sedona",
"sedona-adbc",
+ "sedona-common",
"sedona-datasource",
"sedona-expr",
"sedona-extension",
diff --git a/c/sedona-extension/src/execution_plan.rs b/c/sedona-extension/src/execution_plan.rs
index c35697deb..6bf055797 100644
--- a/c/sedona-extension/src/execution_plan.rs
+++ b/c/sedona-extension/src/execution_plan.rs
@@ -37,7 +37,10 @@ use datafusion_physical_plan::{
use sedona_common::{sedona_internal_datafusion_err, sedona_internal_err};
use serde::{Deserialize, Serialize};
+use crate::export_sendable_record_batch_stream::drive_stream_to_handler;
+use crate::extension::FFI_ArrowAsyncDeviceStreamHandler;
use crate::extension::{SedonaCError, SedonaCExecutionPlan, SedonaCExecutionPlanArgs};
+use crate::import_sendable_record_batch_stream::ImportedAsyncDeviceStream;
use crate::metrics::SerializableExecutionPlanMetricsSet;
use crate::runtime::RuntimeHandle;
use crate::set_ffi_error;
@@ -148,7 +151,7 @@ impl From<ExportedExecutionPlan> for SedonaCExecutionPlan {
get_property: Some(c_exec_plan_get_property),
with_property: None,
execute: Some(c_exec_plan_execute),
- execute_async: None,
+ execute_async: Some(c_exec_plan_execute_async),
reserved: null_mut(),
release: Some(c_exec_plan_release),
private_data: Box::into_raw(boxed) as *mut c_void,
@@ -272,6 +275,56 @@ unsafe extern "C" fn c_exec_plan_execute(
}
}
+unsafe extern "C" fn c_exec_plan_execute_async(
+ self_: *const SedonaCExecutionPlan,
+ args: *mut SedonaCExecutionPlanArgs,
+ out: *mut c_void,
+ err: *mut SedonaCError,
+) -> c_int {
+ debug_assert!(!self_.is_null(), "self pointer is null");
+ debug_assert!(!args.is_null(), "args pointer is null");
+ debug_assert!(!out.is_null(), "out pointer is null");
+ let self_ref = &*self_;
+ let args_ref = &*args;
+ debug_assert!(!self_ref.private_data.is_null(), "private_data is null");
+ let plan = &*(self_ref.private_data as *const ExportedExecutionPlan);
+
+ let args_slice = if args_ref.args.is_null() || args_ref.args_len == 0 {
+ &[]
+ } else {
+ std::slice::from_raw_parts(args_ref.args, args_ref.args_len)
+ };
+ let execute_args: ExecuteArgs = match serde_json::from_slice(args_slice) {
+ Ok(args) => args,
+ Err(e) => {
+ set_ffi_error!(err, "Failed to parse execute args: {}", e);
+ return libc::EINVAL;
+ }
+ };
+
+ match plan.execute(execute_args.partition) {
+ Ok(stream) => {
+ // The driver owns sendable guards for the handler and producer, so
+ // slow streams remain ordinary async tasks instead of occupying a
+ // blocking-pool thread for their entire lifetime.
+ let driver =
+ drive_stream_to_handler(stream, out as *mut FFI_ArrowAsyncDeviceStreamHandler);
+ // A returned stream may outlive the exported plan and context.
+ // Keep their runtime alive until the driver has released the handler.
+ let runtime = plan.runtime.clone();
+ plan.runtime.spawn(async move {
+ let _runtime = runtime;
+ driver.await;
+ });
+ ERRNO_OK
+ }
+ Err(e) => {
+ set_ffi_error!(err, "{}", e);
+ libc::EINVAL
+ }
+ }
+}
+
unsafe extern "C" fn c_exec_plan_release(self_: *mut SedonaCExecutionPlan) {
debug_assert!(!self_.is_null(), "self pointer is null");
let self_ref = &mut *self_;
@@ -296,6 +349,7 @@ pub struct ImportedSedonaCExec {
cancel_checker: Option<Arc<dyn Fn() -> bool + Send + Sync>>,
/// Interval for periodic cancellation checking during stream consumption.
check_interval: Option<Duration>,
+ use_async: bool,
}
impl Debug for ImportedSedonaCExec {
@@ -351,6 +405,7 @@ impl ImportedSedonaCExec {
name,
cancel_checker: None,
check_interval: None,
+ use_async: false,
})
}
@@ -377,6 +432,15 @@ impl ImportedSedonaCExec {
self
}
+ /// Select the Arrow async-device stream callback for execution.
+ ///
+ /// Synchronous execution remains the default for compatibility. When this
+ /// is enabled, `execute_async` must be implemented by the imported plan.
+ pub fn with_async_execution(mut self, use_async: bool) -> Self {
+ self.use_async = use_async;
+ self
+ }
+
fn get_debug_string(&self) -> Result<String> {
get_plan_string_property(&self.inner, "debug_string")
}
@@ -451,10 +515,6 @@ impl ExecutionPlan for ImportedSedonaCExec {
partition: usize,
_context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
- let Some(execute) = self.inner.execute else {
- return sedona_internal_err!("SedonaCExecutionPlan does not have execute");
- };
-
let args = ExecuteArgs { partition };
let args_bytes = serde_json::to_vec(&args).map_err(|e| {
sedona_internal_datafusion_err!("Failed to serialize execute args: {}", e)
@@ -470,9 +530,36 @@ impl ExecutionPlan for ImportedSedonaCExec {
reserved: null_mut(),
};
- let mut ffi_stream = FFI_ArrowArrayStream::empty();
let mut err = SedonaCError::default();
+ if self.use_async {
+ let Some(execute_async) = self.inner.execute_async else {
+ return sedona_internal_err!("SedonaCExecutionPlan does not have execute_async");
+ };
+ let (stream, handler) =
+ ImportedAsyncDeviceStream::new_with_schema(2, self.schema.clone());
+ let code = unsafe {
+ execute_async(
+ &self.inner,
+ &mut ffi_args,
+ handler.as_ptr().cast(),
+ &mut err,
+ )
+ };
+ if code != ERRNO_OK {
+ return exec_err!("Failed to execute plan asynchronously: {}", err);
+ }
+ // From this point the producer owns the handler and releases it
+ // after end-of-stream, error, or cancellation.
+ let _ = handler.into_raw();
+ return Ok(stream.into_sendable());
+ }
+
+ let Some(execute) = self.inner.execute else {
+ return sedona_internal_err!("SedonaCExecutionPlan does not have execute");
+ };
+ let mut ffi_stream = FFI_ArrowArrayStream::empty();
+
let code = unsafe { execute(&self.inner, &mut ffi_args, &mut ffi_stream, &mut err) };
if code != ERRNO_OK {
@@ -584,7 +671,7 @@ mod tests {
metrics::{ExecutionPlanMetricsSet, MetricBuilder},
stream::RecordBatchStreamAdapter,
};
- use futures::{stream, StreamExt};
+ use futures::{future, pin_mut, stream, StreamExt};
use std::fmt::Formatter;
/// A dummy ExecutionPlan with fixed, predictable values for testing FFI roundtrip.
@@ -933,4 +1020,99 @@ ImportedSedonaCExec
assert_batches_eq!(expected2, &batches2);
});
}
+
+ #[test]
+ fn test_execution_plan_roundtrip_execute_async() {
+ let (imported, task_ctx, runtime) = setup_imported_plan();
+ let imported = imported.with_async_execution(true);
+
+ runtime.block_on(async {
+ let stream = imported.execute(2, task_ctx).unwrap();
+ let batches = stream
+ .collect::<Vec<_>>()
+ .await
+ .into_iter()
+ .collect::<Result<Vec<_>>>()
+ .unwrap();
+
+ assert_batches_eq!(
+ [
+ "+----+-------+",
+ "| id | value |",
+ "+----+-------+",
+ "| 21 | 100 |",
+ "| 22 | 200 |",
+ "| 23 | 300 |",
+ "+----+-------+",
+ ],
+ &batches
+ );
+ });
+ }
+
+ #[test]
+ fn test_async_execution_does_not_use_the_blocking_pool() {
+ let runtime = Arc::new(RuntimeHandle::new(
+ tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .max_blocking_threads(1)
+ .build()
+ .unwrap(),
+ ));
+ let task_ctx = Arc::new(TaskContext::default());
+ let exported = ExportedExecutionPlan::new(
+ Arc::new(DummyExec::new()),
+ task_ctx.clone(),
+ runtime.clone(),
+ );
+ let imported = ImportedSedonaCExec::try_new(exported.into())
+ .unwrap()
+ .with_async_execution(true);
+
+ // Occupy the runtime's only blocking worker. The async FFI driver must
+ // still start and complete on the current-thread scheduler.
+ let (blocker_started_tx, blocker_started_rx) = std::sync::mpsc::channel();
+ let (release_blocker_tx, release_blocker_rx) = std::sync::mpsc::channel();
+ let blocker = runtime.spawn_blocking(move || {
+ blocker_started_tx.send(()).unwrap();
+ release_blocker_rx.recv().unwrap();
+ });
+ blocker_started_rx
+ .recv_timeout(Duration::from_secs(2))
+ .expect("blocking worker did not start");
+
+ let (timeout_tx, timeout_rx) = futures::channel::oneshot::channel();
+ std::thread::spawn(move || {
+ std::thread::sleep(Duration::from_secs(2));
+ let _ = timeout_tx.send(());
+ });
+
+ let completed = runtime.block_on(async {
+ let collection = async {
+ let stream = imported.execute(0, task_ctx).unwrap();
+ stream
+ .collect::<Vec<_>>()
+ .await
+ .into_iter()
+ .collect::<Result<Vec<_>>>()
+ .unwrap()
+ };
+ pin_mut!(collection);
+ pin_mut!(timeout_rx);
+ match future::select(collection, timeout_rx).await {
+ future::Either::Left((batches, _)) => {
+ assert_eq!(batches.len(), 1);
+ true
+ }
+ future::Either::Right(_) => false,
+ }
+ });
+
+ release_blocker_tx.send(()).unwrap();
+ runtime.block_on(blocker).unwrap();
+ assert!(
+ completed,
+ "async FFI execution waited for the blocking pool"
+ );
+ }
}
diff --git a/c/sedona-extension/src/export_sendable_record_batch_stream.rs b/c/sedona-extension/src/export_sendable_record_batch_stream.rs
new file mode 100644
index 000000000..f607d0faf
--- /dev/null
+++ b/c/sedona-extension/src/export_sendable_record_batch_stream.rs
@@ -0,0 +1,597 @@
+// 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 wrapper for [`SendableRecordBatchStream`] using the Arrow C Device Data Interface.
+//!
+//! This module provides interoperability between DataFusion's async record batch streams
+//! and C code via the [`FFI_ArrowAsyncDeviceStreamHandler`] interface.
+//!
+//! The implementation follows the Arrow Async Device Stream specification, including:
+//! - Back-pressure via `request()` callback
+//! - Consumer-initiated cancellation via `cancel()` callback
+//! - Proper end-of-stream signaling
+
+use std::ffi::{c_int, c_void, CString};
+use std::future::Future;
+use std::ptr::null_mut;
+use std::ptr::NonNull;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::task::Poll;
+
+use arrow_array::ffi::{to_ffi, FFI_ArrowSchema};
+use arrow_array::RecordBatch;
+use futures::task::AtomicWaker;
+use futures::StreamExt;
+
+pub use datafusion_execution::SendableRecordBatchStream;
+
+use crate::extension::{
+ FFI_ArrowAsyncDeviceStreamHandler, FFI_ArrowAsyncProducer, FFI_ArrowAsyncTask,
+ FFI_ArrowDeviceArray, ARROW_DEVICE_CPU,
+};
+
+/// Yields control back to the async executor once, then resumes.
+/// This allows other tasks (e.g., the consumer) to make progress.
+async fn yield_once() {
+ let mut yielded = false;
+ futures::future::poll_fn(|cx| {
+ if yielded {
+ Poll::Ready(())
+ } else {
+ yielded = true;
+ cx.waker().wake_by_ref();
+ Poll::Pending
+ }
+ })
+ .await
+}
+
+/// Shared state between the async driver and the FFI producer callbacks.
+struct ProducerState {
+ /// Number of batches requested by consumer (back-pressure).
+ requested: AtomicU64,
+ /// Set to true when consumer calls cancel().
+ cancelled: AtomicBool,
+ /// Waker to wake when requests are available or cancelled.
+ waker: AtomicWaker,
+}
+
+impl ProducerState {
+ fn new() -> Self {
+ Self {
+ requested: AtomicU64::new(0),
+ cancelled: AtomicBool::new(false),
+ waker: AtomicWaker::new(),
+ }
+ }
+
+ fn request(&self, n: u64) {
+ self.requested.fetch_add(n, Ordering::SeqCst);
+ self.wake();
+ }
+
+ fn cancel(&self) {
+ self.cancelled.store(true, Ordering::SeqCst);
+ self.wake();
+ }
+
+ fn is_cancelled(&self) -> bool {
+ self.cancelled.load(Ordering::SeqCst)
+ }
+
+ /// Try to consume one request slot. Returns true if successful.
+ fn try_consume_request(&self) -> bool {
+ loop {
+ let current = self.requested.load(Ordering::SeqCst);
+ if current == 0 {
+ return false;
+ }
+ if self
+ .requested
+ .compare_exchange(current, current - 1, Ordering::SeqCst, Ordering::SeqCst)
+ .is_ok()
+ {
+ return true;
+ }
+ }
+ }
+
+ fn has_requests(&self) -> bool {
+ self.requested.load(Ordering::SeqCst) > 0
+ }
+
+ fn wake(&self) {
+ self.waker.wake();
+ }
+}
+
+/// Drives a [`SendableRecordBatchStream`] and pushes results to an [`FFI_ArrowAsyncDeviceStreamHandler`].
+///
+/// This function implements the full Arrow Async Device Stream producer protocol:
+/// - Creates an `FFI_ArrowAsyncProducer` and sets it on the handler before calling `on_schema`
+/// - Respects back-pressure: waits for consumer to call `producer.request(n)` before sending batches
+/// - Handles consumer cancellation via `producer.cancel()`
+/// - Signals end-of-stream by calling `on_next_task` with `NULL`
+/// - Calls `on_error` if the producer encounters an error
+///
+/// # Safety
+///
+/// A non-null `handler` must point to a properly initialized
+/// [`FFI_ArrowAsyncDeviceStreamHandler`] that remains valid until its release
+/// callback is invoked. This function takes exclusive responsibility for
+/// invoking and releasing the handler, including if the returned future is
+/// dropped without being polled. The caller must not release it independently.
+/// Its callbacks and private state must support serialized calls from any
+/// thread, since the returned future is `Send`.
+///
+/// # Example
+///
+/// ```ignore
+/// use sedona_extension::export_sendable_record_batch_stream::drive_stream_to_handler;
+///
+/// // The consumer should call handler.producer.request(n) to receive batches
+/// unsafe { drive_stream_to_handler(stream, handler) }.await;
+/// ```
+///
+/// Raw handler ownership must be explicitly acknowledged by the caller:
+///
+/// ```compile_fail,E0133
+/// use sedona_extension::export_sendable_record_batch_stream::drive_stream_to_handler;
+/// use sedona_extension::extension::FFI_ArrowAsyncDeviceStreamHandler;
+/// use datafusion_execution::SendableRecordBatchStream;
+/// fn start(stream: SendableRecordBatchStream, handler: *mut FFI_ArrowAsyncDeviceStreamHandler) {
+/// let _driver = drive_stream_to_handler(stream, handler);
+/// }
+/// ```
+pub unsafe fn drive_stream_to_handler(
+ stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+) -> impl Future<Output = ()> + Send {
+ // Construct the handler guard before creating the future so dropping an
+ // unpolled spawned task still releases the handler.
+ let handler = SendableHandlerGuard::new(handler);
+
+ async move {
+ let Some(handler) = handler else {
+ return;
+ };
+
+ let state = Arc::new(ProducerState::new());
+ let producer = SendableProducerGuard::new(create_producer(Arc::clone(&state)));
+ let mut resources = SendableDriverResources { handler, producer };
+
+ let _result = drive_stream_inner(
+ stream,
+ resources.handler_address(),
+ resources.producer_address(),
+ state,
+ )
+ .await;
+
+ // Normal completion (success or callback rejection).
+ resources.release_handler();
+ }
+}
+
+/// Inner implementation that returns a Result for cleaner control flow.
+async fn drive_stream_inner(
+ mut stream: SendableRecordBatchStream,
+ handler_address: usize,
+ producer_address: usize,
+ state: Arc<ProducerState>,
+) -> Result<(), ()> {
+ // Set the producer on the handler BEFORE calling on_schema (per Arrow spec)
+ unsafe {
+ (*(handler_address as *mut FFI_ArrowAsyncDeviceStreamHandler)).producer =
+ producer_address as *mut FFI_ArrowAsyncProducer;
+ }
+
+ // Send the schema
+ let schema = stream.schema();
+ let mut ffi_schema = match FFI_ArrowSchema::try_from(schema.as_ref()) {
+ Ok(s) => s,
+ Err(e) => {
+ call_on_error(handler_address, 1, &e.to_string());
+ return Err(());
+ }
+ };
+
+ let on_schema =
+ unsafe { (*(handler_address as *mut FFI_ArrowAsyncDeviceStreamHandler)).on_schema };
+ if let Some(on_schema) = on_schema {
+ let ret = unsafe {
+ on_schema(
+ handler_address as *mut FFI_ArrowAsyncDeviceStreamHandler,
+ &mut ffi_schema,
+ )
+ };
+ // Ownership of the schema contents transfers to the handler when the
+ // callback is invoked, regardless of its return value.
+ std::mem::forget(ffi_schema);
+ if ret != 0 {
+ // A rejected callback may only be followed by release.
+ return Err(());
+ }
+ }
+
+ // Stream batches with back-pressure
+ loop {
+ // Check for cancellation
+ if state.is_cancelled() {
+ // Per Arrow spec: successful cancel should NOT call on_error,
+ // just signal end of stream and release
+ signal_end_of_stream(handler_address);
+ return Ok(());
+ }
+
+ // Wait for consumer to request batches (back-pressure)
+ // Register before checking the state so a concurrent request cannot be
+ // lost between the state check and waker registration.
+ futures::future::poll_fn(|cx| {
+ state.waker.register(cx.waker());
+ if state.has_requests() || state.is_cancelled() {
+ Poll::Ready(())
+ } else {
+ Poll::Pending
+ }
+ })
+ .await;
+
+ // Re-check cancellation after waiting
+ if state.is_cancelled() {
+ signal_end_of_stream(handler_address);
+ return Ok(());
+ }
+
+ // Consume one request slot
+ if !state.try_consume_request() {
+ continue;
+ }
+
+ // Poll cancellation and the source together. A source is allowed to
+ // remain pending indefinitely, so cancellation must wake and interrupt
+ // this wait instead of only being observed between batches.
+ let next = futures::future::poll_fn(|cx| {
+ state.waker.register(cx.waker());
+ if state.is_cancelled() {
+ return Poll::Ready(Err(()));
+ }
+
+ match stream.poll_next_unpin(cx) {
+ Poll::Ready(batch) => Poll::Ready(Ok(batch)),
+ Poll::Pending if state.is_cancelled() => Poll::Ready(Err(())),
+ Poll::Pending => Poll::Pending,
+ }
+ })
+ .await;
+
+ let next = match next {
+ Ok(next) => next,
+ Err(()) => {
+ signal_end_of_stream(handler_address);
+ return Ok(());
+ }
+ };
+
+ match next {
+ Some(Ok(batch)) => {
+ match send_batch_to_handler(handler_address, batch) {
+ Ok(()) => {}
+ Err(SendBatchError::HandlerRejected) => {
+ // A rejected callback may only be followed by release.
+ return Err(());
+ }
+ Err(SendBatchError::Producer(error)) => {
+ call_on_error(handler_address, 1, &error);
+ return Err(());
+ }
+ }
+ // Yield to allow consumer to process the batch.
+ // This enables single-task usage (producer and consumer in same task).
+ yield_once().await;
+ }
+ Some(Err(e)) => {
+ call_on_error(handler_address, 1, &e.to_string());
+ return Err(());
+ }
+ None => {
+ // End of stream
+ signal_end_of_stream(handler_address);
+ return Ok(());
+ }
+ }
+ }
+}
+
+/// Create an FFI producer with request/cancel callbacks.
+fn create_producer(state: Arc<ProducerState>) -> FFI_ArrowAsyncProducer {
+ // Convert Arc to raw pointer - this increments the refcount
+ let private_data = Arc::into_raw(state) as *mut c_void;
+
+ FFI_ArrowAsyncProducer {
+ device_type: ARROW_DEVICE_CPU,
+ request: Some(producer_request),
+ cancel: Some(producer_cancel),
+ additional_metadata: std::ptr::null(),
+ private_data,
+ }
+}
+
+unsafe extern "C" fn producer_request(self_: *mut FFI_ArrowAsyncProducer, n: i64) {
+ if self_.is_null() {
+ return;
+ }
+ let producer = &*self_;
+ if producer.private_data.is_null() {
+ return;
+ }
+ let state = &*(producer.private_data as *const ProducerState);
+ if n > 0 {
+ state.request(n as u64);
+ }
+}
+
+unsafe extern "C" fn producer_cancel(self_: *mut FFI_ArrowAsyncProducer) {
+ if self_.is_null() {
+ return;
+ }
+ let producer = &*self_;
+ if producer.private_data.is_null() {
+ return;
+ }
+ let state = &*(producer.private_data as *const ProducerState);
+ state.cancel();
+}
+
+unsafe extern "C" fn producer_release(self_: *mut FFI_ArrowAsyncProducer) {
+ if self_.is_null() {
+ return;
+ }
+ let producer = &mut *self_;
+ if producer.private_data.is_null() {
+ return;
+ }
+ // Drop the Arc<ProducerState>
+ let _ = Arc::from_raw(producer.private_data as *const ProducerState);
+ producer.private_data = null_mut();
+}
+
+/// Owns this implementation's producer while allowing the async driver to move
+/// between Tokio worker threads. The box keeps the pointer exposed to the
+/// consumer stable even when this wrapper moves.
+struct SendableProducerGuard(Box<FFI_ArrowAsyncProducer>);
+
+impl SendableProducerGuard {
+ fn new(producer: FFI_ArrowAsyncProducer) -> Self {
+ Self(Box::new(producer))
+ }
+
+ fn address(&mut self) -> usize {
+ self.0.as_mut() as *mut FFI_ArrowAsyncProducer as usize
+ }
+}
+
+// SAFETY: This wrapper is only constructed for the producer defined in this
+// module. Its private state is an Arc containing atomics and an AtomicWaker,
+// its callbacks are thread-safe, and the boxed producer has a stable address.
+unsafe impl Send for SendableProducerGuard {}
+
+impl Drop for SendableProducerGuard {
+ fn drop(&mut self) {
+ unsafe { producer_release(self.0.as_mut()) };
+ }
+}
+
+/// Owns the producer-side right to invoke and release one foreign handler.
+/// Dropping an armed guard reports cancellation and releases it exactly once.
+struct SendableHandlerGuard {
+ handler: NonNull<FFI_ArrowAsyncDeviceStreamHandler>,
+ armed: bool,
+}
+
+impl SendableHandlerGuard {
+ fn new(handler: *mut FFI_ArrowAsyncDeviceStreamHandler) -> Option<Self> {
+ NonNull::new(handler).map(|handler| Self {
+ handler,
+ armed: true,
+ })
+ }
+
+ fn address(&self) -> usize {
+ self.handler.as_ptr() as usize
+ }
+
+ fn release(&mut self) {
+ if !self.armed {
+ return;
+ }
+
+ // Disarm before invoking foreign code: after release returns the
+ // handler may already have been freed.
+ self.armed = false;
+ let handler = self.handler.as_ptr();
+ let release = unsafe { (*handler).release };
+ if let Some(release) = release {
+ unsafe { release(handler) };
+ }
+ }
+}
+
+// SAFETY: The guard is uniquely owned by the producer task. Arrow requires
+// handler callbacks to be serialized but permits consecutive callbacks to run
+// on different threads. No reference into the handler is retained across an
+// await or a thread migration.
+unsafe impl Send for SendableHandlerGuard {}
+
+impl Drop for SendableHandlerGuard {
+ fn drop(&mut self) {
+ if self.armed {
+ call_on_error(
+ self.address(),
+ libc::ECANCELED,
+ "Stream driver dropped unexpectedly",
+ );
+ self.release();
+ }
+ }
+}
+
+/// Keeps the handler alive through its release callback, then drops the
+/// producer. Field order is significant because Rust drops fields in
+/// declaration order.
+struct SendableDriverResources {
+ handler: SendableHandlerGuard,
+ producer: SendableProducerGuard,
+}
+
+impl SendableDriverResources {
+ fn handler_address(&self) -> usize {
+ self.handler.address()
+ }
+
+ fn producer_address(&mut self) -> usize {
+ self.producer.address()
+ }
+
+ fn release_handler(&mut self) {
+ self.handler.release();
+ }
+}
+
+/// Signal end of stream by calling on_next_task with NULL.
+fn signal_end_of_stream(handler_address: usize) {
+ let handler = handler_address as *mut FFI_ArrowAsyncDeviceStreamHandler;
+ let on_next_task = unsafe { (*handler).on_next_task };
+ if let Some(on_next_task) = on_next_task {
+ // Per Arrow spec: pass NULL task pointer to signal end of stream
+ unsafe {
+ on_next_task(handler, std::ptr::null_mut(), std::ptr::null());
+ }
+ }
+}
+
+enum SendBatchError {
+ HandlerRejected,
+ Producer(String),
+}
+
+fn send_batch_to_handler(handler_address: usize, batch: RecordBatch) -> Result<(), SendBatchError> {
+ let handler = handler_address as *mut FFI_ArrowAsyncDeviceStreamHandler;
+ let on_next_task = unsafe { (*handler).on_next_task };
+ let Some(on_next_task) = on_next_task else {
+ return Err(SendBatchError::Producer(
+ "on_next_task callback is null".to_string(),
+ ));
+ };
+
+ // Create a task that wraps the batch
+ // The handler moves the task contents if it accepts the callback. Keep the
+ // stack wrapper from dropping the moved private data in that case.
+ let mut task = create_async_task(batch).map_err(SendBatchError::Producer)?;
+ let ret = unsafe { on_next_task(handler, &mut task, std::ptr::null()) };
+ if ret != 0 {
+ return Err(SendBatchError::HandlerRejected);
+ }
+ std::mem::forget(task);
+ Ok(())
+}
+
+fn call_on_error(handler_address: usize, code: c_int, message: &str) {
+ let handler = handler_address as *mut FFI_ArrowAsyncDeviceStreamHandler;
+ let on_error = unsafe { (*handler).on_error };
+ if let Some(on_error) = on_error {
+ let c_message = CString::new(message).unwrap_or_else(|_| CString::new("error").unwrap());
+ unsafe {
+ on_error(
+ handler,
+ code,
+ c_message.as_ptr(),
+ std::ptr::null(), // no metadata
+ );
+ }
+ }
+}
+
+/// Private data for an async task wrapping a RecordBatch
+struct AsyncTaskPrivate {
+ batch: Option<RecordBatch>,
+}
+
+fn create_async_task(batch: RecordBatch) -> Result<FFI_ArrowAsyncTask, String> {
+ let private_data = Box::new(AsyncTaskPrivate { batch: Some(batch) });
+
+ Ok(FFI_ArrowAsyncTask {
+ extract_data: Some(async_task_extract_data),
+ private_data: Box::into_raw(private_data) as *mut c_void,
+ })
+}
+
+unsafe extern "C" fn async_task_extract_data(
+ self_: *mut FFI_ArrowAsyncTask,
+ out: *mut FFI_ArrowDeviceArray,
+) -> c_int {
+ if self_.is_null() {
+ return 1;
+ }
+
+ let task = &mut *self_;
+ if task.private_data.is_null() {
+ return 1;
+ }
+
+ // extract_data is the task's one-shot cleanup callback. Take ownership of
+ // the private allocation before doing any work so all exit paths release it.
+ let mut private = Box::from_raw(task.private_data as *mut AsyncTaskPrivate);
+ task.private_data = null_mut();
+ task.extract_data = None;
+
+ // A null output means the consumer is abandoning the task.
+ if out.is_null() {
+ return 0;
+ }
+
+ // Take the batch (can only extract once)
+ let Some(batch) = private.batch.take() else {
+ return 1;
+ };
+
+ // Convert to struct array and then to FFI
+ let struct_array: arrow_array::StructArray = batch.into();
+ let (ffi_array, _ffi_schema) = match to_ffi(&struct_array.into()) {
+ Ok(result) => result,
+ Err(_) => return 1,
+ };
+
+ // Create device array (CPU)
+ let device_array = FFI_ArrowDeviceArray::from(ffi_array);
+ std::ptr::write(out, device_array);
+ 0
+}
+
+impl Drop for FFI_ArrowAsyncTask {
+ fn drop(&mut self) {
+ // ArrowAsyncTask private_data is producer-owned. Consumers must use
+ // extract_data(NULL), even when abandoning the task, so foreign tasks
+ // are cleaned up by the producer that created them.
+ if let Some(extract_data) = self.extract_data.take() {
+ unsafe {
+ extract_data(self, null_mut());
+ }
+ }
+ }
+}
diff --git a/c/sedona-extension/src/extension.rs b/c/sedona-extension/src/extension.rs
index fce3a7d35..e941bd64a 100644
--- a/c/sedona-extension/src/extension.rs
+++ b/c/sedona-extension/src/extension.rs
@@ -25,6 +25,104 @@ use arrow_array::{
ffi::{FFI_ArrowArray, FFI_ArrowSchema},
ffi_stream::FFI_ArrowArrayStream,
};
+use arrow_schema::ArrowError;
+
+/// CPU device identifier from the Arrow C Device Data Interface.
+pub const ARROW_DEVICE_CPU: i32 = 1;
+
+/// Rust representation of `ArrowDeviceArray` used by the experimental Arrow
+/// asynchronous device stream interface.
+#[repr(C)]
+pub struct FFI_ArrowDeviceArray {
+ pub array: FFI_ArrowArray,
+ pub device_id: i64,
+ pub device_type: i32,
+ pub sync_event: *mut c_void,
+ pub reserved: [i64; 3],
+}
+
+impl TryFrom<FFI_ArrowDeviceArray> for FFI_ArrowArray {
+ type Error = ArrowError;
+
+ fn try_from(value: FFI_ArrowDeviceArray) -> Result<Self, Self::Error> {
+ if value.device_type != ARROW_DEVICE_CPU {
+ return Err(ArrowError::CDataInterface(format!(
+ "Unsupported Arrow device type: {}",
+ value.device_type
+ )));
+ }
+ if !value.sync_event.is_null() {
+ return Err(ArrowError::CDataInterface(
+ "CPU ArrowDeviceArray has a non-null sync event".to_string(),
+ ));
+ }
+ Ok(value.array)
+ }
+}
+
+impl From<FFI_ArrowArray> for FFI_ArrowDeviceArray {
+ fn from(array: FFI_ArrowArray) -> Self {
+ Self {
+ array,
+ device_id: -1,
+ device_type: ARROW_DEVICE_CPU,
+ sync_event: null_mut(),
+ reserved: [0; 3],
+ }
+ }
+}
+
+/// Rust representation of `ArrowAsyncProducer`.
+#[repr(C)]
+pub struct FFI_ArrowAsyncProducer {
+ pub device_type: i32,
+ pub request: Option<unsafe extern "C" fn(*mut FFI_ArrowAsyncProducer, i64)>,
+ pub cancel: Option<unsafe extern "C" fn(*mut FFI_ArrowAsyncProducer)>,
+ pub additional_metadata: *const c_char,
+ pub private_data: *mut c_void,
+}
+
+/// Rust representation of `ArrowAsyncDeviceStreamHandler`.
+#[repr(C)]
+pub struct FFI_ArrowAsyncDeviceStreamHandler {
+ pub on_schema: Option<
+ unsafe extern "C" fn(*mut FFI_ArrowAsyncDeviceStreamHandler, *mut FFI_ArrowSchema) -> c_int,
+ >,
+ pub on_next_task: Option<
+ unsafe extern "C" fn(
+ *mut FFI_ArrowAsyncDeviceStreamHandler,
+ *mut FFI_ArrowAsyncTask,
+ *const c_char,
+ ) -> c_int,
+ >,
+ pub on_error: Option<
+ unsafe extern "C" fn(
+ *mut FFI_ArrowAsyncDeviceStreamHandler,
+ c_int,
+ *const c_char,
+ *const c_char,
+ ),
+ >,
+ pub release: Option<unsafe extern "C" fn(*mut FFI_ArrowAsyncDeviceStreamHandler)>,
+ pub producer: *mut FFI_ArrowAsyncProducer,
+ pub private_data: *mut c_void,
+}
+
+impl Drop for FFI_ArrowAsyncDeviceStreamHandler {
+ fn drop(&mut self) {
+ if let Some(release) = self.release.take() {
+ unsafe { release(self) };
+ }
+ }
+}
+
+/// Rust representation of `ArrowAsyncTask`.
+#[repr(C)]
+pub struct FFI_ArrowAsyncTask {
+ pub extract_data:
+ Option<unsafe extern "C" fn(*mut FFI_ArrowAsyncTask, *mut FFI_ArrowDeviceArray) -> c_int>,
+ pub private_data: *mut c_void,
+}
/// Raw FFI representation of the SedonaCScalarKernel
///
diff --git a/c/sedona-extension/src/import_sendable_record_batch_stream.rs b/c/sedona-extension/src/import_sendable_record_batch_stream.rs
new file mode 100644
index 000000000..ce66e0c1f
--- /dev/null
+++ b/c/sedona-extension/src/import_sendable_record_batch_stream.rs
@@ -0,0 +1,1027 @@
+// 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.
+
+//! Import an async device stream from FFI into a [`SendableRecordBatchStream`].
+//!
+//! This module allows Rust code to consume async record batch streams provided by
+//! C/FFI producers via the Arrow C Device Data Interface.
+//!
+//! # Usage
+//!
+//! 1. Create an [`ImportedAsyncDeviceStream`] which provides an [`FFI_ArrowAsyncDeviceStreamHandler`]
+//! 2. Pass the handler pointer to the FFI producer
+//! 3. Use the stream as a normal [`SendableRecordBatchStream`]
+//!
+//! ```ignore
+//! use sedona_extension::import_sendable_record_batch_stream::ImportedAsyncDeviceStream;
+//!
+//! let (stream, handler) = ImportedAsyncDeviceStream::new(16);
+//!
+//! // Pass handler to FFI producer...
+//! // ffi_producer_start(handler.as_ptr());
+//!
+//! // Consume the stream
+//! while let Some(batch) = stream.next().await {
+//! let batch = batch?;
+//! // process batch
+//! }
+//! ```
+
+use std::ffi::{c_int, c_void, CStr};
+use std::pin::Pin;
+use std::ptr::null_mut;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll};
+
+use arrow_array::ffi::{from_ffi_and_data_type, FFI_ArrowArray, FFI_ArrowSchema};
+use arrow_array::{Array, RecordBatch, RecordBatchOptions, StructArray};
+use arrow_schema::{ArrowError, DataType, Schema, SchemaRef};
+use datafusion_common::Result;
+use datafusion_execution::RecordBatchStream;
+use futures::channel::mpsc;
+use futures::Stream;
+
+use crate::extension::{
+ FFI_ArrowAsyncDeviceStreamHandler, FFI_ArrowAsyncProducer, FFI_ArrowAsyncTask,
+ FFI_ArrowDeviceArray, ARROW_DEVICE_CPU,
+};
+
+/// Messages sent from FFI callbacks to the stream.
+enum StreamMessage {
+ /// Schema received from producer.
+ Schema(SchemaRef),
+ /// A record batch task received.
+ Task(FFI_ArrowAsyncTask),
+ /// Error from producer.
+ Error(ArrowError),
+}
+
+/// State shared between callbacks and stream for producer management.
+/// This is the only state that requires mutex protection.
+struct ProducerState {
+ /// Producer pointer for requesting more data.
+ producer: *mut FFI_ArrowAsyncProducer,
+ /// Number of batches to request at a time (for back-pressure).
+ prefetch_count: u64,
+ /// Remember cancellation even before the producer connects.
+ cancelled: bool,
+}
+
+// Safety: Producer pointer is only dereferenced while holding the lock
+unsafe impl Send for ProducerState {}
+unsafe impl Sync for ProducerState {}
+
+/// Shared state between the handler callbacks and the stream.
+struct ImportedStreamState {
+ /// Sender for messages from callbacks to stream (lock-free).
+ sender: mpsc::UnboundedSender<StreamMessage>,
+ /// Number of outstanding requests (for back-pressure).
+ pending_requests: AtomicU64,
+ /// Set to true when the stream is dropped but producer hasn't called release yet.
+ abandoned: AtomicBool,
+ /// Set to true when handler_release is called (handler was freed by producer).
+ handler_released: AtomicBool,
+ /// Set to true when a producer connects (calls on_schema).
+ producer_connected: AtomicBool,
+ /// Producer state (needs mutex for FFI calls).
+ producer_state: Mutex<ProducerState>,
+}
+
+// Safety: All fields are Send+Sync (atomics, Mutex, channel sender).
+// The producer pointer inside ProducerState is protected by the Mutex.
+unsafe impl Send for ImportedStreamState {}
+unsafe impl Sync for ImportedStreamState {}
+
+impl ImportedStreamState {
+ fn new(prefetch_count: u64, sender: mpsc::UnboundedSender<StreamMessage>) -> Self {
+ Self {
+ sender,
+ pending_requests: AtomicU64::new(0),
+ abandoned: AtomicBool::new(false),
+ handler_released: AtomicBool::new(false),
+ producer_connected: AtomicBool::new(false),
+ producer_state: Mutex::new(ProducerState {
+ producer: null_mut(),
+ prefetch_count,
+ cancelled: false,
+ }),
+ }
+ }
+
+ /// Request more data from the producer if needed.
+ fn maybe_request_more(&self) {
+ // Lock cannot be poisoned: we never panic while holding it
+ let producer_state = self
+ .producer_state
+ .lock()
+ .expect("producer_state mutex poisoned");
+ if producer_state.producer.is_null() || producer_state.cancelled || self.sender.is_closed()
+ {
+ return;
+ }
+
+ let pending = self.pending_requests.load(Ordering::Acquire);
+ let prefetch = producer_state.prefetch_count;
+
+ // Request more when we're running low. Round the threshold up so a
+ // prefetch count of one still issues its initial request.
+ let refill_threshold = prefetch / 2 + prefetch % 2;
+ if pending < refill_threshold {
+ let to_request = prefetch - pending;
+ if let Some(request_fn) = unsafe { (*producer_state.producer).request } {
+ unsafe { request_fn(producer_state.producer, to_request as i64) };
+ self.pending_requests
+ .fetch_add(to_request, Ordering::Release);
+ }
+ }
+ }
+
+ /// Returns false if cancellation preceded the producer's connection.
+ fn set_producer(&self, producer: *mut FFI_ArrowAsyncProducer) -> bool {
+ // Lock cannot be poisoned: we never panic while holding it
+ let mut state = self
+ .producer_state
+ .lock()
+ .expect("producer_state mutex poisoned");
+ self.producer_connected.store(true, Ordering::Release);
+ if state.cancelled {
+ return false;
+ }
+ state.producer = producer;
+ true
+ }
+
+ fn clear_producer(&self) {
+ // Lock cannot be poisoned: we never panic while holding it
+ let mut state = self
+ .producer_state
+ .lock()
+ .expect("producer_state mutex poisoned");
+ state.producer = null_mut();
+ }
+
+ fn cancel(&self) {
+ // Lock cannot be poisoned: we never panic while holding it
+ let mut state = self
+ .producer_state
+ .lock()
+ .expect("producer_state mutex poisoned");
+ if state.cancelled {
+ return;
+ }
+ state.cancelled = true;
+ if !state.producer.is_null() {
+ if let Some(cancel_fn) = unsafe { (*state.producer).cancel } {
+ unsafe { cancel_fn(state.producer) };
+ }
+ }
+ }
+}
+
+/// A [`RecordBatchStream`] that consumes data from an FFI async producer.
+///
+/// Create with [`ImportedAsyncDeviceStream::new`], then pass the handler pointer
+/// to the FFI producer. The stream will yield batches as the producer sends them.
+///
+/// # Dropping
+///
+/// When the stream is dropped, it signals cancellation to the producer. The handler
+/// remains valid until the producer calls its `release` callback. This allows the
+/// idiomatic DataFusion pattern of dropping streams to cancel queries.
+///
+/// # Performance
+///
+/// Uses a lock-free channel for message passing from FFI callbacks to the stream.
+/// Only acquires a mutex when requesting more batches from the producer (approximately
+/// once per `prefetch_count / 2` batches).
+pub struct ImportedAsyncDeviceStream {
+ state: Arc<ImportedStreamState>,
+ /// Receiver for messages from callbacks (lock-free).
+ receiver: mpsc::UnboundedReceiver<StreamMessage>,
+ /// Cached schema for the RecordBatchStream trait.
+ schema: Option<SchemaRef>,
+ /// DataType for converting FFI arrays.
+ schema_struct_type: Option<DataType>,
+}
+
+// Safety: The state uses atomic operations and mutex-protected producer access.
+unsafe impl Send for ImportedAsyncDeviceStream {}
+
+/// RAII wrapper for the FFI handler.
+///
+/// This wrapper ensures the handler is properly cleaned up even if the FFI producer
+/// never calls `release()`. It provides safe access to the raw pointer for FFI calls.
+pub struct AsyncDeviceStreamHandler {
+ /// Raw pointer to the handler (heap-allocated).
+ ptr: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ /// Shared state with the stream - needed to check if already released.
+ state: Arc<ImportedStreamState>,
+}
+
+// Safety: The handler pointer is only accessed from one thread at a time.
+// The state Arc provides synchronization.
+unsafe impl Send for AsyncDeviceStreamHandler {}
+
+impl AsyncDeviceStreamHandler {
+ /// Get the raw pointer to pass to FFI code.
+ ///
+ /// The pointer remains valid until either:
+ /// - The FFI producer calls the handler's `release` callback, or
+ /// - This wrapper is dropped
+ #[inline]
+ pub fn as_ptr(&self) -> *mut FFI_ArrowAsyncDeviceStreamHandler {
+ self.ptr
+ }
+
+ /// Transfer responsibility for releasing this handler to the producer.
+ pub fn into_raw(mut self) -> *mut FFI_ArrowAsyncDeviceStreamHandler {
+ let ptr = self.ptr;
+ self.ptr = null_mut();
+ ptr
+ }
+}
+
+impl Drop for AsyncDeviceStreamHandler {
+ fn drop(&mut self) {
+ if self.ptr.is_null() {
+ return;
+ }
+
+ // Check via shared state if the handler was already released by the producer.
+ // We MUST check this before touching self.ptr because handler_release frees it.
+ if self.state.handler_released.load(Ordering::Acquire) {
+ // Producer already called release and freed the handler
+ return;
+ }
+
+ // If a producer has connected, it will call release when done.
+ // We must NOT free the handler or we'll cause a use-after-free.
+ if self.state.producer_connected.load(Ordering::Acquire) {
+ return;
+ }
+
+ // No producer ever connected - clean it up ourselves.
+ // This can happen if the handler was never passed to FFI code.
+ //
+ // We need to:
+ // 1. Mark the stream as ended
+ // 2. Drop the Arc reference in private_data
+ // 3. Free the handler struct
+
+ let handler = unsafe { &mut *self.ptr };
+
+ self.state.sender.close_channel();
+
+ if !handler.private_data.is_null() {
+ let state_ptr = handler.private_data as *const ImportedStreamState;
+ let _ = unsafe { Arc::from_raw(state_ptr) };
+ handler.private_data = null_mut();
+ }
+
+ // Clear release to prevent FFI_ArrowAsyncDeviceStreamHandler's Drop
+ // from calling it (which would double-free)
+ handler.release = None;
+
+ let _ = unsafe { Box::from_raw(self.ptr) };
+ }
+}
+
+impl ImportedAsyncDeviceStream {
+ /// Create a new imported stream.
+ ///
+ /// Returns the stream and an RAII handler wrapper that should be used to pass
+ /// the handler pointer to the FFI producer.
+ ///
+ /// # Arguments
+ ///
+ /// * `prefetch_count` - Number of batches to request ahead for back-pressure.
+ /// A larger value reduces latency but uses more memory.
+ pub fn new(prefetch_count: u64) -> (Self, AsyncDeviceStreamHandler) {
+ Self::new_inner(prefetch_count, None)
+ }
+
+ /// Create a stream whose schema is known by the caller before the producer
+ /// invokes `on_schema`.
+ ///
+ /// DataFusion requires [`RecordBatchStream::schema`] to be available as soon
+ /// as `ExecutionPlan::execute` returns, so execution-plan imports use this
+ /// constructor with the schema obtained from `SedonaCExecutionPlan`.
+ pub fn new_with_schema(
+ prefetch_count: u64,
+ schema: SchemaRef,
+ ) -> (Self, AsyncDeviceStreamHandler) {
+ Self::new_inner(prefetch_count, Some(schema))
+ }
+
+ fn new_inner(
+ prefetch_count: u64,
+ schema: Option<SchemaRef>,
+ ) -> (Self, AsyncDeviceStreamHandler) {
+ // A zero-sized window can never make progress; treat it as the minimum
+ // useful prefetch size.
+ let prefetch_count = prefetch_count.max(1);
+ let (sender, receiver) = mpsc::unbounded();
+ let state = Arc::new(ImportedStreamState::new(prefetch_count, sender));
+
+ // Handler is heap-allocated and will be freed when:
+ // 1. Producer calls release(), or
+ // 2. AsyncDeviceStreamHandler is dropped (fallback cleanup)
+ let handler_ptr = Box::into_raw(Box::new(FFI_ArrowAsyncDeviceStreamHandler {
+ on_schema: Some(handler_on_schema),
+ on_next_task: Some(handler_on_next_task),
+ on_error: Some(handler_on_error),
+ release: Some(handler_release),
+ producer: null_mut(),
+ private_data: Arc::into_raw(Arc::clone(&state)) as *mut c_void,
+ }));
+
+ let stream = Self {
+ state: Arc::clone(&state),
+ receiver,
+ schema_struct_type: schema
+ .as_ref()
+ .map(|schema| DataType::Struct(schema.fields().iter().cloned().collect())),
+ schema,
+ };
+
+ let handler = AsyncDeviceStreamHandler {
+ ptr: handler_ptr,
+ state,
+ };
+
+ (stream, handler)
+ }
+
+ /// Convert this stream into a [`SendableRecordBatchStream`].
+ ///
+ /// This is a convenience method equivalent to `Box::pin(stream)`.
+ pub fn into_sendable(self) -> datafusion_execution::SendableRecordBatchStream {
+ Box::pin(self)
+ }
+
+ /// Cancel the stream, signaling to the producer to stop sending data.
+ /// If the producer has not connected yet, its schema callback is rejected.
+ pub fn cancel(&self) {
+ self.state.cancel();
+ }
+
+ /// Poll for the next message, handling schema initialization.
+ fn poll_next_inner(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<RecordBatch>>> {
+ // Channel closure both wakes the receiver and preserves queued messages.
+ // A separate end flag could race with this poll and hide a queued error
+ // or batch behind a premature end-of-stream result.
+ match Pin::new(&mut self.receiver).poll_next(cx) {
+ Poll::Ready(Some(msg)) => match msg {
+ StreamMessage::Schema(schema) => {
+ self.schema_struct_type =
+ Some(DataType::Struct(schema.fields().iter().cloned().collect()));
+ self.schema = Some(schema);
+
+ // Request initial batch of data (acquires lock)
+ self.state.maybe_request_more();
+
+ // Continue polling for actual data
+ self.poll_next_inner(cx)
+ }
+ StreamMessage::Task(mut task) => {
+ // Decrement pending (lock-free)
+ self.state.pending_requests.fetch_sub(1, Ordering::Release);
+ // Maybe request more (acquires lock only if needed)
+ self.state.maybe_request_more();
+
+ // Extract data from the task
+ let batch = self.extract_batch_from_task(&mut task);
+ Poll::Ready(Some(batch))
+ }
+ StreamMessage::Error(e) => Poll::Ready(Some(Err(e.into()))),
+ },
+ Poll::Ready(None) => Poll::Ready(None),
+ Poll::Pending => Poll::Pending,
+ }
+ }
+
+ /// Extract a RecordBatch from an FFI task.
+ fn extract_batch_from_task(&self, task: &mut FFI_ArrowAsyncTask) -> Result<RecordBatch> {
+ let Some(extract_data) = task.extract_data else {
+ return Err(ArrowError::CDataInterface("extract_data is null".to_string()).into());
+ };
+
+ let mut device_array = FFI_ArrowDeviceArray {
+ array: FFI_ArrowArray::empty(),
+ device_id: 0,
+ device_type: 0,
+ sync_event: null_mut(),
+ reserved: [0; 3],
+ };
+
+ let ret = unsafe { extract_data(task, &mut device_array) };
+ // Extraction consumes the task even on error. Foreign producers are
+ // not required to clear the callback or private_data themselves.
+ task.extract_data = None;
+ if ret != 0 {
+ return Err(ArrowError::CDataInterface("extract_data failed".to_string()).into());
+ }
+
+ // Only CPU device supported
+ if device_array.device_type != ARROW_DEVICE_CPU {
+ return Err(ArrowError::CDataInterface(format!(
+ "Unsupported device type: {}",
+ device_array.device_type
+ ))
+ .into());
+ }
+
+ // Convert to RecordBatch
+ let Some(ref struct_type) = self.schema_struct_type else {
+ return Err(ArrowError::CDataInterface("Schema not yet received".to_string()).into());
+ };
+
+ let array_data =
+ unsafe { from_ffi_and_data_type(device_array.array, struct_type.clone())? };
+ let struct_array: StructArray = array_data.into();
+ let Some(schema) = self.schema.clone() else {
+ return Err(ArrowError::CDataInterface("Schema not yet received".to_string()).into());
+ };
+ let row_count = struct_array.len();
+ let (_, columns, nulls) = struct_array.into_parts();
+ if nulls.is_some_and(|nulls| nulls.null_count() != 0) {
+ return Err(ArrowError::CDataInterface(
+ "Cannot convert nullable struct array to record batch".to_string(),
+ )
+ .into());
+ }
+
+ let options = RecordBatchOptions::new().with_row_count(Some(row_count));
+ Ok(RecordBatch::try_new_with_options(
+ schema, columns, &options,
+ )?)
+ }
+}
+
+impl Stream for ImportedAsyncDeviceStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.poll_next_inner(cx)
+ }
+}
+
+impl RecordBatchStream for ImportedAsyncDeviceStream {
+ fn schema(&self) -> SchemaRef {
+ self.schema
+ .clone()
+ .unwrap_or_else(|| Arc::new(Schema::empty()))
+ }
+}
+
+impl Drop for ImportedAsyncDeviceStream {
+ fn drop(&mut self) {
+ // Mark as abandoned so callbacks know to return errors (lock-free)
+ self.state.abandoned.store(true, Ordering::Release);
+
+ // Signal cancellation to the producer (acquires lock)
+ self.state.cancel();
+ }
+}
+
+// FFI callback implementations
+
+unsafe extern "C" fn handler_on_schema(
+ self_: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ schema: *mut FFI_ArrowSchema,
+) -> c_int {
+ if schema.is_null() {
+ return 1;
+ }
+
+ // The handler owns the schema contents as soon as this callback is
+ // invoked, including on every error path.
+ let ffi_schema = FFI_ArrowSchema::from_raw(schema);
+
+ if self_.is_null() {
+ return 1;
+ }
+
+ let handler = &mut *self_;
+ if handler.private_data.is_null() {
+ return 1;
+ }
+
+ // Get shared state
+ let state_ptr = handler.private_data as *const ImportedStreamState;
+ let state_arc = Arc::from_raw(state_ptr);
+
+ // If stream was dropped, tell producer to stop (lock-free check)
+ if state_arc.abandoned.load(Ordering::Acquire) {
+ let _ = Arc::into_raw(state_arc);
+ return 1; // Error code signals producer to stop and release
+ }
+
+ // Store the producer pointer (acquires lock)
+ if !state_arc.set_producer(handler.producer) {
+ // Rejecting on_schema obliges the producer to release the handler.
+ // Avoid calling cancel here: a foreign cancel callback may invoke
+ // on_error synchronously and reenter this handler.
+ state_arc.sender.close_channel();
+ let _ = Arc::into_raw(state_arc);
+ return libc::ECANCELED;
+ }
+
+ // Import the schema
+ let result = match Schema::try_from(&ffi_schema) {
+ Ok(s) => {
+ // Send through channel (lock-free)
+ let _ = state_arc
+ .sender
+ .unbounded_send(StreamMessage::Schema(Arc::new(s)));
+ 0
+ }
+ Err(e) => {
+ let _ = state_arc.sender.unbounded_send(StreamMessage::Error(e));
+ state_arc.sender.close_channel();
+ 1
+ }
+ };
+
+ // Don't drop the Arc, just forget about this reference
+ let _ = Arc::into_raw(state_arc);
+ result
+}
+
+unsafe extern "C" fn handler_on_next_task(
+ self_: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ task: *mut FFI_ArrowAsyncTask,
+ _metadata: *const std::ffi::c_char,
+) -> c_int {
+ if self_.is_null() {
+ return 1;
+ }
+
+ let handler = &*self_;
+ if handler.private_data.is_null() {
+ return 1;
+ }
+
+ let state_ptr = handler.private_data as *const ImportedStreamState;
+ let state_arc = Arc::from_raw(state_ptr);
+
+ // If stream was dropped, tell producer to stop (lock-free check)
+ if state_arc.abandoned.load(Ordering::Acquire) {
+ let _ = Arc::into_raw(state_arc);
+ return 1; // Error code signals producer to stop and release
+ }
+
+ // Send message through channel (lock-free)
+ if task.is_null() {
+ // NULL task signals end of stream
+ state_arc.sender.close_channel();
+ } else {
+ // Take ownership of the task by copying it
+ let task_copy = std::ptr::read(task);
+ let _ = state_arc
+ .sender
+ .unbounded_send(StreamMessage::Task(task_copy));
+ }
+
+ let _ = Arc::into_raw(state_arc);
+ 0
+}
+
+unsafe extern "C" fn handler_on_error(
+ self_: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ code: c_int,
+ message: *const std::ffi::c_char,
+ _metadata: *const std::ffi::c_char,
+) {
+ if self_.is_null() {
+ return;
+ }
+
+ let handler = &*self_;
+ if handler.private_data.is_null() {
+ return;
+ }
+
+ let state_ptr = handler.private_data as *const ImportedStreamState;
+ let state_arc = Arc::from_raw(state_ptr);
+
+ let error_msg = if message.is_null() {
+ format!("FFI error code {}", code)
+ } else {
+ let c_str = CStr::from_ptr(message);
+ c_str.to_string_lossy().into_owned()
+ };
+
+ // Send error through channel (lock-free)
+ let _ = state_arc
+ .sender
+ .unbounded_send(StreamMessage::Error(ArrowError::CDataInterface(error_msg)));
+ state_arc.sender.close_channel();
+
+ let _ = Arc::into_raw(state_arc);
+}
+
+unsafe extern "C" fn handler_release(self_: *mut FFI_ArrowAsyncDeviceStreamHandler) {
+ if self_.is_null() {
+ return;
+ }
+
+ let handler = &mut *self_;
+ if handler.private_data.is_null() {
+ // Already released - free the handler struct
+ // Clear release to prevent Drop from calling us again
+ handler.release = None;
+ let _ = Box::from_raw(self_);
+ return;
+ }
+
+ // Get the state and clear the producer pointer so stream Drop won't use it
+ let state_ptr = handler.private_data as *const ImportedStreamState;
+ let state_arc = Arc::from_raw(state_ptr);
+
+ state_arc.clear_producer();
+ state_arc.sender.close_channel();
+
+ // Mark handler as released so AsyncDeviceStreamHandler::drop knows not to free it
+ state_arc.handler_released.store(true, Ordering::Release);
+
+ // Drop the Arc reference (will be freed when stream also drops its reference)
+ drop(state_arc);
+ handler.private_data = null_mut();
+
+ // Clear release to prevent Drop impl from calling us again, then free
+ handler.release = None;
+ let _ = Box::from_raw(self_);
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::export_sendable_record_batch_stream::drive_stream_to_handler;
+ use arrow_array::{Int32Array, StringArray};
+ use arrow_schema::Field;
+ use datafusion_execution::RecordBatchStream;
+ use futures::StreamExt;
+ use std::pin::Pin;
+ use std::sync::Arc;
+
+ /// Create a test schema with id and name columns.
+ fn test_schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("name", DataType::Utf8, true),
+ ]))
+ }
+
+ /// Create a test batch with the given row count.
+ fn make_batch(schema: &SchemaRef, start_id: i32, count: i32) -> RecordBatch {
+ let ids: Vec<i32> = (start_id..start_id + count).collect();
+ let names: Vec<Option<&str>> = ids
+ .iter()
+ .map(|i| Some(if i % 2 == 0 { "even" } else { "odd" }))
+ .collect();
+
+ RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(Int32Array::from(ids)),
+ Arc::new(StringArray::from(names)),
+ ],
+ )
+ .unwrap()
+ }
+
+ /// A simple SendableRecordBatchStream from a list of batches.
+ struct TestStream {
+ schema: SchemaRef,
+ batches: std::collections::VecDeque<Result<RecordBatch>>,
+ }
+
+ impl TestStream {
+ fn new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Pin<Box<Self>> {
+ Box::pin(Self {
+ schema,
+ batches: batches.into_iter().map(Ok).collect(),
+ })
+ }
+
+ fn with_error(
+ schema: SchemaRef,
+ batches: Vec<RecordBatch>,
+ error_after: usize,
+ ) -> Pin<Box<Self>> {
+ let mut results: std::collections::VecDeque<Result<RecordBatch>> =
+ batches.into_iter().map(Ok).collect();
+ if error_after < results.len() {
+ results.insert(
+ error_after,
+ Err(datafusion_common::DataFusionError::External(
+ "Test error".into(),
+ )),
+ );
+ } else {
+ results.push_back(Err(datafusion_common::DataFusionError::External(
+ "Test error".into(),
+ )));
+ }
+ Box::pin(Self {
+ schema,
+ batches: results,
+ })
+ }
+ }
+
+ impl Stream for TestStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ let this = self.as_mut().get_mut();
+ Poll::Ready(this.batches.pop_front())
+ }
+ }
+
+ impl RecordBatchStream for TestStream {
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+ }
+
+ #[tokio::test]
+ async fn test_empty_stream() {
+ let schema = test_schema();
+ let source_stream = TestStream::new(schema.clone(), vec![]);
+
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(4);
+
+ let consumer_future = async {
+ let mut received = vec![];
+ while let Some(result) = consumer.next().await {
+ received.push(result);
+ }
+ (received, consumer.schema())
+ };
+
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler.as_ptr()) };
+
+ let ((received, result_schema), _) = futures::join!(consumer_future, producer_future);
+
+ assert!(received.is_empty());
+ assert_eq!(result_schema.fields().len(), 2);
+ }
+
+ #[tokio::test]
+ async fn test_single_batch() {
+ let schema = test_schema();
+ let batch = make_batch(&schema, 1, 5);
+ let source_stream = TestStream::new(schema.clone(), vec![batch.clone()]);
+
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(4);
+
+ let consumer_future = async {
+ let mut received = vec![];
+ while let Some(result) = consumer.next().await {
+ received.push(result);
+ }
+ received
+ };
+
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler.as_ptr()) };
+
+ let (received, _) = futures::join!(consumer_future, producer_future);
+
+ assert_eq!(received.len(), 1);
+ let received_batch = received.into_iter().next().unwrap().unwrap();
+ assert_eq!(received_batch.num_rows(), 5);
+ assert_eq!(received_batch.schema(), batch.schema());
+
+ // Verify data
+ let ids = received_batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(ids.values(), &[1, 2, 3, 4, 5]);
+ }
+
+ #[tokio::test]
+ async fn test_multiple_batches_with_backpressure() {
+ let schema = test_schema();
+
+ // Create more batches than the prefetch count
+ let prefetch_count = 2u64;
+ let num_batches = 10;
+ let batches: Vec<RecordBatch> = (0..num_batches)
+ .map(|i| make_batch(&schema, i * 3, 3))
+ .collect();
+ let expected_total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
+
+ let source_stream = TestStream::new(schema.clone(), batches);
+
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(prefetch_count);
+
+ let consumer_future = async {
+ let mut received = vec![];
+ while let Some(result) = consumer.next().await {
+ received.push(result);
+ }
+ received
+ };
+
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler.as_ptr()) };
+
+ let (received, _) = futures::join!(consumer_future, producer_future);
+
+ assert_eq!(received.len(), num_batches as usize);
+ let total_rows: usize = received.into_iter().map(|r| r.unwrap().num_rows()).sum();
+ assert_eq!(total_rows, expected_total_rows);
+ }
+
+ #[tokio::test]
+ async fn test_stream_error() {
+ let schema = test_schema();
+ let batches = vec![make_batch(&schema, 1, 3), make_batch(&schema, 4, 3)];
+
+ // Error after first batch
+ let source_stream = TestStream::with_error(schema.clone(), batches, 1);
+
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(4);
+
+ let consumer_future = async {
+ let mut received = vec![];
+ while let Some(result) = consumer.next().await {
+ received.push(result);
+ }
+ received
+ };
+
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler.as_ptr()) };
+
+ let (received, _) = futures::join!(consumer_future, producer_future);
+
+ // First batch should succeed
+ assert!(!received.is_empty());
+ assert!(received[0].is_ok());
+
+ // Should have an error somewhere
+ let has_error = received.iter().any(|r| r.is_err());
+ assert!(has_error, "Expected an error in the stream");
+ }
+
+ #[tokio::test]
+ async fn test_cancellation() {
+ let schema = test_schema();
+
+ // Create a large number of batches
+ let batches: Vec<RecordBatch> = (0..100).map(|i| make_batch(&schema, i * 3, 3)).collect();
+
+ let source_stream = TestStream::new(schema.clone(), batches);
+
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+
+ let consumer_future = async {
+ let mut count = 0;
+ while let Some(result) = consumer.next().await {
+ result.expect("should not error");
+ count += 1;
+ if count >= 3 {
+ // Cancel after receiving 3 batches
+ consumer.cancel();
+ break;
+ }
+ }
+ count
+ };
+
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler.as_ptr()) };
+
+ let (count, _) = futures::join!(consumer_future, producer_future);
+
+ // Should have received at least 3 batches before cancellation
+ assert!(count >= 3);
+ }
+
+ #[tokio::test]
+ async fn test_drop_stream_while_producing() {
+ // Test that dropping the consumer stream while the producer is still producing
+ // doesn't cause crashes or hangs - the producer should stop gracefully.
+ let schema = test_schema();
+ let batches: Vec<RecordBatch> = (0..100).map(|i| make_batch(&schema, i * 3, 3)).collect();
+ let source_stream = TestStream::new(schema.clone(), batches);
+
+ let (consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let handler_ptr = handler.as_ptr();
+
+ // Wrap consumer in Option so we can drop it mid-stream
+ let consumer = std::sync::Arc::new(tokio::sync::Mutex::new(Some(consumer)));
+ let consumer_clone = consumer.clone();
+
+ let consumer_future = async move {
+ let mut count = 0;
+ loop {
+ let mut guard = consumer_clone.lock().await;
+ let stream = guard.as_mut().unwrap();
+ match stream.next().await {
+ Some(result) => {
+ result.expect("should not error before drop");
+ count += 1;
+ if count >= 2 {
+ // Drop the stream after receiving 2 batches
+ drop(guard.take());
+ break;
+ }
+ }
+ None => break,
+ }
+ }
+ count
+ };
+
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler_ptr) };
+
+ // Both futures should complete without panic or hang
+ let (count, _) = futures::join!(consumer_future, producer_future);
+ assert_eq!(count, 2);
+ }
+
+ #[tokio::test]
+ async fn test_drop_handler_while_consuming() {
+ // Test that dropping the handler wrapper while the consumer is still reading
+ // doesn't cause crashes. The handler should be properly freed by the producer's
+ // release callback, and the wrapper's Drop should be a no-op.
+ let schema = test_schema();
+ let batches: Vec<RecordBatch> = (0..5).map(|i| make_batch(&schema, i * 3, 3)).collect();
+ let source_stream = TestStream::new(schema.clone(), batches);
+
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(4);
+ let handler_ptr = handler.as_ptr();
+
+ // Start producer
+ let producer_future = unsafe { drive_stream_to_handler(source_stream, handler_ptr) };
+
+ // Consume all batches, but drop the handler wrapper early
+ let consumer_future = async {
+ let mut received = vec![];
+ // Read one batch first
+ if let Some(result) = consumer.next().await {
+ received.push(result);
+ }
+
+ // Drop the handler wrapper while stream is still active
+ // This should NOT free the handler since producer hasn't called release yet
+ drop(handler);
+
+ // Continue consuming - should work fine
+ while let Some(result) = consumer.next().await {
+ received.push(result);
+ }
+ received
+ };
+
+ let (received, _) = futures::join!(consumer_future, producer_future);
+
+ // All 5 batches should be received
+ assert_eq!(received.len(), 5);
+ assert!(received.iter().all(|r| r.is_ok()));
+ }
+
+ #[tokio::test]
+ async fn test_handler_cleanup_when_producer_never_connects() {
+ // Test RAII cleanup when the handler is dropped but the producer never called
+ // any callbacks (never connected). The handler wrapper should clean up properly.
+ let (_consumer, handler) = ImportedAsyncDeviceStream::new(4);
+
+ // Just drop the handler without ever passing it to a producer
+ // This should NOT panic or leak memory
+ drop(handler);
+
+ // Stream should still be usable (though it will never receive data)
+ // The channel should be closed by handler drop
+ }
+
+ #[tokio::test]
+ async fn test_stream_and_handler_both_dropped_before_producer() {
+ // Test cleanup when both stream and handler are dropped before any producer connects
+ let (consumer, handler) = ImportedAsyncDeviceStream::new(4);
+
+ // Drop both without ever starting a producer
+ drop(consumer);
+ drop(handler);
+
+ // Should not panic or leak
+ }
+}
diff --git a/c/sedona-extension/src/lib.rs b/c/sedona-extension/src/lib.rs
index f91a59c7f..be3e5df8b 100644
--- a/c/sedona-extension/src/lib.rs
+++ b/c/sedona-extension/src/lib.rs
@@ -17,8 +17,10 @@
pub mod catalog;
pub mod execution_plan;
+pub mod export_sendable_record_batch_stream;
pub mod expr;
pub mod extension;
+pub mod import_sendable_record_batch_stream;
pub mod metrics;
pub mod runtime;
pub mod scalar_kernel;
diff --git a/c/sedona-extension/src/sedona_extension.h b/c/sedona-extension/src/sedona_extension.h
index 28638b097..c1d5d406a 100644
--- a/c/sedona-extension/src/sedona_extension.h
+++ b/c/sedona-extension/src/sedona_extension.h
@@ -113,6 +113,66 @@ struct ArrowArrayStream {
#endif // ARROW_C_STREAM_INTERFACE
#endif // ARROW_FLAG_DICTIONARY_ORDERED
+#ifndef ARROW_C_DEVICE_DATA_INTERFACE
+#define ARROW_C_DEVICE_DATA_INTERFACE
+
+typedef int32_t ArrowDeviceType;
+#define ARROW_DEVICE_CPU 1
+#define ARROW_DEVICE_CUDA 2
+#define ARROW_DEVICE_CUDA_HOST 3
+#define ARROW_DEVICE_OPENCL 4
+#define ARROW_DEVICE_VULKAN 7
+#define ARROW_DEVICE_METAL 8
+#define ARROW_DEVICE_VPI 9
+#define ARROW_DEVICE_ROCM 10
+#define ARROW_DEVICE_ROCM_HOST 11
+#define ARROW_DEVICE_EXT_DEV 12
+#define ARROW_DEVICE_CUDA_MANAGED 13
+#define ARROW_DEVICE_ONEAPI 14
+#define ARROW_DEVICE_WEBGPU 15
+#define ARROW_DEVICE_HEXAGON 16
+
+struct ArrowDeviceArray {
+ struct ArrowArray array;
+ int64_t device_id;
+ ArrowDeviceType device_type;
+ void* sync_event;
+ int64_t reserved[3];
+};
+
+#endif // ARROW_C_DEVICE_DATA_INTERFACE
+
+#ifndef ARROW_C_ASYNC_STREAM_INTERFACE
+#define ARROW_C_ASYNC_STREAM_INTERFACE
+
+/// Experimental Arrow C async-device stream interface.
+struct ArrowAsyncTask {
+ int (*extract_data)(struct ArrowAsyncTask* self, struct ArrowDeviceArray* out);
+ void* private_data;
+};
+
+struct ArrowAsyncProducer {
+ ArrowDeviceType device_type;
+ void (*request)(struct ArrowAsyncProducer* self, int64_t n);
+ void (*cancel)(struct ArrowAsyncProducer* self);
+ const char* additional_metadata;
+ void* private_data;
+};
+
+struct ArrowAsyncDeviceStreamHandler {
+ int (*on_schema)(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowSchema* stream_schema);
+ int (*on_next_task)(struct ArrowAsyncDeviceStreamHandler* self,
+ struct ArrowAsyncTask* task, const char* metadata);
+ void (*on_error)(struct ArrowAsyncDeviceStreamHandler* self, int code,
+ const char* message, const char* metadata);
+ void (*release)(struct ArrowAsyncDeviceStreamHandler* self);
+ struct ArrowAsyncProducer* producer;
+ void* private_data;
+};
+
+#endif // ARROW_C_ASYNC_STREAM_INTERFACE
+
/// \brief Simple ABI-stable scalar function implementation
///
/// This object is not thread safe: callers must take care to serialize
@@ -341,9 +401,9 @@ struct SedonaCExecutionPlan {
/// \brief Resolve an asynchronous stream for one partition from this plan
///
- /// This is not currently implemented and must be NULL. In the future,
- /// out must point to a caller-supplied struct ArrowAsyncDeviceStreamHandler
- /// as specified in the Arrow C Device Async Stream specification.
+ /// `out` must point to a caller-supplied ArrowAsyncDeviceStreamHandler.
+ /// Implementations may leave this callback NULL when async execution is not
+ /// supported; callers must check it before use.
int (*execute_async)(const struct SedonaCExecutionPlan* self,
struct SedonaCExecutionPlanArgs* args, void* out,
struct SedonaCError* err);
diff --git a/c/sedona-extension/src/table_provider.rs b/c/sedona-extension/src/table_provider.rs
index ad525de65..91edd7254 100644
--- a/c/sedona-extension/src/table_provider.rs
+++ b/c/sedona-extension/src/table_provider.rs
@@ -438,6 +438,7 @@ pub struct ImportedTableProvider {
cancel_checker: Option<Arc<dyn Fn() -> bool + Send + Sync>>,
/// Interval for periodic cancellation checking during stream consumption.
check_interval: Option<Duration>,
+ use_async_execution: bool,
}
impl Debug for ImportedTableProvider {
@@ -478,6 +479,7 @@ impl ImportedTableProvider {
table_reference,
cancel_checker: None,
check_interval: None,
+ use_async_execution: false,
})
}
@@ -509,6 +511,14 @@ impl ImportedTableProvider {
self
}
+ /// Select async-device-stream execution for plans returned by this provider.
+ ///
+ /// The synchronous Arrow C stream path remains the default.
+ pub fn with_async_execution(mut self, use_async: bool) -> Self {
+ self.use_async_execution = use_async;
+ self
+ }
+
fn get_table_type(provider: &SedonaCTableProvider) -> Result<TableType> {
let type_str = match get_table_provider_string_property(provider, "table_type") {
Ok(s) => s,
@@ -717,6 +727,8 @@ impl TableProvider for ImportedTableProvider {
let mut exec = ImportedSedonaCExec::try_new(ffi_plan)?;
+ exec = exec.with_async_execution(self.use_async_execution);
+
// Pipe through the cancel checker and interval if configured
if let Some(ref checker) = self.cancel_checker {
let checker = checker.clone();
@@ -897,6 +909,14 @@ mod tests {
/// 3. Importing it and registering in a new context
/// 4. Running the SQL query and asserting results
fn test_roundtrip_query(sql: &str, expected: &[&str]) -> Result<()> {
+ test_roundtrip_query_with_mode(sql, expected, false)
+ }
+
+ fn test_roundtrip_query_with_mode(
+ sql: &str,
+ expected: &[&str],
+ use_async_execution: bool,
+ ) -> Result<()> {
let runtime = test_runtime();
runtime.block_on(async {
let ctx = create_test_context().await?;
@@ -910,7 +930,8 @@ mod tests {
let ffi_provider: SedonaCTableProvider = exported.into();
// Import the table provider
- let imported = ImportedTableProvider::try_new(ffi_provider)?;
+ let imported = ImportedTableProvider::try_new(ffi_provider)?
+ .with_async_execution(use_async_execution);
// Create a new context and register the imported table
let ctx2 = SessionContext::new();
@@ -942,6 +963,26 @@ mod tests {
.unwrap();
}
+ #[test]
+ fn test_roundtrip_simple_select_async() {
+ test_roundtrip_query_with_mode(
+ "SELECT id, value_a FROM imported_data ORDER BY id LIMIT 5",
+ &[
+ "+----+---------+",
+ "| id | value_a |",
+ "+----+---------+",
+ "| 1 | 100 |",
+ "| 2 | 200 |",
+ "| 3 | 300 |",
+ "| 4 | 400 |",
+ "| 5 | 500 |",
+ "+----+---------+",
+ ],
+ true,
+ )
+ .unwrap();
+ }
+
#[test]
fn test_roundtrip_projection() {
test_roundtrip_query(
diff --git a/c/sedona-extension/tests/async_stream_interop.rs b/c/sedona-extension/tests/async_stream_interop.rs
new file mode 100644
index 000000000..d68c211b9
--- /dev/null
+++ b/c/sedona-extension/tests/async_stream_interop.rs
@@ -0,0 +1,398 @@
+// 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.
+
+// Foreign callbacks follow the C ABI without Sedona-specific cleanup behavior.
+use std::ffi::c_int;
+use std::pin::Pin;
+use std::ptr::{null, null_mut};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow_array::ffi::{to_ffi, FFI_ArrowSchema};
+use arrow_array::{Int32Array, RecordBatch, StructArray};
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use datafusion_execution::TaskContext;
+use datafusion_physical_plan::{empty::EmptyExec, ExecutionPlan};
+use futures::task::{waker, ArcWake};
+use futures::{Stream, StreamExt};
+use sedona_extension::execution_plan::{ExportedExecutionPlan, ImportedSedonaCExec};
+use sedona_extension::extension::{
+ FFI_ArrowAsyncDeviceStreamHandler, FFI_ArrowAsyncProducer, FFI_ArrowAsyncTask,
+ FFI_ArrowDeviceArray, ARROW_DEVICE_CPU,
+};
+use sedona_extension::import_sendable_record_batch_stream::ImportedAsyncDeviceStream;
+use sedona_extension::runtime::RuntimeHandle;
+
+fn schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]))
+}
+unsafe extern "C" fn request(_: *mut FFI_ArrowAsyncProducer, _: i64) {}
+unsafe extern "C" fn cancel(_: *mut FFI_ArrowAsyncProducer) {}
+fn producer() -> FFI_ArrowAsyncProducer {
+ FFI_ArrowAsyncProducer {
+ device_type: ARROW_DEVICE_CPU,
+ request: Some(request),
+ cancel: Some(cancel),
+ additional_metadata: null(),
+ private_data: null_mut(),
+ }
+}
+unsafe fn connect(
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ producer: &mut FFI_ArrowAsyncProducer,
+) {
+ (*handler).producer = producer;
+ let mut exported = FFI_ArrowSchema::try_from(schema().as_ref()).unwrap();
+ assert_eq!(((*handler).on_schema.unwrap())(handler, &mut exported), 0);
+ assert!(exported.release.is_none());
+}
+
+#[test]
+fn imported_schema_must_mark_original_released() {
+ let (_consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let mut producer = producer();
+ let ptr = handler.as_ptr();
+ unsafe {
+ (*ptr).producer = &mut producer;
+ }
+ let mut exported = FFI_ArrowSchema::try_from(schema().as_ref()).unwrap();
+ unsafe {
+ assert_eq!(((*ptr).on_schema.unwrap())(ptr, &mut exported), 0);
+ }
+ let released = exported.release.is_none();
+ // Keep a regression from double-freeing before the assertion can report it.
+ std::mem::forget(exported);
+ unsafe {
+ ((*ptr).release.unwrap())(ptr);
+ }
+ assert!(
+ released,
+ "C++ SchemaExportGuard will release already-freed schema contents"
+ );
+}
+
+struct ForeignTask {
+ calls: AtomicUsize,
+ fail: bool,
+}
+unsafe extern "C" fn extract(
+ task: *mut FFI_ArrowAsyncTask,
+ out: *mut FFI_ArrowDeviceArray,
+) -> c_int {
+ // Like Arrow C++ extract_data, leave the task's fields untouched.
+ // Keep the counter allocation live so a double invocation is observable safely.
+ let state = &*((*task).private_data as *const ForeignTask);
+ state.calls.fetch_add(1, Ordering::SeqCst);
+ if out.is_null() {
+ return 0;
+ }
+ if state.fail {
+ return libc::EINVAL;
+ }
+ let batch = RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(vec![42]))]).unwrap();
+ let data = StructArray::from(batch).into();
+ let (array, _) = to_ffi(&data).unwrap();
+ std::ptr::write(out, FFI_ArrowDeviceArray::from(array));
+ 0
+}
+fn check_extract_once(fail: bool) {
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let mut producer = producer();
+ let ptr = handler.as_ptr();
+ unsafe {
+ connect(ptr, &mut producer);
+ }
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+ assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
+ let state = ForeignTask {
+ calls: AtomicUsize::new(0),
+ fail,
+ };
+ let mut task = FFI_ArrowAsyncTask {
+ extract_data: Some(extract),
+ private_data: (&state as *const ForeignTask).cast_mut().cast(),
+ };
+ unsafe {
+ assert_eq!(((*ptr).on_next_task.unwrap())(ptr, &mut task, null()), 0);
+ }
+ std::mem::forget(task);
+ let result = Pin::new(&mut consumer).poll_next(&mut cx);
+ match result {
+ Poll::Ready(Some(result)) => {
+ assert_eq!(result.is_err(), fail);
+ if let Ok(batch) = result {
+ assert_eq!(batch.num_rows(), 1);
+ let values = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(values.value(0), 42);
+ }
+ }
+ other => panic!("expected task result, got {other:?}"),
+ }
+ unsafe {
+ ((*ptr).release.unwrap())(ptr);
+ }
+ assert_eq!(
+ state.calls.load(Ordering::SeqCst),
+ 1,
+ "extract_data must be called exactly once"
+ );
+}
+#[test]
+fn foreign_successful_task_is_extracted_once() {
+ check_extract_once(false);
+}
+#[test]
+fn foreign_failed_task_is_extracted_once() {
+ check_extract_once(true);
+}
+
+#[derive(Default)]
+struct WakeCounter(AtomicUsize);
+impl ArcWake for WakeCounter {
+ fn wake_by_ref(this: &Arc<Self>) {
+ this.0.fetch_add(1, Ordering::SeqCst);
+ }
+}
+#[test]
+fn release_without_eos_must_wake_waiting_consumer() {
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let mut producer = producer();
+ let ptr = handler.as_ptr();
+ unsafe {
+ connect(ptr, &mut producer);
+ }
+ let wakes = Arc::new(WakeCounter::default());
+ let waker = waker(wakes.clone());
+ let mut cx = Context::from_waker(&waker);
+ assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
+ consumer.cancel();
+ wakes.0.store(0, Ordering::SeqCst);
+ // A conforming producer may finish cancellation with release alone.
+ unsafe {
+ ((*ptr).release.unwrap())(ptr);
+ }
+ assert!(
+ wakes.0.load(Ordering::SeqCst) > 0,
+ "consumer is left asleep after release"
+ );
+ assert!(matches!(
+ Pin::new(&mut consumer).poll_next(&mut cx),
+ Poll::Ready(None)
+ ));
+}
+
+#[test]
+fn stream_keeps_producer_runtime_alive_after_plan_drop() {
+ let runtime = Arc::new(RuntimeHandle::new(
+ tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ .unwrap(),
+ ));
+ let ctx = Arc::new(TaskContext::default());
+ let exported = ExportedExecutionPlan::new(
+ Arc::new(EmptyExec::new(schema())),
+ ctx.clone(),
+ runtime.clone(),
+ );
+ let imported = ImportedSedonaCExec::try_new(exported.into())
+ .unwrap()
+ .with_async_execution(true);
+ let mut consumer = imported.execute(0, ctx).unwrap();
+ let weak = Arc::downgrade(&runtime);
+ drop(runtime);
+ drop(imported);
+
+ // Recover a temporary driver for this current-thread runtime. The task
+ // itself must own the runtime before and throughout stream consumption.
+ let runtime = weak.upgrade().expect("stream lost its producer runtime");
+ runtime.block_on(async {
+ assert!(
+ tokio::time::timeout(std::time::Duration::from_secs(2), consumer.next())
+ .await
+ .unwrap()
+ .is_none()
+ );
+ });
+ drop(consumer);
+ drop(runtime);
+ assert!(
+ weak.upgrade().is_none(),
+ "finished stream retained its runtime"
+ );
+}
+
+#[test]
+fn dropping_stream_releases_retained_runtime() {
+ let runtime = Arc::new(RuntimeHandle::new(
+ tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()
+ .unwrap(),
+ ));
+ let ctx = Arc::new(TaskContext::default());
+ let exported = ExportedExecutionPlan::new(
+ Arc::new(EmptyExec::new(schema())),
+ ctx.clone(),
+ runtime.clone(),
+ );
+ let imported = ImportedSedonaCExec::try_new(exported.into())
+ .unwrap()
+ .with_async_execution(true);
+ let consumer = imported.execute(0, ctx).unwrap();
+ let weak = Arc::downgrade(&runtime);
+ drop(imported);
+ drop(consumer);
+ runtime.block_on(async {
+ tokio::time::timeout(std::time::Duration::from_secs(2), async {
+ while Arc::strong_count(&runtime) > 1 {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .expect("cancelled driver retained its runtime");
+ });
+ drop(runtime);
+ assert!(weak.upgrade().is_none());
+}
+
+#[test]
+fn queued_batches_are_drained_before_error_or_eof() {
+ for fail in [false, true] {
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let mut producer = producer();
+ let ptr = handler.as_ptr();
+ unsafe {
+ connect(ptr, &mut producer);
+ }
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+ assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
+ let state = ForeignTask {
+ calls: AtomicUsize::new(0),
+ fail: false,
+ };
+ for _ in 0..2 {
+ let mut task = FFI_ArrowAsyncTask {
+ extract_data: Some(extract),
+ private_data: (&state as *const ForeignTask).cast_mut().cast(),
+ };
+ unsafe {
+ assert_eq!(((*ptr).on_next_task.unwrap())(ptr, &mut task, null()), 0);
+ }
+ std::mem::forget(task);
+ }
+ unsafe {
+ if fail {
+ ((*ptr).on_error.unwrap())(ptr, libc::EIO, c"expected error".as_ptr(), null());
+ } else {
+ assert_eq!(((*ptr).on_next_task.unwrap())(ptr, null_mut(), null()), 0);
+ }
+ ((*ptr).release.unwrap())(ptr);
+ }
+ // Extract tasks after the producer has already released the handler.
+ for _ in 0..2 {
+ assert!(matches!(
+ Pin::new(&mut consumer).poll_next(&mut cx),
+ Poll::Ready(Some(Ok(_)))
+ ));
+ }
+ if fail {
+ let Poll::Ready(Some(Err(error))) = Pin::new(&mut consumer).poll_next(&mut cx) else {
+ panic!("producer error was lost");
+ };
+ assert!(error.to_string().contains("expected error"));
+ }
+ assert!(matches!(
+ Pin::new(&mut consumer).poll_next(&mut cx),
+ Poll::Ready(None)
+ ));
+ assert_eq!(state.calls.load(Ordering::SeqCst), 2);
+ }
+}
+
+#[test]
+fn cancellation_before_schema_rejects_the_connection() {
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ consumer.cancel();
+ consumer.cancel();
+ let mut producer = producer();
+ let ptr = handler.as_ptr();
+ unsafe {
+ (*ptr).producer = &mut producer;
+ }
+ let mut exported = FFI_ArrowSchema::try_from(schema().as_ref()).unwrap();
+ let result = unsafe { ((*ptr).on_schema.unwrap())(ptr, &mut exported) };
+ let released = exported.release.is_none();
+ std::mem::forget(exported);
+ // A rejected on_schema must only be followed by release.
+ unsafe {
+ ((*ptr).release.unwrap())(ptr);
+ }
+ assert_eq!(result, libc::ECANCELED);
+ assert!(released, "rejected schema was not consumed");
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+ assert!(matches!(
+ Pin::new(&mut consumer).poll_next(&mut cx),
+ Poll::Ready(None)
+ ));
+}
+
+#[test]
+fn racing_error_must_not_be_reported_as_clean_eof() {
+ let (send, receive) = std::sync::mpsc::channel::<usize>();
+ let (done_tx, done_rx) = std::sync::mpsc::sync_channel(0);
+ let producer = std::thread::spawn(move || {
+ while let Ok(address) = receive.recv() {
+ let ptr = address as *mut FFI_ArrowAsyncDeviceStreamHandler;
+ unsafe {
+ ((*ptr).on_error.unwrap())(ptr, libc::EIO, c"expected error".as_ptr(), null());
+ ((*ptr).release.unwrap())(ptr);
+ }
+ done_tx.send(()).unwrap();
+ }
+ });
+ let mut lost = 0;
+ for _ in 0..50_000 {
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let address = handler.into_raw() as usize;
+ send.send(address).unwrap();
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+ loop {
+ match Pin::new(&mut consumer).poll_next(&mut cx) {
+ Poll::Ready(None) => {
+ lost += 1;
+ break;
+ }
+ Poll::Ready(Some(Err(_))) => break,
+ Poll::Ready(Some(Ok(_))) => panic!("unexpected batch"),
+ Poll::Pending => std::hint::spin_loop(),
+ }
+ }
+ done_rx.recv().unwrap();
+ }
+ drop(send);
+ producer.join().unwrap();
+ assert_eq!(
+ lost, 0,
+ "on_error raced with the ended flag and became clean EOF"
+ );
+}
diff --git a/c/sedona-extension/tests/async_stream_regressions.rs b/c/sedona-extension/tests/async_stream_regressions.rs
new file mode 100644
index 000000000..9fa8f1aa9
--- /dev/null
+++ b/c/sedona-extension/tests/async_stream_regressions.rs
@@ -0,0 +1,546 @@
+// 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.
+
+use std::collections::HashMap;
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow_array::{Int32Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use datafusion_common::Result;
+use datafusion_execution::RecordBatchStream;
+use futures::channel::oneshot;
+use futures::{Stream, StreamExt};
+use sedona_extension::export_sendable_record_batch_stream::drive_stream_to_handler;
+use sedona_extension::extension::{
+ FFI_ArrowAsyncDeviceStreamHandler, FFI_ArrowAsyncProducer, FFI_ArrowAsyncTask,
+ FFI_ArrowDeviceArray, ARROW_DEVICE_CPU,
+};
+use sedona_extension::import_sendable_record_batch_stream::ImportedAsyncDeviceStream;
+
+fn schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new(
+ "value",
+ DataType::Int32,
+ false,
+ )]))
+}
+
+struct PendingStream {
+ schema: SchemaRef,
+ polls: Arc<AtomicUsize>,
+ drops: Arc<AtomicUsize>,
+}
+
+impl Stream for PendingStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.polls.fetch_add(1, Ordering::SeqCst);
+ Poll::Pending
+ }
+}
+
+impl RecordBatchStream for PendingStream {
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+}
+
+impl Drop for PendingStream {
+ fn drop(&mut self) {
+ self.drops.fetch_add(1, Ordering::SeqCst);
+ }
+}
+
+struct DelayedEmptyStream {
+ schema: SchemaRef,
+ ready: oneshot::Receiver<()>,
+ polls: Arc<AtomicUsize>,
+ drops: Arc<AtomicUsize>,
+}
+
+impl Stream for DelayedEmptyStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ self.polls.fetch_add(1, Ordering::SeqCst);
+ match Pin::new(&mut self.ready).poll(cx) {
+ Poll::Pending => Poll::Pending,
+ Poll::Ready(_) => Poll::Ready(None),
+ }
+ }
+}
+
+impl RecordBatchStream for DelayedEmptyStream {
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+}
+
+impl Drop for DelayedEmptyStream {
+ fn drop(&mut self) {
+ self.drops.fetch_add(1, Ordering::SeqCst);
+ }
+}
+
+#[test]
+fn request_racing_waker_registration_is_not_lost() {
+ use std::task::{RawWaker, RawWakerVTable, Waker};
+
+ struct WakeState {
+ producer: AtomicPtr<FFI_ArrowAsyncProducer>,
+ wakes: AtomicUsize,
+ }
+
+ unsafe fn clone_waker(data: *const ()) -> RawWaker {
+ let state = &*(data as *const WakeState);
+ let producer = state.producer.load(Ordering::SeqCst);
+ if !producer.is_null() {
+ ((*producer).request.unwrap())(producer, 1);
+ }
+ Arc::increment_strong_count(data as *const WakeState);
+ RawWaker::new(data, &VTABLE)
+ }
+
+ unsafe fn wake(data: *const ()) {
+ let state = Arc::from_raw(data as *const WakeState);
+ state.wakes.fetch_add(1, Ordering::SeqCst);
+ }
+
+ unsafe fn wake_by_ref(data: *const ()) {
+ (*(data as *const WakeState))
+ .wakes
+ .fetch_add(1, Ordering::SeqCst);
+ }
+
+ unsafe fn drop_waker(data: *const ()) {
+ drop(Arc::from_raw(data as *const WakeState));
+ }
+
+ static VTABLE: RawWakerVTable = RawWakerVTable::new(clone_waker, wake, wake_by_ref, drop_waker);
+
+ unsafe extern "C" fn on_schema(
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ schema: *mut arrow_array::ffi::FFI_ArrowSchema,
+ ) -> std::ffi::c_int {
+ drop(arrow_array::ffi::FFI_ArrowSchema::from_raw(schema));
+ (*((*handler).private_data as *const WakeState))
+ .producer
+ .store((*handler).producer, Ordering::SeqCst);
+ 0
+ }
+
+ let wake_state = Arc::new(WakeState {
+ producer: AtomicPtr::new(std::ptr::null_mut()),
+ wakes: AtomicUsize::new(0),
+ });
+ let waker = unsafe {
+ Waker::from_raw(RawWaker::new(
+ Arc::into_raw(wake_state.clone()).cast(),
+ &VTABLE,
+ ))
+ };
+ let mut handler = FFI_ArrowAsyncDeviceStreamHandler {
+ on_schema: Some(on_schema),
+ on_next_task: None,
+ on_error: None,
+ release: None,
+ producer: std::ptr::null_mut(),
+ private_data: Arc::as_ptr(&wake_state) as *mut _,
+ };
+ let source_polls = Arc::new(AtomicUsize::new(0));
+ let source = Box::pin(PendingStream {
+ schema: schema(),
+ polls: source_polls.clone(),
+ drops: Arc::new(AtomicUsize::new(0)),
+ });
+ let mut driver = Box::pin(unsafe { drive_stream_to_handler(source, &mut handler) });
+ let mut cx = Context::from_waker(&waker);
+ let result = driver.as_mut().poll(&mut cx);
+
+ assert!(
+ result.is_ready()
+ || wake_state.wakes.load(Ordering::SeqCst) > 0
+ || source_polls.load(Ordering::SeqCst) > 0,
+ "request must schedule another poll or be observed before returning Pending"
+ );
+}
+
+#[test]
+fn cancellation_interrupts_pending_source() {
+ let polls = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ let source = Box::pin(PendingStream {
+ schema: schema(),
+ polls: polls.clone(),
+ drops: drops.clone(),
+ });
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let mut driver = Box::pin(unsafe { drive_stream_to_handler(source, handler.as_ptr()) });
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+
+ assert!(driver.as_mut().poll(&mut cx).is_pending());
+ assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
+ assert!(driver.as_mut().poll(&mut cx).is_pending());
+ assert_eq!(polls.load(Ordering::SeqCst), 1);
+ consumer.cancel();
+
+ assert!(driver.as_mut().poll(&mut cx).is_ready());
+ assert_eq!(drops.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn cancellation_before_first_poll_does_not_poll_source() {
+ let polls = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ let source = Box::pin(PendingStream {
+ schema: schema(),
+ polls: polls.clone(),
+ drops: drops.clone(),
+ });
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ consumer.cancel();
+ // SAFETY: into_raw transfers the handler to this driver's sole ownership.
+ unsafe { drive_stream_to_handler(source, handler.into_raw()) }.await;
+
+ assert!(consumer.next().await.is_none());
+ assert_eq!(polls.load(Ordering::SeqCst), 0);
+ assert_eq!(drops.load(Ordering::SeqCst), 1);
+}
+
+#[test]
+fn prefetch_one_requests_a_batch() {
+ let polls = Arc::new(AtomicUsize::new(0));
+ let source = Box::pin(PendingStream {
+ schema: schema(),
+ polls: polls.clone(),
+ drops: Arc::new(AtomicUsize::new(0)),
+ });
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(1);
+ let mut driver = Box::pin(unsafe { drive_stream_to_handler(source, handler.as_ptr()) });
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+
+ assert!(driver.as_mut().poll(&mut cx).is_pending());
+ assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
+ assert!(driver.as_mut().poll(&mut cx).is_pending());
+ assert_eq!(polls.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn delayed_empty_stream_completes_without_a_batch() {
+ let polls = Arc::new(AtomicUsize::new(0));
+ let drops = Arc::new(AtomicUsize::new(0));
+ let (ready_tx, ready_rx) = oneshot::channel();
+ let source = Box::pin(DelayedEmptyStream {
+ schema: schema(),
+ ready: ready_rx,
+ polls: polls.clone(),
+ drops: drops.clone(),
+ });
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+
+ let consumption = async {
+ assert!(consumer.next().await.is_none());
+ };
+ let release_source = async {
+ // Ensure the driver and consumer both observe the pending phase before
+ // allowing the source to complete without producing a batch.
+ tokio::task::yield_now().await;
+ tokio::task::yield_now().await;
+ ready_tx.send(()).unwrap();
+ };
+
+ futures::join!(
+ consumption,
+ unsafe { drive_stream_to_handler(source, handler.as_ptr()) },
+ release_source
+ );
+
+ assert!(polls.load(Ordering::SeqCst) >= 2);
+ assert_eq!(drops.load(Ordering::SeqCst), 1);
+}
+
+static FOREIGN_TASK_CLEANUPS: AtomicUsize = AtomicUsize::new(0);
+
+unsafe extern "C" fn foreign_request(_: *mut FFI_ArrowAsyncProducer, _: i64) {}
+unsafe extern "C" fn foreign_cancel(_: *mut FFI_ArrowAsyncProducer) {}
+unsafe extern "C" fn foreign_extract(
+ _: *mut FFI_ArrowAsyncTask,
+ out: *mut FFI_ArrowDeviceArray,
+) -> std::ffi::c_int {
+ assert!(out.is_null());
+ FOREIGN_TASK_CLEANUPS.fetch_add(1, Ordering::SeqCst);
+ 0
+}
+
+#[test]
+fn abandoning_a_foreign_task_calls_extract_null() {
+ FOREIGN_TASK_CLEANUPS.store(0, Ordering::SeqCst);
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let mut producer = FFI_ArrowAsyncProducer {
+ device_type: ARROW_DEVICE_CPU,
+ request: Some(foreign_request),
+ cancel: Some(foreign_cancel),
+ additional_metadata: std::ptr::null(),
+ private_data: std::ptr::null_mut(),
+ };
+ let ptr = handler.as_ptr();
+ unsafe {
+ (*ptr).producer = &mut producer;
+ }
+ let mut ffi_schema = arrow_array::ffi::FFI_ArrowSchema::try_from(schema().as_ref()).unwrap();
+ assert_eq!(
+ unsafe { ((*ptr).on_schema.unwrap())(ptr, &mut ffi_schema) },
+ 0
+ );
+ assert!(ffi_schema.release.is_none());
+ let mut cx = Context::from_waker(futures::task::noop_waker_ref());
+ assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
+
+ let mut task = FFI_ArrowAsyncTask {
+ extract_data: Some(foreign_extract),
+ private_data: std::ptr::null_mut(),
+ };
+ assert_eq!(
+ unsafe { ((*ptr).on_next_task.unwrap())(ptr, &mut task, std::ptr::null()) },
+ 0
+ );
+ std::mem::forget(task);
+ drop(consumer);
+ unsafe {
+ ((*ptr).release.unwrap())(ptr);
+ }
+
+ assert_eq!(FOREIGN_TASK_CLEANUPS.load(Ordering::SeqCst), 1);
+}
+
+struct ReadyStream {
+ schema: SchemaRef,
+ batch: Option<RecordBatch>,
+}
+
+impl Stream for ReadyStream {
+ type Item = Result<RecordBatch>;
+
+ fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ Poll::Ready(self.batch.take().map(Ok))
+ }
+}
+
+impl RecordBatchStream for ReadyStream {
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+}
+
+#[derive(Default)]
+struct RejectingHandlerState {
+ schema_calls: AtomicUsize,
+ next_calls: AtomicUsize,
+ error_calls: AtomicUsize,
+ release_calls: AtomicUsize,
+}
+
+unsafe extern "C" fn reject_schema(
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ schema: *mut arrow_array::ffi::FFI_ArrowSchema,
+) -> std::ffi::c_int {
+ drop(arrow_array::ffi::FFI_ArrowSchema::from_raw(schema));
+ let state = &*((*handler).private_data as *const RejectingHandlerState);
+ state.schema_calls.fetch_add(1, Ordering::SeqCst);
+ libc::EINVAL
+}
+
+unsafe extern "C" fn accept_schema_and_request_one(
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ schema: *mut arrow_array::ffi::FFI_ArrowSchema,
+) -> std::ffi::c_int {
+ drop(arrow_array::ffi::FFI_ArrowSchema::from_raw(schema));
+ let state = &*((*handler).private_data as *const RejectingHandlerState);
+ state.schema_calls.fetch_add(1, Ordering::SeqCst);
+ let producer = (*handler).producer;
+ ((*producer).request.unwrap())(producer, 1);
+ 0
+}
+
+unsafe extern "C" fn reject_task(
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ _task: *mut FFI_ArrowAsyncTask,
+ _metadata: *const std::ffi::c_char,
+) -> std::ffi::c_int {
+ let state = &*((*handler).private_data as *const RejectingHandlerState);
+ state.next_calls.fetch_add(1, Ordering::SeqCst);
+ libc::EINVAL
+}
+
+unsafe extern "C" fn count_error(
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ _code: std::ffi::c_int,
+ _message: *const std::ffi::c_char,
+ _metadata: *const std::ffi::c_char,
+) {
+ let state = &*((*handler).private_data as *const RejectingHandlerState);
+ state.error_calls.fetch_add(1, Ordering::SeqCst);
+}
+
+unsafe extern "C" fn count_release(handler: *mut FFI_ArrowAsyncDeviceStreamHandler) {
+ let state = &*((*handler).private_data as *const RejectingHandlerState);
+ state.release_calls.fetch_add(1, Ordering::SeqCst);
+ (*handler).release = None;
+}
+
+#[test]
+fn unpolled_sendable_driver_releases_handler_once() {
+ fn assert_send<T: Send>(_: &T) {}
+
+ let state = RejectingHandlerState::default();
+ let mut handler = FFI_ArrowAsyncDeviceStreamHandler {
+ on_schema: Some(reject_schema),
+ on_next_task: Some(reject_task),
+ on_error: Some(count_error),
+ release: Some(count_release),
+ producer: std::ptr::null_mut(),
+ private_data: (&state as *const RejectingHandlerState).cast_mut().cast(),
+ };
+ let source_drops = Arc::new(AtomicUsize::new(0));
+ let source = Box::pin(PendingStream {
+ schema: schema(),
+ polls: Arc::new(AtomicUsize::new(0)),
+ drops: source_drops.clone(),
+ });
+
+ let driver = unsafe { drive_stream_to_handler(source, &mut handler) };
+ assert_send(&driver);
+ drop(driver);
+
+ assert_eq!(state.schema_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.next_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.error_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(state.release_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(source_drops.load(Ordering::SeqCst), 1);
+}
+
+#[test]
+fn runtime_shutdown_before_first_poll_releases_handler_once() {
+ let state = RejectingHandlerState::default();
+ let mut handler = FFI_ArrowAsyncDeviceStreamHandler {
+ on_schema: Some(reject_schema),
+ on_next_task: Some(reject_task),
+ on_error: Some(count_error),
+ release: Some(count_release),
+ producer: std::ptr::null_mut(),
+ private_data: (&state as *const RejectingHandlerState).cast_mut().cast(),
+ };
+ let source_drops = Arc::new(AtomicUsize::new(0));
+ let source = Box::pin(PendingStream {
+ schema: schema(),
+ polls: Arc::new(AtomicUsize::new(0)),
+ drops: source_drops.clone(),
+ });
+ let runtime = tokio::runtime::Builder::new_current_thread()
+ .build()
+ .unwrap();
+
+ runtime.spawn(unsafe { drive_stream_to_handler(source, &mut handler) });
+ drop(runtime);
+
+ assert_eq!(state.schema_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.next_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.error_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(state.release_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(source_drops.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn rejected_schema_is_consumed_and_only_followed_by_release() {
+ let state = RejectingHandlerState::default();
+ let mut handler = FFI_ArrowAsyncDeviceStreamHandler {
+ on_schema: Some(reject_schema),
+ on_next_task: Some(reject_task),
+ on_error: Some(count_error),
+ release: Some(count_release),
+ producer: std::ptr::null_mut(),
+ private_data: (&state as *const RejectingHandlerState).cast_mut().cast(),
+ };
+ let source = Box::pin(ReadyStream {
+ schema: schema(),
+ batch: None,
+ });
+
+ unsafe { drive_stream_to_handler(source, &mut handler) }.await;
+
+ assert_eq!(state.schema_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(state.next_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.error_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.release_calls.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn rejected_task_is_only_followed_by_release() {
+ let state = RejectingHandlerState::default();
+ let mut handler = FFI_ArrowAsyncDeviceStreamHandler {
+ on_schema: Some(accept_schema_and_request_one),
+ on_next_task: Some(reject_task),
+ on_error: Some(count_error),
+ release: Some(count_release),
+ producer: std::ptr::null_mut(),
+ private_data: (&state as *const RejectingHandlerState).cast_mut().cast(),
+ };
+ let batch = RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap();
+ let source = Box::pin(ReadyStream {
+ schema: schema(),
+ batch: Some(batch),
+ });
+
+ unsafe { drive_stream_to_handler(source, &mut handler) }.await;
+
+ assert_eq!(state.schema_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(state.next_calls.load(Ordering::SeqCst), 1);
+ assert_eq!(state.error_calls.load(Ordering::SeqCst), 0);
+ assert_eq!(state.release_calls.load(Ordering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn roundtrip_preserves_schema_metadata() {
+ let metadata = HashMap::from([("review_source".to_owned(), "table-v1".to_owned())]);
+ let schema = Arc::new(schema().as_ref().clone().with_metadata(metadata));
+ let batch =
+ RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap();
+ let source = Box::pin(ReadyStream {
+ schema: schema.clone(),
+ batch: Some(batch),
+ });
+ let (mut consumer, handler) = ImportedAsyncDeviceStream::new(2);
+ let consumption = async {
+ let mut batches = vec![];
+ while let Some(batch) = consumer.next().await {
+ batches.push(batch.unwrap());
+ }
+ (batches, consumer.schema())
+ };
+
+ let ((batches, actual_schema), _) = futures::join!(consumption, unsafe {
+ drive_stream_to_handler(source, handler.as_ptr())
+ });
+
+ assert_eq!(actual_schema, schema);
+ assert_eq!(batches[0].schema(), schema);
+}
diff --git a/python/sedonadb/Cargo.toml b/python/sedonadb/Cargo.toml
index 3b1c89c56..262aa2adf 100644
--- a/python/sedonadb/Cargo.toml
+++ b/python/sedonadb/Cargo.toml
@@ -48,6 +48,7 @@ datafusion-ffi = { workspace = true }
futures = { workspace = true }
pyo3 = { workspace = true }
sedona = { workspace = true }
+sedona-common = { workspace = true }
sedona-adbc = { workspace = true }
sedona-datasource = { workspace = true }
sedona-extension = { workspace = true }
diff --git a/python/sedonadb/src/context.rs b/python/sedonadb/src/context.rs
index cd81d8667..5525ed5da 100644
--- a/python/sedonadb/src/context.rs
+++ b/python/sedonadb/src/context.rs
@@ -26,6 +26,7 @@ use datafusion_expr::ScalarUDFImpl;
use pyo3::prelude::*;
use sedona::context::SedonaContext;
use sedona::context_builder::SedonaContextBuilder;
+use sedona_common::SedonaOptions;
use sedona_datasource::format::ExternalFormatFactory;
use sedona_extension::runtime::RuntimeHandle;
@@ -126,8 +127,17 @@ impl InternalContext {
obj: &Bound<PyAny>,
requested_schema: Option<&Bound<PyAny>>,
) -> Result<InternalDataFrame, PySedonaError> {
+ let use_async_execution = self
+ .inner
+ .ctx
+ .state()
+ .config_options()
+ .extensions
+ .get::<SedonaOptions>()
+ .map(|options| options.ffi.use_async)
+ .unwrap_or(false);
let (provider, table_reference) =
- import_table_provider_from_any(py, obj, requested_schema)?;
+ import_table_provider_from_any(py, obj, requested_schema, use_async_execution)?;
let df = if let Some(table_reference) = table_reference {
let plan =
LogicalPlanBuilder::scan(table_reference, provider_as_source(provider), None)?
diff --git a/python/sedonadb/src/import_from.rs b/python/sedonadb/src/import_from.rs
index 638b53064..d28b7b5f2 100644
--- a/python/sedonadb/src/import_from.rs
+++ b/python/sedonadb/src/import_from.rs
@@ -50,9 +50,10 @@ pub fn import_table_provider_from_any<'py>(
py: Python<'py>,
obj: &Bound<PyAny>,
requested_schema: Option<&Bound<PyAny>>,
+ use_async_execution: bool,
) -> Result<(Arc<dyn TableProvider>, Option<TableReference>), PySedonaError> {
if obj.hasattr("__sedonadb_table_provider__")? {
- import_sedona_ffi_table_provider(obj)
+ import_sedona_ffi_table_provider(obj, use_async_execution)
} else if obj.hasattr("__arrow_c_stream__")? {
let reader = import_arrow_array_stream(py, obj, requested_schema)?;
Ok((Arc::new(RecordBatchReaderProvider::new(reader)), None))
@@ -65,6 +66,7 @@ pub fn import_table_provider_from_any<'py>(
pub fn import_sedona_ffi_table_provider(
obj: &Bound<PyAny>,
+ use_async_execution: bool,
) -> Result<(Arc<dyn TableProvider>, Option<TableReference>), PySedonaError> {
let capsule = obj.getattr("__sedonadb_table_provider__")?.call0()?;
let contents =
@@ -79,7 +81,8 @@ pub fn import_sedona_ffi_table_provider(
provider
};
// try_new validates the release callback
- let provider = ImportedTableProvider::try_new(ffi_provider)?;
+ let provider =
+ ImportedTableProvider::try_new(ffi_provider)?.with_async_execution(use_async_execution);
// Add a Python-aware cancel checker that checks for Ctrl+C signals
// Use a 2 second interval to match the StreamingRecordBatchReader behavior
diff --git a/python/sedonadb/tests/test_dataframe_ffi.py b/python/sedonadb/tests/test_dataframe_ffi.py
index c08aebdc3..2881191c7 100644
--- a/python/sedonadb/tests/test_dataframe_ffi.py
+++ b/python/sedonadb/tests/test_dataframe_ffi.py
@@ -18,12 +18,10 @@
import pandas as pd
import pyarrow as pa
import pytest
-
import sedonadb
from sedonadb import udf
from sedonadb.testing import skip_if_not_exists
-
# Test cases: (producer_sql, consumer_sql)
FFI_TEST_CASES = [
# Simple select all
@@ -101,8 +99,9 @@ FFI_TEST_CASES = [
]
+@pytest.mark.parametrize("use_async", [True, False])
@pytest.mark.parametrize("producer_sql,consumer_sql", FFI_TEST_CASES)
-def test_ffi_roundtrip(geoarrow_data, producer_sql, consumer_sql):
+def test_ffi_roundtrip(geoarrow_data, producer_sql, consumer_sql, use_async):
# Use a real file with a reasonable number of rows so that parallelism
# and multiple batches are invoked
path = geoarrow_data / "ns-water" / "files" / "ns-water_water-point_geo.parquet"
@@ -110,6 +109,7 @@ def test_ffi_roundtrip(geoarrow_data, producer_sql, consumer_sql):
sd_producer = sedonadb.connect()
sd_consumer = sedonadb.connect()
+ sd_consumer.sql(f"SET sedona.ffi.use_async TO {str(use_async).lower()}").execute()
sd_producer.read_parquet(path).to_view("water_point")
df_producer = sd_producer.sql(producer_sql)
@@ -125,12 +125,14 @@ def test_ffi_roundtrip(geoarrow_data, producer_sql, consumer_sql):
pd.testing.assert_frame_equal(result_no_ffi, result_over_ffi)
-def test_filter_pushdown_into_ffi_producer(geoarrow_data):
+@pytest.mark.parametrize("use_async", [True, False])
+def test_filter_pushdown_into_ffi_producer(geoarrow_data, use_async):
path = geoarrow_data / "ns-water" / "files" / "ns-water_water-point_geo.parquet"
skip_if_not_exists(path)
sd_producer = sedonadb.connect()
sd_consumer = sedonadb.connect()
+ sd_consumer.sql(f"SET sedona.ffi.use_async TO {str(use_async).lower()}").execute()
# Producer exposes all columns without filtering
sd_producer.read_parquet(path).to_view("water_point")
@@ -167,7 +169,8 @@ def test_filter_pushdown_into_ffi_producer(geoarrow_data):
)
-def test_udf_filter_not_pushed_down_into_ffi_producer(geoarrow_data):
+@pytest.mark.parametrize("use_async", [True, False])
+def test_udf_filter_not_pushed_down_into_ffi_producer(geoarrow_data, use_async):
"""Verify that filters using consumer-side UDFs are NOT pushed down into the FFI producer.
UDFs are session-specific, so a UDF registered on the consumer cannot be
@@ -178,6 +181,7 @@ def test_udf_filter_not_pushed_down_into_ffi_producer(geoarrow_data):
sd_producer = sedonadb.connect()
sd_consumer = sedonadb.connect()
+ sd_consumer.sql(f"SET sedona.ffi.use_async TO {str(use_async).lower()}").execute()
# Define a simple UDF on the consumer side only
@udf.arrow_udf(pa.bool_(), [pa.int64()])
diff --git a/rust/sedona-common/src/option.rs b/rust/sedona-common/src/option.rs
index 65716a109..2a02b7c0e 100644
--- a/rust/sedona-common/src/option.rs
+++ b/rust/sedona-common/src/option.rs
@@ -51,6 +51,9 @@ config_namespace! {
/// Options for raster execution
pub raster: RasterOptions, default = RasterOptions::default()
+
+ /// Options for exchanging tables and execution plans across FFI
+ pub ffi: FfiOptions, default = FfiOptions::default()
}
}
@@ -96,6 +99,15 @@ config_namespace! {
}
}
+config_namespace! {
+ /// Configuration options for Sedona's FFI interfaces.
+ pub struct FfiOptions {
+ /// Use the experimental Arrow async-device stream interface when
+ /// importing a table provider or execution plan across FFI.
+ pub use_async: bool, default = false
+ }
+}
+
config_namespace! {
/// Configuration options for GDAL usage.
pub struct GdalOptions {
@@ -758,4 +770,13 @@ mod tests {
"Unexpected error message: {err_msg}"
);
}
+
+ #[test]
+ fn test_ffi_use_async_option() {
+ let mut options = SedonaOptions::default();
+ assert!(!options.ffi.use_async);
+
+ <SedonaOptions as ConfigField>::set(&mut options, "ffi.use_async", "true").unwrap();
+ assert!(options.ffi.use_async);
+ }
}