Document CUDA stream helpers

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

vortex-cuda/ffi/src/lib.rs

@@ -32,6 +32,7 @@ use vortex_ffi::vx_session_ref;
 const VX_CUDA_OK: c_int = 0;
 const VX_CUDA_ERR: c_int = 1;
 
+/// Return a session with CUDA state, adding default CUDA support when needed.
 fn session_with_cuda(session: &VortexSession) -> VortexResult<VortexSession> {
     if session.get_opt::<CudaSession>().is_some() {
         return Ok(session.clone());

vortex-cuda/src/arrow/mod.rs

@@ -253,16 +253,19 @@ impl DeviceArrayExt for ArrayRef {
     }
 }
 
+// POSIX EIO for Arrow stream producer/export failures.
 const ARROW_STREAM_EIO: c_int = 5;
+// POSIX EINVAL for invalid Arrow stream callback arguments or released streams.
 const ARROW_STREAM_EINVAL: c_int = 22;
 
 #[derive(Clone, Debug, PartialEq)]
-enum DeviceExportSchema {
+enum ArrowDeviceStreamSchema {
     Schema(Schema),
     Field(Field),
 }
 
-impl DeviceExportSchema {
+impl ArrowDeviceStreamSchema {
+    /// Interpret an Arrow C schema as the stream schema shape for `dtype`.
     fn from_ffi(schema: &FFI_ArrowSchema, dtype: &DType) -> VortexResult<Self> {
         if matches!(dtype, DType::Struct(..)) {
             Ok(Self::Schema(Schema::try_from(schema)?))
@@ -271,6 +274,7 @@ impl DeviceExportSchema {
         }
     }
 
+    /// Build the Arrow stream schema for an empty stream with the given dtype.
     fn from_dtype(dtype: &DType, ctx: &mut CudaExecutionCtx) -> VortexResult<Self> {
         let dtype = arrow_device_export_dtype(dtype);
         if let DType::Struct(struct_dtype, _) = &dtype {
@@ -282,6 +286,7 @@ impl DeviceExportSchema {
         }
     }
 
+    /// Export this stream schema as an owned Arrow C schema.
     fn to_ffi(&self) -> VortexResult<FFI_ArrowSchema> {
         match self {
             Self::Schema(schema) => Ok(FFI_ArrowSchema::try_from(schema)?),
@@ -292,48 +297,41 @@ impl DeviceExportSchema {
 
 type ArrayStreamIterator = Box<dyn Iterator<Item = VortexResult<ArrayRef>>>;
 
-struct FuturesArrayStreamIterator {
-    stream: SendableArrayStream,
-}
-
-impl Iterator for FuturesArrayStreamIterator {
-    type Item = VortexResult<ArrayRef>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        futures::executor::block_on(self.stream.next())
-    }
-}
-
 struct DeviceArrayStreamPrivateData {
     array_iter: ArrayStreamIterator,
     ctx: CudaExecutionCtx,
     dtype: DType,
-    schema: Option<DeviceExportSchema>,
+    schema: Option<ArrowDeviceStreamSchema>,
     pending_array: Option<ArrowDeviceArray>,
     device_id: i64,
     last_error: Option<CString>,
 }
 
 impl DeviceArrayStreamPrivateData {
+    /// Clear the last stream error before a new callback invocation.
     fn clear_error(&mut self) {
         self.last_error = None;
     }
 
+    /// Store the last stream error and return the Arrow callback error code.
     fn set_error(&mut self, error: impl ToString) -> c_int {
-        let message = error.to_string().replace('\0', "\\0");
-        self.last_error = CString::new(message).ok();
+        self.last_error = CString::new(error.to_string()).ok();
         ARROW_STREAM_EIO
     }
 
-    fn ensure_schema(&mut self) -> VortexResult<&DeviceExportSchema> {
+    /// Initialize and return the stream schema, exporting the first batch if needed.
+    fn ensure_schema(&mut self) -> VortexResult<&ArrowDeviceStreamSchema> {
         if self.schema.is_none() {
             match self.array_iter.next() {
                 Some(array) => {
                     let array = self.export_batch(array?)?;
                     self.pending_array = Some(array);
                 }
                 None => {
-                    self.schema = Some(DeviceExportSchema::from_dtype(&self.dtype, &mut self.ctx)?);
+                    self.schema = Some(ArrowDeviceStreamSchema::from_dtype(
+                        &self.dtype,
+                        &mut self.ctx,
+                    )?);
                 }
             }
         }
@@ -343,21 +341,26 @@ impl DeviceArrayStreamPrivateData {
             .ok_or_else(|| vortex_err!("ArrowDeviceArrayStream schema was not initialized"))
     }
 
+    /// Export and return the next device batch, or `None` at end of stream.
     fn next_array(&mut self) -> VortexResult<Option<ArrowDeviceArray>> {
         if let Some(array) = self.pending_array.take() {
             return Ok(Some(array));
         }
 
         let Some(array) = self.array_iter.next() else {
             if self.schema.is_none() {
-                self.schema = Some(DeviceExportSchema::from_dtype(&self.dtype, &mut self.ctx)?);
+                self.schema = Some(ArrowDeviceStreamSchema::from_dtype(
+                    &self.dtype,
+                    &mut self.ctx,
+                )?);
             }
             return Ok(None);
         };
 
         self.export_batch(array?).map(Some)
     }
 
+    /// Export one Vortex array batch and validate it against the stream schema and device.
     fn export_batch(&mut self, array: ArrayRef) -> VortexResult<ArrowDeviceArray> {
         vortex_ensure!(
             array.dtype() == &self.dtype,
@@ -372,7 +375,7 @@ impl DeviceArrayStreamPrivateData {
             mut schema,
             mut array,
         } = exported;
-        let batch_schema = DeviceExportSchema::from_ffi(&schema, &self.dtype);
+        let batch_schema = ArrowDeviceStreamSchema::from_ffi(&schema, &self.dtype);
         release_schema(&mut schema);
         let batch_schema = match batch_schema {
             Ok(batch_schema) => batch_schema,
@@ -418,6 +421,7 @@ impl DeviceArrayStreamPrivateData {
 }
 
 impl Drop for DeviceArrayStreamPrivateData {
+    /// Release any pending batch exported while initializing the stream schema.
     fn drop(&mut self) {
         if let Some(mut array) = self.pending_array.take() {
             release_device_array(&mut array);
@@ -440,13 +444,15 @@ pub trait DeviceArrayStreamExt {
 }
 
 impl DeviceArrayStreamExt for SendableArrayStream {
+    /// Export this stream by adapting it to a blocking iterator.
     fn export_device_array_stream(
         self,
         session: &VortexSession,
     ) -> VortexResult<ArrowDeviceArrayStream> {
         let dtype = self.dtype().clone();
+        let mut stream = self;
         export_device_array_stream_from_iter(
-            FuturesArrayStreamIterator { stream: self },
+            std::iter::from_fn(move || futures::executor::block_on(stream.next())),
             dtype,
             session,
         )
@@ -503,6 +509,7 @@ pub fn export_device_array_stream_from_iter_with_ctx(
     }
 }
 
+/// Return the private stream state for a live Arrow device stream.
 unsafe fn device_stream_private_data<'a>(
     stream: *mut ArrowDeviceArrayStream,
 ) -> Option<&'a mut DeviceArrayStreamPrivateData> {
@@ -515,6 +522,7 @@ unsafe fn device_stream_private_data<'a>(
     }
 }
 
+/// Create the Arrow end-of-stream marker for the stream's CUDA device.
 fn released_device_array(device_id: i64) -> ArrowDeviceArray {
     ArrowDeviceArray {
         array: ArrowArray::empty(),
@@ -525,18 +533,21 @@ fn released_device_array(device_id: i64) -> ArrowDeviceArray {
     }
 }
 
+/// Release an Arrow C schema if it is live.
 fn release_schema(schema: &mut FFI_ArrowSchema) {
     if let Some(release) = schema.release {
         unsafe { release(schema) };
     }
 }
 
+/// Release an Arrow device array if it is live.
 fn release_device_array(array: &mut ArrowDeviceArray) {
     if let Some(release) = array.array.release {
         unsafe { release(&raw mut array.array) };
     }
 }
 
+/// Implement `ArrowDeviceArrayStream.get_schema` for Vortex CUDA streams.
 unsafe extern "C" fn device_stream_get_schema(
     stream: *mut ArrowDeviceArrayStream,
     out: *mut ArrowSchema,
@@ -560,6 +571,7 @@ unsafe extern "C" fn device_stream_get_schema(
     }
 }
 
+/// Implement `ArrowDeviceArrayStream.get_next` for Vortex CUDA streams.
 unsafe extern "C" fn device_stream_get_next(
     stream: *mut ArrowDeviceArrayStream,
     out: *mut ArrowDeviceArray,
@@ -586,6 +598,7 @@ unsafe extern "C" fn device_stream_get_next(
     }
 }
 
+/// Implement `ArrowDeviceArrayStream.get_last_error` for Vortex CUDA streams.
 unsafe extern "C" fn device_stream_get_last_error(
     stream: *mut ArrowDeviceArrayStream,
 ) -> *const c_char {
@@ -599,6 +612,7 @@ unsafe extern "C" fn device_stream_get_last_error(
         .map_or(ptr::null(), |error| error.as_ptr())
 }
 
+/// Implement `ArrowDeviceArrayStream.release` for Vortex CUDA streams.
 unsafe extern "C" fn device_stream_release(stream: *mut ArrowDeviceArrayStream) {
     let Some(stream_ref) = (unsafe { stream.as_mut() }) else {
         return;
@@ -875,18 +889,21 @@ mod tests {
     use crate::arrow::ArrowSchema;
     use crate::arrow::DeviceArrayStreamExt;
 
+    /// Release an Arrow C schema in stream tests if it is live.
     unsafe fn release_schema(schema: &mut FFI_ArrowSchema) {
         if let Some(release) = schema.release {
             unsafe { release(schema) };
         }
     }
 
+    /// Release an Arrow device array in stream tests if it is live.
     unsafe fn release_device_array(array: &mut ArrowDeviceArray) {
         if let Some(release) = array.array.release {
             unsafe { release(&raw mut array.array) };
         }
     }
 
+    /// Create a zeroed placeholder Arrow device array for callback outputs.
     fn empty_device_array() -> ArrowDeviceArray {
         ArrowDeviceArray {
             array: ArrowArray::empty(),
@@ -897,6 +914,7 @@ mod tests {
         }
     }
 
+    /// Verify schema, batch, EOS, and idempotent release stream behavior.
     #[cuda_test]
     fn test_export_device_array_stream_schema_next_eos_release() -> VortexResult<()> {
         let session = VortexSession::default().with_some(CudaSession::try_default()?);

vortex-test/e2e-cuda/src/lib.rs

@@ -219,6 +219,7 @@ fn dictionary_array() -> VortexArrayRef {
     .into_array()
 }
 
+/// Build the shared cuDF interop test array used by array and stream exports.
 fn cudf_test_array() -> Result<VortexArrayRef, String> {
     let primitive = primitive_array()?;
     // cuDF supports Arrow decimal device imports through Decimal128. Decimal256 is intentionally
@@ -292,6 +293,8 @@ fn cudf_test_array() -> Result<VortexArrayRef, String> {
     .into_array())
 }
 
+/// Export the shared cuDF test array as one Arrow device array.
+///
 /// # Safety
 /// `schema_ptr` and `array_ptr` must be valid writable pointers.
 #[unsafe(no_mangle)]
@@ -302,6 +305,7 @@ pub unsafe extern "C" fn export_array(
     ffi_boundary("export_array", || export_array_inner(schema_ptr, array_ptr))
 }
 
+/// Implement `export_array` inside the panic-catching FFI boundary.
 fn export_array_inner(schema_ptr: &mut FFI_ArrowSchema, array_ptr: &mut ArrowDeviceArray) -> i32 {
     let mut ctx = match CudaSession::create_execution_ctx(&SESSION) {
         Ok(ctx) => ctx,
@@ -332,6 +336,8 @@ fn export_array_inner(schema_ptr: &mut FFI_ArrowSchema, array_ptr: &mut ArrowDev
     }
 }
 
+/// Export the shared cuDF test array as an Arrow device array stream.
+///
 /// # Safety
 /// `stream_ptr` must be a valid writable pointer.
 #[unsafe(no_mangle)]
@@ -341,6 +347,7 @@ pub unsafe extern "C" fn export_device_stream(stream_ptr: &mut ArrowDeviceArrayS
     })
 }
 
+/// Implement `export_device_stream` inside the panic-catching FFI boundary.
 fn export_device_stream_inner(stream_ptr: &mut ArrowDeviceArrayStream) -> i32 {
     let array = match cudf_test_array() {
         Ok(array) => array,