You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "paleolimbot (via GitHub)" <gi...@apache.org> on 2026/09/22 01:38:50 UTC
[PR] feat(sedona-extension): implement async FFI execution [sedona-db]
paleolimbot opened a new pull request, #1333:
URL: https://github.com/apache/sedona-db/pull/1333
Reimplements the async stream work from #914 against current main and completes SedonaCExecutionPlan.execute_async. Synchronous execution remains the default; SET sedona.ffi.use_async TO true selects async imports at runtime. Includes backpressure, cancellation, ownership cleanup, Rust unit tests, and Python FFI coverage through both paths.
Supersedes and closes #914.
Python example:
producer = sedonadb.connect()
consumer = sedonadb.connect()
consumer.sql("SET sedona.ffi.use_async TO true").execute()
df = producer.sql("SELECT 1 AS value")
print(consumer.create_data_frame(df).to_pandas())
# value
# 0 1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-extension): implement async FFI execution [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on PR #1333:
URL: https://github.com/apache/sedona-db/pull/1333#issuecomment-5823968704
Thank you for these!
I went through and fixed these issues, and also had a different model review it and compare with the Arrow C++ implementation of this to find a few more issues. There are a lot of corner cases here which is why it's not the default quite yet.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-extension): implement async FFI execution [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot merged PR #1333:
URL: https://github.com/apache/sedona-db/pull/1333
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(c/sedona-extension): implement async FFI execution [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on code in PR #1333:
URL: https://github.com/apache/sedona-db/pull/1333#discussion_r4087977900
##########
c/sedona-extension/src/export_sendable_record_batch_stream.rs:
##########
@@ -0,0 +1,472 @@
+// 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::ptr::null_mut;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::task::{Poll, Waker};
+
+use arrow_array::ffi::{to_ffi, FFI_ArrowSchema};
+use arrow_array::RecordBatch;
+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: std::sync::Mutex<Option<Waker>>,
+}
+
+impl ProducerState {
+ fn new() -> Self {
+ Self {
+ requested: AtomicU64::new(0),
+ cancelled: AtomicBool::new(false),
+ waker: std::sync::Mutex::new(None),
+ }
+ }
+
+ 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 register_waker(&self, waker: Waker) {
+ // Lock cannot be poisoned: we never panic while holding it
+ *self.waker.lock().expect("waker mutex poisoned") = Some(waker);
+ }
+
+ fn wake(&self) {
+ // Lock cannot be poisoned: we never panic while holding it
+ if let Some(waker) = self.waker.lock().expect("waker mutex poisoned").take() {
+ 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 an error occurs (including cancellation)
+///
+/// # Safety
+///
+/// The `handler` pointer must be valid and point to a properly initialized
+/// [`FFI_ArrowAsyncDeviceStreamHandler`].
+///
+/// # 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
+/// drive_stream_to_handler(stream, handler).await;
+/// ```
+pub async fn drive_stream_to_handler(
+ stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+) {
+ if handler.is_null() {
+ return;
+ }
+
+ // Create shared state for producer callbacks
+ let state = Arc::new(ProducerState::new());
+
+ // Create the producer
+ let mut producer = ProducerGuard(create_producer(Arc::clone(&state)));
+
+ // Use a guard to handle unexpected drops (panics)
+ let mut guard = StreamDriverGuard::new(handler);
+
+ let _result = drive_stream_inner(stream, handler, &mut producer.0, state).await;
+
+ // Normal completion (success or error) - disarm the guard and release handler
+ guard.disarm();
+ release_handler(handler);
+}
+
+/// Inner implementation that returns a Result for cleaner control flow.
+async fn drive_stream_inner(
+ mut stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ producer: &mut FFI_ArrowAsyncProducer,
+ state: Arc<ProducerState>,
+) -> Result<(), ()> {
+ let handler_ref = unsafe { &mut *handler };
+
+ // Set the producer on the handler BEFORE calling on_schema (per Arrow spec)
+ handler_ref.producer = producer;
+
+ // 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_ref, 1, &e.to_string());
+ return Err(());
+ }
+ };
+
+ if let Some(on_schema) = handler_ref.on_schema {
+ let ret = unsafe { on_schema(handler, &mut ffi_schema) };
+ if ret != 0 {
+ call_on_error(handler_ref, ret, "on_schema callback failed");
+ return Err(());
+ }
+ // The handler now owns the schema
+ std::mem::forget(ffi_schema);
+ }
+
+ // 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_ref, handler);
+ return Ok(());
+ }
+
+ // Wait for consumer to request batches (back-pressure)
+ // Use poll_fn to properly register our waker with ProducerState
+ futures::future::poll_fn(|cx| {
+ if state.has_requests() || state.is_cancelled() {
+ Poll::Ready(())
+ } else {
+ state.register_waker(cx.waker().clone());
Review Comment:
Could a request get lost between the check above and registering the waker? If `request(1)` arrives there, it increments the counter while there's no waker to notify. We then return `Pending`; the consumer is waiting for a batch before requesting more.
I forced that ordering by making the waker's clone call `request(1)`. The result was **`Pending`, with 0 wake calls and 0 source polls**. Registering the waker before checking the state, or checking again afterward, seems like it would close this gap.
<details>
<summary>Repro</summary>
Save this as `c/sedona-extension/tests/pr1333_request_wake.rs` and run from the repository root:
```sh
cargo test -p sedona-extension --test pr1333_request_wake -- --nocapture
```
```rust
use arrow_array::RecordBatch;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion_common::Result;
use datafusion_execution::RecordBatchStream;
use futures::Stream;
use sedona_extension::export_sendable_record_batch_stream::drive_stream_to_handler;
use std::{future::Future, pin::Pin, sync::{atomic::{AtomicUsize, Ordering}, Arc}, task::{Context, Poll}};
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);
}
}
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Int32,
false,
)]))
}
#[test]
fn request_racing_waker_registration_is_not_lost() {
use sedona_extension::extension::{FFI_ArrowAsyncDeviceStreamHandler, FFI_ArrowAsyncProducer};
use std::task::{RawWaker, RawWakerVTable, Waker};
struct WakeState {
producer: std::sync::atomic::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);
// Force the precise interleaving: request arrives after has_requests()
// returned false, but before register_waker() stores its cloned waker.
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(std::ptr::read(schema));
(*((*handler).private_data as *const WakeState))
.producer
.store((*handler).producer, Ordering::SeqCst);
0
}
let wake_state = Arc::new(WakeState {
producer: std::sync::atomic::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(drive_stream_to_handler(source, &mut handler));
let mut cx = Context::from_waker(&waker);
let result = driver.as_mut().poll(&mut cx);
println!(
"request arrived while registering waker: poll={result:?}, wake calls={}, source polls={}",
wake_state.wakes.load(Ordering::SeqCst),
source_polls.load(Ordering::SeqCst)
);
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"
);
}
```
</details>
##########
c/sedona-extension/src/export_sendable_record_batch_stream.rs:
##########
@@ -0,0 +1,472 @@
+// 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::ptr::null_mut;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::task::{Poll, Waker};
+
+use arrow_array::ffi::{to_ffi, FFI_ArrowSchema};
+use arrow_array::RecordBatch;
+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: std::sync::Mutex<Option<Waker>>,
+}
+
+impl ProducerState {
+ fn new() -> Self {
+ Self {
+ requested: AtomicU64::new(0),
+ cancelled: AtomicBool::new(false),
+ waker: std::sync::Mutex::new(None),
+ }
+ }
+
+ 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 register_waker(&self, waker: Waker) {
+ // Lock cannot be poisoned: we never panic while holding it
+ *self.waker.lock().expect("waker mutex poisoned") = Some(waker);
+ }
+
+ fn wake(&self) {
+ // Lock cannot be poisoned: we never panic while holding it
+ if let Some(waker) = self.waker.lock().expect("waker mutex poisoned").take() {
+ 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 an error occurs (including cancellation)
+///
+/// # Safety
+///
+/// The `handler` pointer must be valid and point to a properly initialized
+/// [`FFI_ArrowAsyncDeviceStreamHandler`].
+///
+/// # 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
+/// drive_stream_to_handler(stream, handler).await;
+/// ```
+pub async fn drive_stream_to_handler(
+ stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+) {
+ if handler.is_null() {
+ return;
+ }
+
+ // Create shared state for producer callbacks
+ let state = Arc::new(ProducerState::new());
+
+ // Create the producer
+ let mut producer = ProducerGuard(create_producer(Arc::clone(&state)));
+
+ // Use a guard to handle unexpected drops (panics)
+ let mut guard = StreamDriverGuard::new(handler);
+
+ let _result = drive_stream_inner(stream, handler, &mut producer.0, state).await;
+
+ // Normal completion (success or error) - disarm the guard and release handler
+ guard.disarm();
+ release_handler(handler);
+}
+
+/// Inner implementation that returns a Result for cleaner control flow.
+async fn drive_stream_inner(
+ mut stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ producer: &mut FFI_ArrowAsyncProducer,
+ state: Arc<ProducerState>,
+) -> Result<(), ()> {
+ let handler_ref = unsafe { &mut *handler };
+
+ // Set the producer on the handler BEFORE calling on_schema (per Arrow spec)
+ handler_ref.producer = producer;
+
+ // 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_ref, 1, &e.to_string());
+ return Err(());
+ }
+ };
+
+ if let Some(on_schema) = handler_ref.on_schema {
+ let ret = unsafe { on_schema(handler, &mut ffi_schema) };
+ if ret != 0 {
+ call_on_error(handler_ref, ret, "on_schema callback failed");
+ return Err(());
+ }
+ // The handler now owns the schema
+ std::mem::forget(ffi_schema);
+ }
+
+ // 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_ref, handler);
+ return Ok(());
+ }
+
+ // Wait for consumer to request batches (back-pressure)
+ // Use poll_fn to properly register our waker with ProducerState
+ futures::future::poll_fn(|cx| {
+ if state.has_requests() || state.is_cancelled() {
+ Poll::Ready(())
+ } else {
+ state.register_waker(cx.waker().clone());
+ Poll::Pending
+ }
+ })
+ .await;
+
+ // Re-check cancellation after waiting
+ if state.is_cancelled() {
+ signal_end_of_stream(handler_ref, handler);
+ return Ok(());
+ }
+
+ // Consume one request slot
+ if !state.try_consume_request() {
+ continue;
+ }
+
+ // Get next batch
+ match stream.next().await {
Review Comment:
Could this also wait for cancellation while `next()` is pending? I connected a source that never returns a batch, requested data, then called `consumer.cancel()`. Even after polling the driver again, it stayed **`Pending`**, and the source's drop count was **0**. It looks like cancellation only gets checked once the source yields, so a stalled read can keep the producer task and its resources alive indefinitely.
<details>
<summary>Repro</summary>
Save this as `c/sedona-extension/tests/pr1333_cancel_pending.rs` and run from the repository root:
```sh
cargo test -p sedona-extension --test pr1333_cancel_pending -- --nocapture
```
```rust
use arrow_array::RecordBatch;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion_common::Result;
use datafusion_execution::RecordBatchStream;
use futures::Stream;
use sedona_extension::{export_sendable_record_batch_stream::drive_stream_to_handler, import_sendable_record_batch_stream::ImportedAsyncDeviceStream};
use std::{future::Future, pin::Pin, sync::{atomic::{AtomicUsize, Ordering}, Arc}, task::{Context, Poll}};
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);
}
}
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Int32,
false,
)]))
}
#[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(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()); // on_schema; waits for request
assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending()); // requests two batches
assert!(driver.as_mut().poll(&mut cx).is_pending()); // source next() is pending
assert_eq!(polls.load(Ordering::SeqCst), 1);
consumer.cancel();
let after_cancel = driver.as_mut().poll(&mut cx);
println!(
"driver after cancel={after_cancel:?}, source polls={}, source drops={}",
polls.load(Ordering::SeqCst),
drops.load(Ordering::SeqCst)
);
assert!(
after_cancel.is_ready(),
"cancel should finish the driver even if the source never returns a batch"
);
assert_eq!(drops.load(Ordering::SeqCst), 1);
}
```
</details>
##########
c/sedona-extension/src/import_sendable_record_batch_stream.rs:
##########
@@ -0,0 +1,1016 @@
+// 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::{RecordBatch, StructArray};
+use arrow_schema::{ArrowError, DataType, Schema, SchemaRef};
+use datafusion_common::Result;
+use datafusion_execution::RecordBatchStream;
+use futures::channel::mpsc;
+use futures::task::AtomicWaker;
+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),
+ /// End of stream (NULL task received).
+ EndOfStream,
+ /// 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,
+}
+
+// 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>,
+ /// Waker for the stream (lock-free).
+ waker: AtomicWaker,
+ /// Whether the stream has ended.
+ ended: AtomicBool,
+ /// 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,
+ waker: AtomicWaker::new(),
+ ended: AtomicBool::new(false),
+ 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,
+ }),
+ }
+ }
+
+ fn wake(&self) {
+ self.waker.wake();
+ }
+
+ /// Request more data from the producer if needed.
+ fn maybe_request_more(&self) {
+ if self.ended.load(Ordering::Acquire) {
+ return;
+ }
+
+ // 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() {
+ return;
+ }
+
+ let pending = self.pending_requests.load(Ordering::Acquire);
+ let prefetch = producer_state.prefetch_count;
+
+ // Request more when we're running low
+ if pending < prefetch / 2 {
+ 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);
+ }
+ }
+ }
+
+ fn set_producer(&self, producer: *mut FFI_ArrowAsyncProducer) {
+ // Lock cannot be poisoned: we never panic while holding it
+ let mut state = self
+ .producer_state
+ .lock()
+ .expect("producer_state mutex poisoned");
+ state.producer = producer;
+ self.producer_connected.store(true, Ordering::Release);
+ }
+
+ 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 state = self
+ .producer_state
+ .lock()
+ .expect("producer_state mutex poisoned");
+ 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.ended.store(true, Ordering::Release);
+
+ 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) {
+ 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.
+ 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>>> {
+ // Register waker (lock-free)
+ self.state.waker.register(cx.waker());
+
+ // Poll the channel for messages first (lock-free)
+ // This ensures we drain all messages before returning None due to ended flag
+ 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::EndOfStream => {
+ self.state.ended.store(true, Ordering::Release);
+ Poll::Ready(None)
+ }
+ StreamMessage::Error(e) => {
+ self.state.ended.store(true, Ordering::Release);
+ Poll::Ready(Some(Err(e.into())))
+ }
+ },
+ Poll::Ready(None) => {
+ // Channel closed (shouldn't happen in normal operation)
+ self.state.ended.store(true, Ordering::Release);
+ Poll::Ready(None)
+ }
+ Poll::Pending => {
+ // No messages in channel - check if stream has ended (lock-free)
+ if self.state.ended.load(Ordering::Acquire) {
+ Poll::Ready(None)
+ } else {
+ 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) };
+ 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();
+ Ok(struct_array.into())
Review Comment:
Could we build this batch with the schema received in `on_schema`? I roundtripped one batch with schema metadata `{"review_source": "table-v1"}`. `stream.schema()` kept that map, but the returned batch's metadata was **`{}`**. It looks like converting the `StructArray` directly reconstructs the schema from just the fields; the synchronous reader preserves the original schema.
<details>
<summary>Repro</summary>
Save this as `c/sedona-extension/tests/pr1333_schema_metadata.rs` and run from the repository root:
```sh
cargo test -p sedona-extension --test pr1333_schema_metadata -- --nocapture
```
```rust
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion_common::Result;
use datafusion_execution::RecordBatchStream;
use futures::{Stream, StreamExt};
use sedona_extension::{export_sendable_record_batch_stream::drive_stream_to_handler, import_sendable_record_batch_stream::ImportedAsyncDeviceStream};
use std::{collections::HashMap, pin::Pin, sync::Arc, task::{Context, Poll}};
struct ReadyStream {
schema: SchemaRef,
batches: Vec<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.batches.pop().map(Ok))
}
}
impl RecordBatchStream for ReadyStream {
fn schema(&self) -> SchemaRef {
self.schema.clone()
}
}
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Int32,
false,
)]))
}
#[tokio::test]
async fn async_roundtrip_preserves_schema_metadata() {
let mut metadata = HashMap::new();
metadata.insert("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(),
batches: vec![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,
drive_stream_to_handler(source, handler.as_ptr())
);
assert_eq!(actual_schema, schema);
println!(
"stream schema metadata={:?}; batch schema metadata={:?}",
actual_schema.metadata(),
batches[0].schema().metadata()
);
assert_eq!(batches[0].schema(), schema);
}
```
</details>
##########
c/sedona-extension/src/export_sendable_record_batch_stream.rs:
##########
@@ -0,0 +1,472 @@
+// 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::ptr::null_mut;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::Arc;
+use std::task::{Poll, Waker};
+
+use arrow_array::ffi::{to_ffi, FFI_ArrowSchema};
+use arrow_array::RecordBatch;
+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: std::sync::Mutex<Option<Waker>>,
+}
+
+impl ProducerState {
+ fn new() -> Self {
+ Self {
+ requested: AtomicU64::new(0),
+ cancelled: AtomicBool::new(false),
+ waker: std::sync::Mutex::new(None),
+ }
+ }
+
+ 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 register_waker(&self, waker: Waker) {
+ // Lock cannot be poisoned: we never panic while holding it
+ *self.waker.lock().expect("waker mutex poisoned") = Some(waker);
+ }
+
+ fn wake(&self) {
+ // Lock cannot be poisoned: we never panic while holding it
+ if let Some(waker) = self.waker.lock().expect("waker mutex poisoned").take() {
+ 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 an error occurs (including cancellation)
+///
+/// # Safety
+///
+/// The `handler` pointer must be valid and point to a properly initialized
+/// [`FFI_ArrowAsyncDeviceStreamHandler`].
+///
+/// # 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
+/// drive_stream_to_handler(stream, handler).await;
+/// ```
+pub async fn drive_stream_to_handler(
+ stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+) {
+ if handler.is_null() {
+ return;
+ }
+
+ // Create shared state for producer callbacks
+ let state = Arc::new(ProducerState::new());
+
+ // Create the producer
+ let mut producer = ProducerGuard(create_producer(Arc::clone(&state)));
+
+ // Use a guard to handle unexpected drops (panics)
+ let mut guard = StreamDriverGuard::new(handler);
+
+ let _result = drive_stream_inner(stream, handler, &mut producer.0, state).await;
+
+ // Normal completion (success or error) - disarm the guard and release handler
+ guard.disarm();
+ release_handler(handler);
+}
+
+/// Inner implementation that returns a Result for cleaner control flow.
+async fn drive_stream_inner(
+ mut stream: SendableRecordBatchStream,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ producer: &mut FFI_ArrowAsyncProducer,
+ state: Arc<ProducerState>,
+) -> Result<(), ()> {
+ let handler_ref = unsafe { &mut *handler };
+
+ // Set the producer on the handler BEFORE calling on_schema (per Arrow spec)
+ handler_ref.producer = producer;
+
+ // 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_ref, 1, &e.to_string());
+ return Err(());
+ }
+ };
+
+ if let Some(on_schema) = handler_ref.on_schema {
+ let ret = unsafe { on_schema(handler, &mut ffi_schema) };
+ if ret != 0 {
+ call_on_error(handler_ref, ret, "on_schema callback failed");
+ return Err(());
+ }
+ // The handler now owns the schema
+ std::mem::forget(ffi_schema);
+ }
+
+ // 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_ref, handler);
+ return Ok(());
+ }
+
+ // Wait for consumer to request batches (back-pressure)
+ // Use poll_fn to properly register our waker with ProducerState
+ futures::future::poll_fn(|cx| {
+ if state.has_requests() || state.is_cancelled() {
+ Poll::Ready(())
+ } else {
+ state.register_waker(cx.waker().clone());
+ Poll::Pending
+ }
+ })
+ .await;
+
+ // Re-check cancellation after waiting
+ if state.is_cancelled() {
+ signal_end_of_stream(handler_ref, handler);
+ return Ok(());
+ }
+
+ // Consume one request slot
+ if !state.try_consume_request() {
+ continue;
+ }
+
+ // Get next batch
+ match stream.next().await {
+ Some(Ok(batch)) => {
+ if let Err(e) = send_batch_to_handler(handler_ref, handler, batch) {
+ call_on_error(handler_ref, 1, &e);
+ 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_ref, 1, &e.to_string());
+ return Err(());
+ }
+ None => {
+ // End of stream
+ signal_end_of_stream(handler_ref, handler);
+ 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();
+}
+
+struct ProducerGuard(FFI_ArrowAsyncProducer);
+
+impl Drop for ProducerGuard {
+ fn drop(&mut self) {
+ unsafe { producer_release(&mut self.0) };
+ }
+}
+
+/// Guard that calls `on_error` if dropped without being disarmed.
+struct StreamDriverGuard {
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ armed: bool,
+}
+
+impl StreamDriverGuard {
+ fn new(handler: *mut FFI_ArrowAsyncDeviceStreamHandler) -> Self {
+ Self {
+ handler,
+ armed: true,
+ }
+ }
+
+ fn disarm(&mut self) {
+ self.armed = false;
+ }
+}
+
+impl Drop for StreamDriverGuard {
+ fn drop(&mut self) {
+ if self.armed && !self.handler.is_null() {
+ let handler_ref = unsafe { &mut *self.handler };
+ call_on_error(
+ handler_ref,
+ libc::ECANCELED,
+ "Stream driver dropped unexpectedly",
+ );
+ // Still need to release the handler
+ release_handler(self.handler);
+ }
+ }
+}
+
+/// Signal end of stream by calling on_next_task with NULL.
+fn signal_end_of_stream(
+ handler_ref: &mut FFI_ArrowAsyncDeviceStreamHandler,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+) {
+ if let Some(on_next_task) = handler_ref.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());
+ }
+ }
+}
+
+/// Release the handler (call its release callback).
+fn release_handler(handler: *mut FFI_ArrowAsyncDeviceStreamHandler) {
+ if handler.is_null() {
+ return;
+ }
+ let handler_ref = unsafe { &*handler };
+ if let Some(release) = handler_ref.release {
+ unsafe { release(handler) };
+ }
+}
+
+fn send_batch_to_handler(
+ handler_ref: &mut FFI_ArrowAsyncDeviceStreamHandler,
+ handler: *mut FFI_ArrowAsyncDeviceStreamHandler,
+ batch: RecordBatch,
+) -> Result<(), String> {
+ let Some(on_next_task) = handler_ref.on_next_task else {
+ return Err("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)?;
+ let ret = unsafe { on_next_task(handler, &mut task, std::ptr::null()) };
+ if ret != 0 {
+ return Err("on_next_task callback failed".to_string());
+ }
+ std::mem::forget(task);
+ Ok(())
+}
+
+fn call_on_error(handler: &mut FFI_ArrowAsyncDeviceStreamHandler, code: c_int, message: &str) {
+ if let Some(on_error) = handler.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() || out.is_null() {
+ return 1;
+ }
+
+ let task = &mut *self_;
+ if task.private_data.is_null() {
+ return 1;
+ }
+
+ let private = &mut *(task.private_data as *mut AsyncTaskPrivate);
+
+ // 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) {
+ if !self.private_data.is_null() {
+ // Drop the private data
+ let _ = unsafe { Box::from_raw(self.private_data as *mut AsyncTaskPrivate) };
Review Comment:
I'm still getting familiar with this interface, so I may be missing something: can we use the task's `extract_data` callback to clean it up here? This type also holds tasks from other producers, whose `private_data` might not be a `Box<AsyncTaskPrivate>` (or might already have been freed by `extract_data`).
I queued a foreign task after `request(2)` and dropped the imported stream. Its cleanup callback was called **0 times**, rather than once with a null output pointer. I also tried `extract_data(task, NULL)` on an exported task: it returned `1` and retained `private_data`. The [Arrow contract](https://arrow.apache.org/docs/format/CDeviceDataInterface.html#the-arrowasynctask-structure) puts cleanup in that callback, including when the data is unwanted. I haven't reproduced an invalid free; the cast here is what worries me about a different producer's allocation.
<details>
<summary>Repro</summary>
Save this as `c/sedona-extension/tests/pr1333_task_cleanup.rs` and run from the repository root:
```sh
cargo test -p sedona-extension --test pr1333_task_cleanup -- --nocapture
```
```rust
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use futures::Stream;
use sedona_extension::import_sendable_record_batch_stream::ImportedAsyncDeviceStream;
use std::{pin::Pin, sync::{atomic::{AtomicUsize, Ordering}, Arc}, task::Context};
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Int32,
false,
)]))
}
static FOREIGN_TASK_CLEANUPS: AtomicUsize = AtomicUsize::new(0);
static FOREIGN_REQUESTS: AtomicUsize = AtomicUsize::new(0);
unsafe extern "C" fn foreign_request(
_: *mut sedona_extension::extension::FFI_ArrowAsyncProducer,
n: i64,
) {
FOREIGN_REQUESTS.fetch_add(n as usize, Ordering::SeqCst);
}
unsafe extern "C" fn foreign_cancel(_: *mut sedona_extension::extension::FFI_ArrowAsyncProducer) {}
unsafe extern "C" fn foreign_extract(
_: *mut sedona_extension::extension::FFI_ArrowAsyncTask,
out: *mut sedona_extension::extension::FFI_ArrowDeviceArray,
) -> std::ffi::c_int {
assert!(
out.is_null(),
"this test abandons the queued task without requesting its data"
);
FOREIGN_TASK_CLEANUPS.fetch_add(1, Ordering::SeqCst);
0
}
#[test]
fn abandoning_a_foreign_task_calls_extract_null() {
use sedona_extension::extension::{
FFI_ArrowAsyncProducer, FFI_ArrowAsyncTask, ARROW_DEVICE_CPU,
};
FOREIGN_TASK_CLEANUPS.store(0, Ordering::SeqCst);
FOREIGN_REQUESTS.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
);
std::mem::forget(ffi_schema); // on_schema takes ownership on this PR
let mut cx = Context::from_waker(futures::task::noop_waker_ref());
assert!(Pin::new(&mut consumer).poll_next(&mut cx).is_pending());
assert_eq!(FOREIGN_REQUESTS.load(Ordering::SeqCst), 2);
// This producer owns its resources elsewhere; private_data is deliberately
// null, so the reproducer doesn't dereference a foreign allocation as Rust.
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);
}
println!(
"foreign extract_data(NULL) cleanup calls={}",
FOREIGN_TASK_CLEANUPS.load(Ordering::SeqCst)
);
assert_eq!(
FOREIGN_TASK_CLEANUPS.load(Ordering::SeqCst),
1,
"the Arrow task must be cleaned by its producer callback when the consumer drops it"
);
}
```
</details>
##########
c/sedona-extension/src/import_sendable_record_batch_stream.rs:
##########
@@ -0,0 +1,1016 @@
+// 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::{RecordBatch, StructArray};
+use arrow_schema::{ArrowError, DataType, Schema, SchemaRef};
+use datafusion_common::Result;
+use datafusion_execution::RecordBatchStream;
+use futures::channel::mpsc;
+use futures::task::AtomicWaker;
+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),
+ /// End of stream (NULL task received).
+ EndOfStream,
+ /// 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,
+}
+
+// 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>,
+ /// Waker for the stream (lock-free).
+ waker: AtomicWaker,
+ /// Whether the stream has ended.
+ ended: AtomicBool,
+ /// 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,
+ waker: AtomicWaker::new(),
+ ended: AtomicBool::new(false),
+ 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,
+ }),
+ }
+ }
+
+ fn wake(&self) {
+ self.waker.wake();
+ }
+
+ /// Request more data from the producer if needed.
+ fn maybe_request_more(&self) {
+ if self.ended.load(Ordering::Acquire) {
+ return;
+ }
+
+ // 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() {
+ return;
+ }
+
+ let pending = self.pending_requests.load(Ordering::Acquire);
+ let prefetch = producer_state.prefetch_count;
+
+ // Request more when we're running low
+ if pending < prefetch / 2 {
Review Comment:
Should this handle a prefetch count of one? With `ImportedAsyncDeviceStream::new(1)`, the initial condition is `0 < 1 / 2`, which is false, so no request is sent. In the repro below, both sides stayed pending and the source was polled **0 times**. The current execution-plan path uses 2, but the public constructor accepts 1.
<details>
<summary>Repro</summary>
Save this as `c/sedona-extension/tests/pr1333_prefetch_one.rs` and run from the repository root:
```sh
cargo test -p sedona-extension --test pr1333_prefetch_one -- --nocapture
```
```rust
use arrow_array::RecordBatch;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion_common::Result;
use datafusion_execution::RecordBatchStream;
use futures::Stream;
use sedona_extension::{export_sendable_record_batch_stream::drive_stream_to_handler, import_sendable_record_batch_stream::ImportedAsyncDeviceStream};
use std::{future::Future, pin::Pin, sync::{atomic::{AtomicUsize, Ordering}, Arc}, task::{Context, Poll}};
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);
}
}
fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![Field::new(
"value",
DataType::Int32,
false,
)]))
}
#[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(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());
println!("prefetch=1: source polls={}", polls.load(Ordering::SeqCst));
assert_eq!(
polls.load(Ordering::SeqCst),
1,
"prefetch 1 must request at least one batch"
);
}
```
</details>
##########
c/sedona-extension/src/sedona_extension.h:
##########
@@ -110,6 +110,53 @@ struct ArrowArrayStream {
};
#endif // ARROW_C_STREAM_INTERFACE
+
+#ifndef ARROW_C_DEVICE_DATA_INTERFACE
Review Comment:
Could these new device/async definitions go outside the existing `ARROW_FLAG_DICTIONARY_ORDERED` guard? Including nanoarrow first defines that macro and skips this whole block, even though it hasn't defined `ArrowAsyncDeviceStreamHandler`.
From the repository root, this fails with **incomplete type `struct ArrowAsyncDeviceStreamHandler`**. Swapping the two includes passes:
```sh
cat > /tmp/pr1333-header.c <<'C'
#include "c/sedona-geoarrow-c/src/nanoarrow/nanoarrow.h"
#include "c/sedona-extension/src/sedona_extension.h"
int main(void) {
return sizeof(struct ArrowAsyncDeviceStreamHandler) == 0;
}
C
clang -I . -fsyntax-only /tmp/pr1333-header.c
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org