FFI: Thread-safe writer

lang/cpp/include/vortex/writer.hpp

@@ -10,20 +10,38 @@
 
 #include <memory>
 #include <string_view>
+#include <thread>
 
 namespace vortex {
 
+// Summary of a written Vortex file
+struct WriteSummary {
+    uint64_t row_count;
+    // File size in bytes
+    uint64_t file_size;
+};
+
 /**
- * Writes arrays into a Vortex file.
+ * Write Arrays into a .vortex file.
  *
  * finish() writes the footer and finalizes the file.
  * Not calling finish() leaves file corrupted.
  *
- * Writer methods are thread-unsafe.
+ * Writer methods are thread-safe.
  */
 class Writer {
 public:
-    static Writer open(const Session &session, std::string_view path, const DataType &dtype);
+    /*
+     * Open a writer for a file at "path". "path" is copied.
+     * "dtype" is used to validate pushed arrays so they all have same schema.
+     *
+     * "concurrent_array_limit" is an upper limit on how many pushed arrays are
+     * buffered before push() blocks. This caps RAM used for buffering.
+     */
+    static Writer open(const Session &session,
+                       std::string_view path,
+                       const DataType &dtype,
+                       size_t concurrent_array_limit = std::thread::hardware_concurrency());
 
     Writer(const Writer &) = delete;
     Writer &operator=(const Writer &) = delete;
@@ -42,15 +60,15 @@ class Writer {
      * Write footer and finalize the file.
      * Throws on failure. Writer is closed afterwards and further uses throws.
      */
-    void finish();
+    WriteSummary finish();
 
 private:
-    explicit Writer(vx_array_sink *sink) : handle_(sink) {
+    explicit Writer(vx_writer *writer) : handle_(writer) {
     }
 
     struct Deleter {
-        void operator()(vx_array_sink *ptr) const noexcept;
+        void operator()(vx_writer *ptr) const noexcept;
     };
-    std::unique_ptr<vx_array_sink, Deleter> handle_;
+    std::unique_ptr<vx_writer, Deleter> handle_;
 };
 } // namespace vortex

lang/cpp/src/writer.cpp

@@ -19,25 +19,31 @@ using detail::Access;
 using detail::throw_on_error;
 using detail::to_view;
 
-void Writer::Deleter::operator()(vx_array_sink *ptr) const noexcept {
-    vx_array_sink_abort(ptr);
+void Writer::Deleter::operator()(vx_writer *ptr) const noexcept {
+    vx_writer_free(ptr);
 }
 
-Writer Writer::open(const Session &session, std::string_view path, const DataType &dtype) {
+Writer Writer::open(const Session &session,
+                    std::string_view path,
+                    const DataType &dtype,
+                    size_t concurrent_array_limit) {
     vx_error *error = nullptr;
-    vx_array_sink *sink =
-        vx_array_sink_open_file(Access::c_ptr(session), to_view(path), Access::c_ptr(dtype), &error);
+    vx_writer *sink = vx_writer_open(Access::c_ptr(session),
+                                     to_view(path),
+                                     Access::c_ptr(dtype),
+                                     concurrent_array_limit,
+                                     &error);
     throw_on_error(error);
     return Writer(sink);
 }
 
 void Writer::push(std::span<const Array> arrays) {
     if (handle_ == nullptr) {
-        throw VortexException("null handle_", ErrorCode::InvalidArgument);
+        throw VortexException("nullptr handle_", ErrorCode::InvalidArgument);
     }
     vx_error *error = nullptr;
     for (const Array &array : arrays) {
-        vx_array_sink_push(handle_.get(), Access::c_ptr(array), &error);
+        vx_writer_push(handle_.get(), Access::c_ptr(array), &error);
         throw_on_error(error);
     }
 }
@@ -50,12 +56,15 @@ void Writer::push(std::initializer_list<Array> arrays) {
     push({arrays.begin(), arrays.end()});
 }
 
-void Writer::finish() {
+WriteSummary Writer::finish() {
     if (handle_ == nullptr) {
-        throw VortexException("finish() called twice", ErrorCode::InvalidArgument);
+        throw VortexException("nullptr handle_", ErrorCode::InvalidArgument);
     }
     vx_error *error = nullptr;
-    vx_array_sink_close(handle_.release(), &error);
+    vx_write_summary summary = {};
+    vx_writer_close(handle_.get(), &summary, &error);
+    handle_.reset(); // do this so further finish() or push() throws
     throw_on_error(error);
+    return {summary.row_count, summary.file_size};
 }
 } // namespace vortex

vortex-cuda/ffi/README.md

@@ -17,8 +17,8 @@ and link the CUDA FFI library (`vortex_cuda_ffi`). Do not pass Vortex handles be
 
 Use `vx_cuda_session_new` to initialize CUDA once and reuse it across exports.
 
-Use `vx_cuda_array_sink_open_file` to open a standard Vortex file sink configured to produce CUDA-readable files.
-Push host-resident arrays and close the sink using the standard `vx_array_sink_*` APIs.
+Use `vx_cuda_writer_open` to open a standard Vortex file writer configured to produce CUDA-readable files.
+Push host-resident arrays and close the writer using the standard `vx_writer_*` APIs.
 
 Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers
 and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer

vortex-cuda/ffi/cinclude/vortex_cuda.h

@@ -65,34 +65,36 @@ struct ArrowDeviceArrayStream {
 vx_session *vx_cuda_session_new(vx_error **error_out);
 
 /**
- * Open a Vortex file sink configured to produce CUDA-readable files.
+ * Open a Vortex file writer configured to produce CUDA-readable files.
  *
- * Push host-resident arrays and close or abort the returned sink with the standard
- * `vx_array_sink_*` functions. This API configures the on-disk encodings and layout; it does not
+ * Push host-resident arrays and close the returned writer with the standard
+ * `vx_writer_*` functions. This API configures the on-disk encodings and layout; it does not
  * move arrays to the GPU during the write.
  */
-vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session,
-                                            vx_view path,
-                                            const vx_dtype *dtype,
-                                            vx_error **error_out);
+vx_writer *
+vx_cuda_writer_open(const vx_session *session, vx_view path, const vx_dtype *dtype, vx_error **error_out);
 
 /**
- * Open a CUDA-readable Vortex file sink with a fixed row block size.
+ * Open a CUDA-readable Vortex file writer with a fixed row block size.
  *
  * `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero uses the default
  * writer strategy: 8,192-row blocks may be coalesced into data blocks targeting 1 MiB. Passing any
  * nonzero value disables this byte-size coalescing, so passing 8,192 is not equivalent to passing
  * zero.
  *
+ * `concurrent_array_limit` is an upper limit on how many pushed arrays may be buffered before
+ * `vx_writer_push` blocks.
+ *
  * Write and scan sizing are independent. To align on-disk row blocks with scan batches, pass the
  * same nonzero value to this function and `vx_cuda_scan_path_arrow_device_stream_batch_rows`; the
  * API does not enforce a match.
  */
-vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session,
-                                                       vx_view path,
-                                                       const vx_dtype *dtype,
-                                                       size_t block_rows,
-                                                       vx_error **error_out);
+vx_writer *vx_cuda_writer_open_block_rows(const vx_session *session,
+                                          vx_view path,
+                                          const vx_dtype *dtype,
+                                          size_t block_rows,
+                                          size_t concurrent_array_limit,
+                                          vx_error **error_out);
 
 /**
  * Options for scanning a CUDA-compatible Vortex file.
@@ -113,7 +115,7 @@ typedef struct vx_cuda_scan_options {
 /**
  * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream.
  *
- * Files written by `vx_cuda_array_sink_open_file` are compatible with this path. Reusing the same
+ * Files written by `vx_cuda_writer_open` are compatible with this path. Reusing the same
  * CUDA session across calls also reuses the pinned host buffers used to stage file reads.
  *
  * On success returns 0 and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The caller
@@ -132,7 +134,7 @@ int vx_cuda_scan_path_arrow_device_stream(const vx_session *session,
  * layout-derived splitting of `vx_cuda_scan_path_arrow_device_stream`.
  *
  * Scan and write sizing are independent. To align scan batches with on-disk row blocks, pass the
- * same nonzero value to this function and `vx_cuda_array_sink_open_file_block_rows`; the API does
+ * same nonzero value to this function and `vx_cuda_writer_open_block_rows`; the API does
  * not enforce a match.
  */
 int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session,

vortex-cuda/ffi/src/lib.rs

@@ -36,8 +36,6 @@ use vortex_ffi::ffi_runtime;
 use vortex_ffi::try_or;
 use vortex_ffi::vx_array;
 use vortex_ffi::vx_array_ref;
-use vortex_ffi::vx_array_sink;
-use vortex_ffi::vx_array_sink_open_file_with_strategy;
 use vortex_ffi::vx_dtype;
 use vortex_ffi::vx_error;
 use vortex_ffi::vx_partition;
@@ -46,6 +44,8 @@ use vortex_ffi::vx_session;
 use vortex_ffi::vx_session_new_with;
 use vortex_ffi::vx_session_ref;
 use vortex_ffi::vx_view;
+use vortex_ffi::vx_writer;
+use vortex_ffi::vx_writer_open_with_strategy;
 
 const VX_CUDA_OK: c_int = 0;
 const VX_CUDA_ERR: c_int = 1;
@@ -98,31 +98,31 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new(
     })
 }
 
-/// Open a Vortex file sink configured to produce CUDA-readable files.
+/// Open a Vortex file writer configured to produce CUDA-readable files.
 ///
-/// Push host-resident arrays and close or abort the returned sink with the standard
-/// `vx_array_sink_*` functions. This function configures the on-disk encodings and layout; it does
+/// Push host-resident arrays and close the returned writer with the standard
+/// `vx_writer_*` functions. This function configures the on-disk encodings and layout; it does
 /// not move arrays to the GPU during the write.
 ///
 /// # Safety
 ///
 /// `session`, `path`, and `dtype` must satisfy the same requirements as
-/// `vx_array_sink_open_file`. If `error_out` is non-null, it must be valid for writing one error
+/// `vx_writer_open`. If `error_out` is non-null, it must be valid for writing one error
 /// pointer.
 #[unsafe(no_mangle)]
-pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file(
+pub unsafe extern "C-unwind" fn vx_cuda_writer_open(
     session: *const vx_session,
     path: vx_view,
     dtype: *const vx_dtype,
     error_out: *mut *mut vx_error,
-) -> *mut vx_array_sink {
-    unsafe { vx_cuda_array_sink_open_file_block_rows(session, path, dtype, 0, error_out) }
+) -> *mut vx_writer {
+    unsafe { vx_cuda_writer_open_block_rows(session, path, dtype, 0, 32, error_out) }
 }
 
-/// Open a CUDA-readable Vortex file sink with a fixed row block size.
+/// Open a CUDA-readable Vortex file writer with a fixed row block size.
 ///
 /// `block_rows` controls the row granularity of CUDA-flat data blocks. Passing zero preserves the
-/// default writer strategy used by [`vx_cuda_array_sink_open_file`]. Any nonzero value disables
+/// default writer strategy used by [`vx_cuda_writer_open`]. Any nonzero value disables
 /// byte-size coalescing so data blocks retain the requested row granularity.
 ///
 /// Write and scan sizing are independent. To align on-disk row blocks with scan batches, pass the
@@ -131,16 +131,17 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file(
 /// # Safety
 ///
 /// `session`, `path`, and `dtype` must satisfy the same requirements as
-/// `vx_array_sink_open_file`. If `error_out` is non-null, it must be valid for writing one error
+/// `vx_writer_open_file`. If `error_out` is non-null, it must be valid for writing one error
 /// pointer.
 #[unsafe(no_mangle)]
-pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows(
+pub unsafe extern "C-unwind" fn vx_cuda_writer_open_block_rows(
     session: *const vx_session,
     path: vx_view,
     dtype: *const vx_dtype,
     block_rows: usize,
+    concurrent_array_limit: usize,
     error_out: *mut *mut vx_error,
-) -> *mut vx_array_sink {
+) -> *mut vx_writer {
     try_or(error_out, ptr::null_mut(), || {
         session_with_cuda(unsafe { vx_session_ref(session) }?)?;
         let mut strategy = WriteStrategyBuilder::default()
@@ -155,7 +156,15 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows(
                 .with_row_block_size(block_rows)
                 .with_data_block_target_bytes(None);
         }
-        unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy.build()) }
+        unsafe {
+            vx_writer_open_with_strategy(
+                session,
+                path,
+                dtype,
+                concurrent_array_limit,
+                strategy.build(),
+            )
+        }
     })
 }
 
@@ -165,7 +174,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows(
 /// buffers and transferred directly to the GPU.
 ///
 /// The file must use encodings and layouts supported by the CUDA execution path, such as files
-/// written by [`vx_cuda_array_sink_open_file`]. Pinned staging buffers are reused across scans made
+/// written by [`vx_cuda_writer_open`]. Pinned staging buffers are reused across scans made
 /// with the same CUDA session.
 ///
 /// On success returns `0` and writes an owned [`ArrowDeviceArrayStream`] to `out_stream`. The
@@ -204,7 +213,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream(
 /// layout-derived splitting used by [`vx_cuda_scan_path_arrow_device_stream`].
 ///
 /// Scan and write sizing are independent. To align scan batches with on-disk row blocks, pass the
-/// same nonzero value to this function and [`vx_cuda_array_sink_open_file_block_rows`].
+/// same nonzero value to this function and [`vx_cuda_writer_open_block_rows`].
 ///
 /// # Safety
 ///

vortex-ffi/Cargo.toml

@@ -26,6 +26,7 @@ async-fs = { workspace = true }
 bytes = { workspace = true }
 futures = { workspace = true }
 mimalloc = { workspace = true, optional = true }
+parking_lot = { workspace = true }
 paste = { workspace = true }
 tracing = { workspace = true, features = ["log"] }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }