vx_array_sink->vx_writer, parallel writes

lang/cpp/include/vortex/writer.hpp

@@ -18,6 +18,8 @@ namespace vortex {
  *
  * finish() writes the footer and finalizes the file.
  * Not calling finish() leaves file corrupted.
+ *
+ * Writer methods are thread-safe.
  */
 class Writer {
 public:
@@ -41,12 +43,12 @@ class Writer {
     void 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

@@ -18,14 +18,13 @@ 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) {
     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), &error);
     throw_on_error(error);
     return Writer(sink);
 }
@@ -35,7 +34,7 @@ void Writer::push(const Array &array) {
         throw VortexException("null handle_", ErrorCode::InvalidArgument);
     }
     vx_error *error = nullptr;
-    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);
 }
 
@@ -44,7 +43,8 @@ void Writer::finish() {
         throw VortexException("finish() called twice", ErrorCode::InvalidArgument);
     }
     vx_error *error = nullptr;
-    vx_array_sink_close(handle_.release(), &error);
+    vx_writer_close(handle_.get(), &error);
+    handle_.reset();
     throw_on_error(error);
 }
 } // namespace vortex

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"] }

vortex-ffi/cinclude/vortex.h

@@ -473,23 +473,6 @@ typedef struct vx_array vx_array;
  */
 typedef struct vx_array_iterator vx_array_iterator;
 
-/**
- * The `sink` interface is used to collect array chunks and place them into a resource
- * (e.g. an array stream or file (`vx_array_sink_open_file`)).
- *
- * ## Thread Safety
- *
- * This struct is **not** thread-safe for concurrent operations. While the underlying
- * `Sender` is thread-safe, the FFI wrapper should only be accessed from a single thread
- * to avoid race conditions between `push` and `close` operations. The `close` operation
- * consumes the sink, making any subsequent operations undefined behavior.
- *
- * Multiple threads may safely hold pointers to the same sink, but only one thread should
- * perform operations on it at a time, and coordination is required to ensure `close` is
- * called exactly once after all `push` operations are complete.
- */
-typedef struct vx_array_sink vx_array_sink;
-
 /**
  * A reference to one or more (possibly remote) paths.
  * Creating vx_data_source opens the first matched path to read the schema.
@@ -525,11 +508,6 @@ typedef struct vx_error vx_error;
  */
 typedef struct vx_expression vx_expression;
 
-/**
- * A handle to a Vortex file encapsulating the footer and logic for instantiating a reader.
- */
-typedef struct vx_file vx_file;
-
 /**
  * A partition is an independent unit of work. Call vx_partition_next repeatedly to
  * retrieve arrays, then free the partition with vx_partition_free.
@@ -569,6 +547,11 @@ typedef struct vx_struct_fields vx_struct_fields;
  */
 typedef struct vx_struct_fields_builder vx_struct_fields_builder;
 
+/**
+ * vx_writer can be used to write vx_arrays into a (possibly remote) file.
+ */
+typedef struct vx_writer vx_writer;
+
 /**
  * Array validity descriptor used by C FFI constructors.
  */
@@ -1619,44 +1602,6 @@ vx_session *vx_session_new(void);
  */
 vx_session *vx_session_clone(const vx_session *session);
 
-/**
- * Increase reference count on vx_file
- */
-const vx_file *vx_file_clone(const vx_file *ptr);
-
-/**
- * Decrease reference count on vx_file or free if there are no other references
- */
-void vx_file_free(const vx_file *ptr);
-
-/**
- * Opens a writable array stream, where sink is used to push values into the stream.
- * To close the stream close the sink with `vx_array_sink_close`.
- * "path" is copied.
- */
-vx_array_sink *
-vx_array_sink_open_file(const vx_session *session, vx_view path, const vx_dtype *dtype, vx_error **error_out);
-
-/**
- * Push an array into a file sink.
- * Does not take ownership of array.
- *
- * Errors if array's DType doesn't match sink's DType.
- */
-void vx_array_sink_push(vx_array_sink *sink, const vx_array *array, vx_error **error_out);
-
-/**
- * Closes an array sink, must be called to ensure all the values pushed to the sink are written
- * to the external resource.
- */
-void vx_array_sink_close(vx_array_sink *sink, vx_error **error_out);
-
-/**
- * Abort an array sink. File footer is not written, and file is left invalid.
- * Don't use sink after this call.
- */
-void vx_array_sink_abort(vx_array_sink *sink);
-
 /**
  * Free a vx_struct_column_builder
  */
@@ -1760,6 +1705,50 @@ void vx_struct_fields_builder_add_field(vx_struct_fields_builder *builder,
  */
 vx_struct_fields *vx_struct_fields_builder_finalize(vx_struct_fields_builder *builder);
 
+/**
+ * Open a writer for a file at "path". "path" is copied.
+ *
+ * "dtype" is used to validate pushed arrays so they would all have the same
+ * schema.
+ */
+vx_writer *
+vx_writer_open(const vx_session *session, vx_view path, const vx_dtype *dtype, vx_error **error_out);
+
+/**
+ * Push an array into a writer. Does not take ownership of array.
+ *
+ * Errors if array's dtype and writer's initialized dtype are different.
+ * Errors if writer has already been closed.
+ *
+ * Thread safe.
+ */
+void vx_writer_push(vx_writer *writer, const vx_array *array, vx_error **error_out);
+
+/**
+ * Close a writer.
+ *
+ * Call to ensure all values pushed to the writer are indeed written. This
+ * call writes the footer to the file. If you don't call this function, file
+ * will be left corrupted.
+ *
+ * Thread-unsafe.
+ *
+ * Errors if writer was already closed.
+ *
+ * Use vx_writer_free to free the writer afterwards.
+ */
+void vx_writer_close(vx_writer *writer, vx_error **error_out);
+
+/**
+ * Release the writer.
+ *
+ * Thread unsafe. Must be called exactly once.
+ *
+ * If vx_writer_close wasn't called before this function, file is left
+ * corrupted.
+ */
+void vx_writer_free(vx_writer *sink);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif // __cplusplus

vortex-ffi/examples/write_sample.c

@@ -106,7 +106,7 @@ int main(int argc, char *argv[]) {
     const vx_dtype *dtype = sample_dtype();
 
     vx_error *error = NULL;
-    vx_array_sink *sink = vx_array_sink_open_file(session, vx_view_from_cstr(output), dtype, &error);
+    vx_writer *writer = vx_writer_open(session, vx_view_from_cstr(output), dtype, &error);
 
     vx_dtype_free(dtype);
     if (error != NULL) {
@@ -118,24 +118,27 @@ int main(int argc, char *argv[]) {
     if (array == NULL) {
         // We already have an error, so we can ignore a potential error
         // from this operation
-        vx_array_sink_close(sink, &error);
+        vx_writer_close(writer, &error);
+        vx_writer_free(writer);
         vx_session_free(session);
         return 1;
     }
 
-    vx_array_sink_push(sink, array, &error);
+    vx_writer_push(writer, array, &error);
     if (error != NULL) {
-        vx_array_sink_close(sink, &error);
+        vx_writer_close(writer, &error);
+        vx_writer_free(writer);
         vx_session_free(session);
         return 1;
     }
     vx_array_free(array);
 
-    vx_array_sink_close(sink, &error);
+    vx_writer_close(writer, &error);
     if (error != NULL) {
         print_error("Error closing output sink", error);
         vx_error_free(error);
     }
+    vx_writer_free(writer);
 
     vx_session_free(session);
     return 0;

vortex-ffi/src/lib.rs

@@ -17,10 +17,10 @@ mod ptype;
 mod scalar;
 mod scan;
 mod session;
-mod sink;
 mod string;
 mod struct_array;
 mod struct_fields;
+mod writer;
 
 use std::sync::Arc;
 use std::sync::LazyLock;
@@ -100,10 +100,11 @@ mod tests {
     use crate::error::vx_error_free;
     use crate::error::vx_error_message;
     use crate::session::vx_session;
-    use crate::sink::vx_array_sink_close;
-    use crate::sink::vx_array_sink_open_file;
-    use crate::sink::vx_array_sink_push;
     use crate::string::vx_view;
+    use crate::writer::vx_writer_free;
+    use crate::writer::vx_writer_close;
+    use crate::writer::vx_writer_open;
+    use crate::writer::vx_writer_push;
 
     /// Panic if error is NULL. Free the error if it's not
     pub(crate) fn assert_error(error: *mut vx_error) {
@@ -164,10 +165,11 @@ mod tests {
         unsafe {
             let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype.clone()));
             let mut error = ptr::null_mut();
-            let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error);
+            let sink = vx_writer_open(session, path, vx_dtype_ptr, &raw mut error);
             let array = vx_array::new(Arc::new(struct_array.clone().into_array()));
-            vx_array_sink_push(sink, array, &raw mut error);
-            vx_array_sink_close(sink, &raw mut error);
+            vx_writer_push(sink, array, &raw mut error);
+            vx_writer_close(sink, &raw mut error);
+            vx_writer_free(sink);
             vx_array_free(array);
             vx_dtype_free(vx_dtype_ptr);
         }