Clarify CUDA stream ownership

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

Author: 0ax1

vortex-cuda/ffi/cinclude/vortex_cuda.h

@@ -82,15 +82,17 @@ int vx_cuda_array_export_arrow_device(const vx_session *session,
                                       vx_error **error_out);
 
 /**
- * Scan a Vortex partition as an Arrow C Device stream.
+ * Consume a Vortex partition and scan it as an Arrow C Device stream.
+ *
+ * This function takes ownership of `partition`. Callers must not free or reuse
+ * it after calling this function, regardless of success or failure.
  *
  * On success returns 0 and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The stream owns
- * the partition and must be released through its embedded Arrow `release` callback. Each produced
- * `ArrowDeviceArray` is exported on one CUDA device and must be released independently by the
- * consumer through its embedded `ArrowArray.release` callback.
+ * the resulting scan iterator and must be released through its embedded Arrow `release` callback.
+ * Each produced `ArrowDeviceArray` must be released through its embedded `ArrowArray.release`
+ * callback.
  *
- * On error returns 1 and, when `error_out` is non-NULL, writes a `vx_error` (free with
- * `vx_error_free`). If `partition` is consumed by this call, callers must not free or reuse it.
+ * On error returns 1 and writes a `vx_error` to `error_out` when non-NULL.
  */
 int vx_cuda_partition_scan_arrow_device(const vx_session *session,
                                         vx_partition *partition,

vortex-cuda/ffi/src/lib.rs

@@ -103,15 +103,18 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_export_arrow_device(
     })
 }
 
-/// Scan a Vortex partition as an Arrow C Device stream.
+/// Consume a Vortex partition and scan it as an Arrow C Device stream.
+///
+/// This function takes ownership of `partition`. Callers must not free or reuse it after calling
+/// this function, regardless of success or failure.
 ///
 /// On success returns `0` and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The stream
-/// owns the partition and must be released through its embedded Arrow `release` callback. Each
-/// produced `ArrowDeviceArray` is exported on one CUDA device and is independently released by the
-/// consumer through the embedded `ArrowArray.release` callback.
+/// owns the resulting scan iterator. The caller must release the stream through its embedded Arrow
+/// `release` callback, and must release each produced `ArrowDeviceArray` through its embedded
+/// `ArrowArray.release` callback.
 ///
 /// On error returns `1` and, when `error_out` is non-null, writes a `vx_error` (free with
-/// `vx_error_free`). If `partition` is consumed by this call, callers must not free or reuse it.
+/// `vx_error_free`).
 ///
 /// # Safety
 ///
@@ -129,9 +132,9 @@ pub unsafe extern "C-unwind" fn vx_cuda_partition_scan_arrow_device(
         vortex_ensure!(!partition.is_null(), "null vx_partition");
         vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output");
 
+        let (dtype, array_iter) = unsafe { vx_partition_into_array_iter(partition) }?;
         let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?;
         let ctx = CudaSession::create_execution_ctx(&session)?;
-        let (dtype, array_iter) = unsafe { vx_partition_into_array_iter(partition) }?;
         let device_stream =
             vortex_cuda::export_device_array_stream_from_iter_with_ctx(array_iter, dtype, ctx);
 

vortex-cuda/src/arrow/mod.rs

@@ -473,9 +473,9 @@ pub fn export_device_array_stream_from_iter(
 /// Export a blocking Vortex array iterator as an [`ArrowDeviceArrayStream`] using an existing CUDA
 /// execution context.
 ///
-/// This is useful for FFI entry points that must finish all fallible CUDA initialization before
-/// consuming an owned input handle. Each yielded array must have `dtype`; every exported batch is
-/// validated to stay on the CUDA device selected by `ctx`.
+/// Use this helper when the caller has already selected the CUDA execution context that must drive
+/// the exported stream. Each yielded array must have `dtype`; every exported batch is validated to
+/// stay on the CUDA device selected by `ctx`.
 pub fn export_device_array_stream_from_iter_with_ctx(
     array_iter: impl Iterator<Item = VortexResult<ArrayRef>> + 'static,
     dtype: DType,