diff --git a/.github/workflows/rust-instrumented.yml b/.github/workflows/rust-instrumented.yml index 5430f1512fc..df566b5c936 100644 --- a/.github/workflows/rust-instrumented.yml +++ b/.github/workflows/rust-instrumented.yml @@ -252,6 +252,7 @@ jobs: ./vortex-ffi/build/examples/dtype '*.vortex' ./vortex-ffi/build/examples/scan '*.vortex' ./vortex-ffi/build/examples/scan_to_arrow '*.vortex' + ./vortex-ffi/build/examples/parallel_write -j 16 '1.vortex' miri: name: "Rust tests (miri)" diff --git a/Cargo.lock b/Cargo.lock index 74236675f9f..3ef5c0716bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10137,6 +10137,7 @@ dependencies = [ "cbindgen", "futures", "mimalloc", + "parking_lot", "paste", "rand 0.10.2", "tempfile", diff --git a/lang/cpp/include/vortex/writer.hpp b/lang/cpp/include/vortex/writer.hpp index da761c334f9..fdbd3bea755 100644 --- a/lang/cpp/include/vortex/writer.hpp +++ b/lang/cpp/include/vortex/writer.hpp @@ -18,10 +18,13 @@ namespace vortex { * * finish() writes the footer and finalizes the file. * Not calling finish() leaves file corrupted. + * + * Writer methods are thread-safe. */ class Writer { public: - static Writer open(const Session &session, std::string_view path, const DataType &dtype); + static Writer + open(const Session &session, std::string_view path, const DataType &dtype, size_t concurrent_array_limit); Writer(const Writer &) = delete; Writer &operator=(const Writer &) = delete; @@ -41,12 +44,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 handle_; + std::unique_ptr handle_; }; } // namespace vortex diff --git a/lang/cpp/src/writer.cpp b/lang/cpp/src/writer.cpp index 4adc5d77e50..00690db9e3c 100644 --- a/lang/cpp/src/writer.cpp +++ b/lang/cpp/src/writer.cpp @@ -18,14 +18,20 @@ 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); } @@ -35,7 +41,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 +50,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 diff --git a/vortex-ffi/Cargo.toml b/vortex-ffi/Cargo.toml index a64aabac03f..b0ff3a0f616 100644 --- a/vortex-ffi/Cargo.toml +++ b/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"] } diff --git a/vortex-ffi/cinclude/vortex.h b/vortex-ffi/cinclude/vortex.h index 1c2caa6fad8..a5d2279439a 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/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,63 @@ 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. + * + * "concurrent_array_limit" is the limit on the number of arrays that are + * processed concurrently. This limits RAM used for processing. + */ +vx_writer *vx_writer_open(const vx_session *session, + vx_view path, + const vx_dtype *dtype, + size_t concurrent_array_limit, + vx_error **error_out); + +/** + * Push an array into a writer. Does not take ownership of array. + * + * Array ordering across concurrent calls to this function is + * non-deterministic: vx_writer_push(array1) called concurrently with + * vx_writer_push(array2) may write array2 first. + * + * 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. + * + * If this function is called concurrently with vx_writer_push, it will block + * until vx_writer_push call finishes. + * + * 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 *writer); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/vortex-ffi/examples/CMakeLists.txt b/vortex-ffi/examples/CMakeLists.txt index 47228d9a903..d02e67f6a55 100644 --- a/vortex-ffi/examples/CMakeLists.txt +++ b/vortex-ffi/examples/CMakeLists.txt @@ -4,15 +4,7 @@ # allow linking with vortex_ffi_shared although it's not in current folder cmake_policy(SET CMP0079 NEW) -add_executable(scan scan.c) -target_link_libraries(scan PRIVATE vortex_ffi_shared) - -add_executable(scan_to_arrow scan_to_arrow.c) -target_link_libraries(scan_to_arrow PRIVATE - nanoarrow_shared vortex_ffi_shared) - -add_executable(dtype dtype.c) -target_link_libraries(dtype PRIVATE vortex_ffi_shared) - -add_executable(write_sample write_sample.c) -target_link_libraries(write_sample PRIVATE vortex_ffi_shared) +foreach(name scan scan_to_arrow dtype write_sample parallel_write) + add_executable(${name} ${name}.c) + target_link_libraries(${name} PRIVATE vortex_ffi_shared nanoarrow_shared) +endforeach() diff --git a/vortex-ffi/examples/parallel_write.c b/vortex-ffi/examples/parallel_write.c new file mode 100644 index 00000000000..2cabd470d4a --- /dev/null +++ b/vortex-ffi/examples/parallel_write.c @@ -0,0 +1,172 @@ +// SPDX-License-Identifier: CC-BY-4.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include "vortex.h" +#include +#include +#include +#include + +#define MAX_THREADS 64 +#define CHUNK_LEN 300000 + +const char *usage = "Write a .vortex file from multiple threads concurrently\n" + "Usage: parallel_write [-j threads] [-n chunks-per-thread] \n"; + +void print_error(const char *what, const vx_error *error) { + const vx_view str = vx_error_message(error); + fprintf(stderr, "%s: %.*s\n", what, (int)str.len, str.ptr); +} + +const vx_array *chunk_array(uint64_t start, uint64_t len) { + uint64_t *data = malloc(len * sizeof(uint64_t)); + for (uint64_t i = 0; i < len; ++i) { + data[i] = (start + i) % 997; + } + + vx_validity validity = {.type = VX_VALIDITY_NON_NULLABLE}; + vx_error *error = NULL; + const vx_array *array = vx_array_new_primitive(PTYPE_U64, data, len, &validity, &error); + free(data); + if (error != NULL) { + print_error("Error creating chunk array", error); + vx_error_free(error); + return NULL; + } + return array; +} + +struct worker { + pthread_t thread_id; + vx_writer *writer; + uint64_t chunks_per_thread; + vx_error *error; +}; + +void *write_thread(void *arg) { + struct worker *worker = arg; + + for (uint64_t chunk = 0; chunk < worker->chunks_per_thread; ++chunk) { + const uint64_t index = worker->thread_id * worker->chunks_per_thread + chunk; + const vx_array *array = chunk_array(index * CHUNK_LEN, CHUNK_LEN); + if (array == NULL) { + return NULL; + } + + vx_writer_push(worker->writer, array, &worker->error); + vx_array_free(array); + if (worker->error != NULL) { + return NULL; + } + } + + printf("Thread %lu finished, pushed %lu chunks\n", worker->thread_id + 1, worker->chunks_per_thread); + return NULL; +} + +vx_error *parallel_write(vx_writer *writer, uint64_t num_threads, uint64_t chunks_per_thread) { + pthread_t threads[MAX_THREADS]; + struct worker infos[MAX_THREADS] = {0}; + + printf("Writing using %lu threads, %lu chunks per thread\n", num_threads, chunks_per_thread); + for (uint64_t id = 0; id < num_threads; ++id) { + struct worker *info = &infos[id]; + info->thread_id = id; + info->writer = writer; + info->chunks_per_thread = chunks_per_thread; + pthread_create(&threads[id], NULL, write_thread, info); + } + + for (uint64_t id = 0; id < num_threads; ++id) { + pthread_join(threads[id], NULL); + } + + for (uint64_t id = 0; id < num_threads; ++id) { + if (infos[id].error != NULL) { + // Don't return other threads' errors, only the first one found + return infos[id].error; + } + } + + return NULL; +} + +int parse_options(int argc, char *argv[], uint64_t *threads, uint64_t *chunks_per_thread, char **output) { + int opt; + while ((opt = getopt(argc, argv, "j:n:")) != -1) { + switch (opt) { + case 'j': + *threads = strtoul(optarg, NULL, 10); + break; + case 'n': + *chunks_per_thread = strtoul(optarg, NULL, 10); + break; + default: + fprintf(stderr, "%s", usage); + return 1; + } + } + + if (*threads < 1 || *threads > MAX_THREADS) { + fprintf(stderr, "Invalid thread count %lu, expected [1; %d]\n", *threads, MAX_THREADS); + return 1; + } + + if (optind + 1 != argc) { + fprintf(stderr, "%s", usage); + return 1; + } + + *output = argv[optind]; + return 0; +} + +int main(int argc, char *argv[]) { + uint64_t threads = 4; + uint64_t chunks_per_thread = 3; + char *output; + if (parse_options(argc, argv, &threads, &chunks_per_thread, &output)) { + return 1; + } + + vx_session *session = vx_session_new(); + if (session == NULL) { + fprintf(stderr, "Failed to create Vortex session\n"); + return 1; + } + + const vx_dtype *dtype = vx_dtype_new_primitive(PTYPE_U64, false); + + vx_error *error = NULL; + vx_writer *writer = vx_writer_open(session, vx_view_from_cstr(output), dtype, threads, &error); + vx_dtype_free(dtype); + if (writer == NULL) { + print_error("Failed to open writer", error); + vx_error_free(error); + vx_session_free(session); + return 1; + } + + // vx_writer supports being pushed to concurrently from multiple threads. + error = parallel_write(writer, threads, chunks_per_thread); + if (error != NULL) { + print_error("Failed to write", error); + vx_error_free(error); + vx_writer_close(writer, &error); + vx_error_free(error); + vx_writer_free(writer); + vx_session_free(session); + return 1; + } + + vx_writer_close(writer, &error); + if (error != NULL) { + print_error("Error closing writer", error); + vx_error_free(error); + } + vx_writer_free(writer); + + printf("Wrote %lu rows to %s\n", threads * chunks_per_thread * CHUNK_LEN, output); + + vx_session_free(session); + return 0; +} diff --git a/vortex-ffi/examples/write_sample.c b/vortex-ffi/examples/write_sample.c index 456d3bba166..bd8336de335 100644 --- a/vortex-ffi/examples/write_sample.c +++ b/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, 32, &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; diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index ef493a27b03..338d70473cc 100644 --- a/vortex-ffi/src/lib.rs +++ b/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_close; + use crate::writer::vx_writer_free; + 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, 32, &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); } diff --git a/vortex-ffi/src/sink.rs b/vortex-ffi/src/writer.rs similarity index 54% rename from vortex-ffi/src/sink.rs rename to vortex-ffi/src/writer.rs index f307cfc8bc0..0ca4fab84a1 100644 --- a/vortex-ffi/src/sink.rs +++ b/vortex-ffi/src/writer.rs @@ -5,14 +5,13 @@ use futures::SinkExt; use futures::TryStreamExt; use futures::channel::mpsc; use futures::channel::mpsc::Sender; +use parking_lot::Mutex; use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; use vortex::error::VortexResult; -use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; -use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; use vortex::io::runtime::BlockingRuntime; @@ -20,7 +19,6 @@ use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; use crate::RUNTIME; -use crate::arc_wrapper; use crate::array::vx_array; use crate::dtype::vx_dtype; use crate::error::try_or_default; @@ -28,130 +26,144 @@ use crate::error::vx_error; use crate::session::vx_session; use crate::string::vx_view; -arc_wrapper!( - /// A handle to a Vortex file encapsulating the footer and logic for instantiating a reader. - VortexFile, - vx_file -); - #[expect(non_camel_case_types)] -/// 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. -pub struct vx_array_sink { - sink: Sender>, - writer: Task>, +/// vx_writer can be used to write vx_arrays into a (possibly remote) file. +pub struct vx_writer { + sender: Mutex>>>, + writer: Mutex>>>, dtype: DType, } -/// 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. +/// 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. +/// +/// "concurrent_array_limit" is the limit on the number of arrays that are +/// processed concurrently. This limits RAM used for processing. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_open_file( +pub unsafe extern "C-unwind" fn vx_writer_open( session: *const vx_session, path: vx_view, dtype: *const vx_dtype, + concurrent_array_limit: usize, error_out: *mut *mut vx_error, -) -> *mut vx_array_sink { +) -> *mut vx_writer { + let session = vx_session::as_ref(session).clone(); try_or_default(error_out, || { - let session = vx_session::as_ref(session).clone(); - - if path.ptr.is_null() { - vortex_bail!("null path"); - } + vortex_ensure!(!path.ptr.is_null()); let path = unsafe { path.as_str() }?.to_string(); let file_dtype = vx_dtype::as_ref(dtype); - // The channel size 32 was chosen arbitrarily. - let (sink, rx) = mpsc::channel(32); + let (sender, receiver) = mpsc::channel(concurrent_array_limit); let dtype = file_dtype.clone(); - let array_stream = ArrayStreamAdapter::new(dtype.clone(), rx.into_stream()); + let array_stream = ArrayStreamAdapter::new(dtype.clone(), receiver.into_stream()); let writer = session.handle().spawn(async move { let mut file = async_fs::File::create(path).await?; session.write_options().write(&mut file, array_stream).await }); - Ok(Box::into_raw(Box::new(vx_array_sink { - sink, - writer, + Ok(Box::into_raw(Box::new(vx_writer { + sender: Mutex::new(Some(sender)), + writer: Mutex::new(Some(writer)), dtype, }))) }) } -/// Push an array into a file sink. -/// Does not take ownership of array. +/// Push an array into a writer. Does not take ownership of array. /// -/// Errors if array's DType doesn't match sink's DType. +/// Array ordering across concurrent calls to this function is +/// non-deterministic: vx_writer_push(array1) called concurrently with +/// vx_writer_push(array2) may write array2 first. +/// +/// Errors if array's dtype and writer's initialized dtype are different. +/// Errors if writer has already been closed. +/// +/// Thread safe. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_push( - sink: *mut vx_array_sink, +pub unsafe extern "C-unwind" fn vx_writer_push( + writer: *mut vx_writer, array: *const vx_array, error_out: *mut *mut vx_error, ) { try_or_default(error_out, || { vortex_ensure!(!array.is_null()); - vortex_ensure!(!sink.is_null()); + vortex_ensure!(!writer.is_null()); let array = vx_array::as_ref(array); - let sink = unsafe { &mut *sink }; + let writer = unsafe { &*writer }; vortex_ensure!( - *array.dtype() == sink.dtype, - "array dtype {} does not match sink dtype {}", + *array.dtype() == writer.dtype, + "array dtype {} does not match writer dtype {}", array.dtype(), - sink.dtype + writer.dtype ); + + let mut sender = writer + .sender + .lock() + .clone() + .ok_or_else(|| vortex_err!("writer is closed"))?; + RUNTIME - .block_on(sink.sink.send(Ok(array.clone()))) - .map_err(|e| vortex_err!("Send error: {e}")) + .block_on(sender.send(Ok(array.clone()))) + .map_err(|e| vortex_err!("Writer already closed: {e}")) }) } -/// Closes an array sink, must be called to ensure all the values pushed to the sink are written -/// to the external resource. +/// 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. +/// +/// If this function is called concurrently with vx_writer_push, it will block +/// until vx_writer_push call finishes. +/// +/// Thread-unsafe. +/// +/// Errors if writer was already closed. +/// +/// Use vx_writer_free to free the writer afterwards. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_close( - sink: *mut vx_array_sink, +pub unsafe extern "C-unwind" fn vx_writer_close( + writer: *mut vx_writer, error_out: *mut *mut vx_error, ) { try_or_default(error_out, || { - let vx_array_sink { - sink, - writer, - dtype: _, - } = *unsafe { Box::from_raw(sink) }; - drop(sink); + vortex_ensure!(!writer.is_null()); + let writer = unsafe { &*writer }; + + drop(writer.sender.lock().take()); + + let writer = writer + .writer + .lock() + .take() + .ok_or_else(|| vortex_err!("writer is closed"))?; RUNTIME.block_on(async { let _footer = writer.await?; VortexResult::Ok(()) - })?; - - Ok(()) + }) }) } -/// Abort an array sink. File footer is not written, and file is left invalid. -/// Don't use sink after this call. +/// Release the writer. +/// +/// Thread unsafe. Must be called exactly once. +/// +/// If vx_writer_close wasn't called before this function, file is left +/// corrupted. #[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_abort(sink: *mut vx_array_sink) { - if sink.is_null() { +pub unsafe extern "C-unwind" fn vx_writer_free(writer: *mut vx_writer) { + if writer.is_null() { return; } - drop(unsafe { Box::from_raw(sink) }); + drop(unsafe { Box::from_raw(writer) }); } #[cfg(test)] @@ -178,7 +190,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] - fn test_sink_basic_workflow() { + fn test_writer_basic() { unsafe { let session = vx_session_new(); @@ -189,22 +201,20 @@ mod tests { let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error); + let writer = vx_writer_open(session, path, vx_dtype_ptr, &raw mut error); assert!(error.is_null()); - assert!(!sink.is_null()); + assert!(!writer.is_null()); - // Create and push an array let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); let vx_array_ptr = vx_array::new(Arc::new(array.into_array())); - vx_array_sink_push(sink, vx_array_ptr, &raw mut error); + vx_writer_push(writer, vx_array_ptr, &raw mut error); assert!(error.is_null()); - // Close the sink - vx_array_sink_close(sink, &raw mut error); + vx_writer_close(writer, &raw mut error); assert!(error.is_null()); - // Cleanup + vx_writer_free(writer); vx_array_free(vx_array_ptr); vx_dtype_free(vx_dtype_ptr); vx_session_free(session); @@ -213,7 +223,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] - fn test_sink_multiple_arrays() { + fn test_writer_multiple_arrays() { unsafe { let session = vx_session_new(); @@ -224,10 +234,9 @@ mod tests { let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error); + let writer = vx_writer_open(session, path, vx_dtype_ptr, &raw mut error); assert!(error.is_null()); - // Push multiple arrays for i in 0..3 { let start = i * 3; let array = PrimitiveArray::new( @@ -236,14 +245,15 @@ mod tests { ); let vx_array_ptr = vx_array::new(Arc::new(array.into_array())); - vx_array_sink_push(sink, vx_array_ptr, &raw mut error); + vx_writer_push(writer, vx_array_ptr, &raw mut error); assert!(error.is_null()); vx_array_free(vx_array_ptr); } - vx_array_sink_close(sink, &raw mut error); + vx_writer_close(writer, &raw mut error); assert!(error.is_null()); + vx_writer_free(writer); vx_dtype_free(vx_dtype_ptr); vx_session_free(session); @@ -252,7 +262,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] - fn test_sink_invalid_path() { + fn test_writer_invalid_path() { unsafe { let session = vx_session_new(); @@ -262,27 +272,23 @@ mod tests { let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(session, invalid_path, vx_dtype_ptr, &raw mut error); + let writer = vx_writer_open(session, invalid_path, vx_dtype_ptr, &raw mut error); - // The sink creation may succeed but close should fail due to invalid path - if !sink.is_null() { - // Push an array + // Creation may succeed but close should fail due to invalid path + if !writer.is_null() { let array = PrimitiveArray::new(buffer![1i32], Validity::NonNullable); let vx_array_ptr = vx_array::new(Arc::new(array.into_array())); - vx_array_sink_push(sink, vx_array_ptr, &raw mut error); + vx_writer_push(writer, vx_array_ptr, &raw mut error); vx_array_free(vx_array_ptr); - // Close should fail due to invalid path - vx_array_sink_close(sink, &raw mut error); + vx_writer_close(writer, &raw mut error); // Either error is set or operation succeeds (depends on filesystem) if !error.is_null() { vx_error_free(error); } - } else { - // Sink creation failed, which is also valid - if !error.is_null() { - vx_error_free(error); - } + vx_writer_free(writer); + } else if !error.is_null() { + vx_error_free(error); } vx_dtype_free(vx_dtype_ptr); @@ -292,7 +298,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] - fn test_sink_abort() { + fn test_writer_free() { let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); @@ -304,15 +310,15 @@ mod tests { let dtype = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(session, path, dtype, &raw mut error); + let writer = vx_writer_open(session, path, dtype, &raw mut error); assert!(error.is_null()); let array = vx_array::new(Arc::new(array.into_array())); - vx_array_sink_push(sink, array, &raw mut error); + vx_writer_push(writer, array, &raw mut error); assert!(error.is_null()); vx_array_free(array); - vx_array_sink_abort(sink); + vx_writer_free(writer); let opts = vx_data_source_options { paths: &raw const path, @@ -330,22 +336,63 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] - fn test_sink_null_path() { + fn test_writer_null_path() { + unsafe { + let session = vx_session_new(); + + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); + + let mut error = std::ptr::null_mut(); + let writer = vx_writer_open(session, vx_view::null(), vx_dtype_ptr, &raw mut error); + + assert!(writer.is_null()); + assert!(!error.is_null()); + + vx_error_free(error); + vx_dtype_free(vx_dtype_ptr); + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_writer_push_after_close() { unsafe { let session = vx_session_new(); + let temp_file = NamedTempFile::new().unwrap(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); let vx_dtype_ptr = vx_dtype::new(Arc::new(dtype)); let mut error = std::ptr::null_mut(); - // This should return null and set error due to null path - let sink = - vx_array_sink_open_file(session, vx_view::null(), vx_dtype_ptr, &raw mut error); + let writer = vx_writer_open(session, path, vx_dtype_ptr, &raw mut error); + assert!(error.is_null()); + + let array = PrimitiveArray::new(buffer![1i32], Validity::NonNullable); + let vx_array_ptr = vx_array::new(Arc::new(array.into_array())); + + vx_writer_push(writer, vx_array_ptr, &raw mut error); + assert!(error.is_null()); - assert!(sink.is_null()); + vx_writer_close(writer, &raw mut error); + assert!(error.is_null()); + + vx_writer_push(writer, vx_array_ptr, &raw mut error); assert!(!error.is_null()); + let message = crate::error::vx_error_message(error); + assert!(message.as_str().unwrap().contains("closed")); + vx_error_free(error); + vx_writer_close(writer, &raw mut error); + assert!(!error.is_null()); vx_error_free(error); + + vx_writer_free(writer); + + vx_array_free(vx_array_ptr); vx_dtype_free(vx_dtype_ptr); vx_session_free(session); } diff --git a/vortex-ffi/test/scan.cpp b/vortex-ffi/test/scan.cpp index e1f629dfde0..4a7b3a147bb 100644 --- a/vortex-ffi/test/scan.cpp +++ b/vortex-ffi/test/scan.cpp @@ -18,6 +18,7 @@ using FFI_ArrowSchema = ArrowSchema; #include #include "common.h" +#include "temp_path.hpp" namespace fs = std::filesystem; using namespace std::string_literals; @@ -28,31 +29,6 @@ using nanoarrow::UniqueArrayStream; using nanoarrow::UniqueArrayView; using nanoarrow::UniqueSchema; -struct TempPath : fs::path { - TempPath() = default; - explicit TempPath(fs::path p) : fs::path(std::move(p)) { - } - - TempPath(const TempPath &) = delete; - TempPath &operator=(const TempPath &) = delete; - - TempPath(TempPath &&other) noexcept : fs::path(std::move(other)) { - } - TempPath &operator=(TempPath &&other) noexcept { - if (this != &other) { - fs::remove(*this); - fs::path::operator=(std::move(other)); - } - return *this; - } - - ~TempPath() { - if (!empty()) { - fs::remove(*this); - } - } -}; - // StructArray { age=u8, height=u16? } [[nodiscard]] const vx_dtype *sample_dtype() { vx_struct_fields_builder *builder = vx_struct_fields_builder_new(); @@ -167,18 +143,21 @@ UniqueArrayStream sample_array_stream() { }; vx_error *error = nullptr; - vx_array_sink *sink = vx_array_sink_open_file(session, vx_view_from_cstr(path.c_str()), dtype, &error); - REQUIRE(sink != nullptr); + vx_writer *writer = vx_writer_open(session, vx_view_from_cstr(path.c_str()), dtype, 32, &error); + REQUIRE(writer != nullptr); require_no_error(error); + defer { + vx_writer_free(writer); + }; const vx_array *array = sample_array(); defer { vx_array_free(array); }; - vx_array_sink_push(sink, array, &error); + vx_writer_push(writer, array, &error); require_no_error(error); - vx_array_sink_close(sink, &error); + vx_writer_close(writer, &error); require_no_error(error); INFO("Written vortex file "s + path.generic_string()); diff --git a/vortex-ffi/test/temp_path.hpp b/vortex-ffi/test/temp_path.hpp new file mode 100644 index 00000000000..5d320e69967 --- /dev/null +++ b/vortex-ffi/test/temp_path.hpp @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include +#include + +namespace fs = std::filesystem; + +struct TempPath : fs::path { + TempPath() = default; + explicit TempPath(fs::path p) : fs::path(std::move(p)) { + } + + TempPath(const TempPath &) = delete; + TempPath &operator=(const TempPath &) = delete; + + TempPath(TempPath &&other) noexcept : fs::path(std::move(other)) { + } + TempPath &operator=(TempPath &&other) noexcept { + if (this != &other) { + fs::remove(*this); + fs::path::operator=(std::move(other)); + } + return *this; + } + + ~TempPath() { + if (!empty()) { + fs::remove(*this); + } + } +}; + +inline TempPath temp_vortex_path() { + return TempPath {fs::temp_directory_path() / + fs::path("sink-test-" + std::to_string(std::random_device {}()) + ".vortex")}; +} diff --git a/vortex-ffi/test/writer.cpp b/vortex-ffi/test/writer.cpp new file mode 100644 index 00000000000..db619c79b19 --- /dev/null +++ b/vortex-ffi/test/writer.cpp @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#include +#include +#include +#include +#include + +#include + +#include "common.h" +#include "temp_path.hpp" + +using Catch::Matchers::ContainsSubstring; + +namespace { + +const vx_array *pattern_array(uint64_t start, uint64_t len) { + std::vector data(len); + for (uint64_t i = 0; i < len; ++i) { + data[i] = (start + i) % 997; + } + + vx_validity validity = {}; + validity.type = VX_VALIDITY_NON_NULLABLE; + + vx_error *error = nullptr; + const vx_array *array = vx_array_new_primitive(PTYPE_U64, data.data(), len, &validity, &error); + require_no_error(error); + return array; +} + +TEST_CASE("Push after close", "[writer]") { + vx_session *session = vx_session_new(); + defer { + vx_session_free(session); + }; + + TempPath path = temp_vortex_path(); + const vx_dtype *dtype = vx_dtype_new_primitive(PTYPE_U64, false); + defer { + vx_dtype_free(dtype); + }; + + vx_error *error = nullptr; + vx_writer *sink = vx_writer_open(session, vx_view_from_cstr(path.c_str()), dtype, 32, &error); + require_no_error(error); + REQUIRE(sink != nullptr); + defer { + vx_writer_free(sink); + }; + + const vx_array *array = pattern_array(0, 1); + defer { + vx_array_free(array); + }; + + vx_writer_push(sink, array, &error); + require_no_error(error); + + vx_writer_close(sink, &error); + require_no_error(error); + + vx_writer_push(sink, array, &error); + REQUIRE(error != nullptr); + REQUIRE_THAT(to_string(error), ContainsSubstring("closed")); + vx_error_free(error); + + vx_writer_close(sink, &error); + REQUIRE(error != nullptr); + vx_error_free(error); +} + +TEST_CASE("Concurrent push", "[writer]") { + vx_session *session = vx_session_new(); + defer { + vx_session_free(session); + }; + + TempPath path = temp_vortex_path(); + const vx_dtype *dtype = vx_dtype_new_primitive(PTYPE_U64, false); + defer { + vx_dtype_free(dtype); + }; + + vx_error *error = nullptr; + vx_writer *sink = vx_writer_open(session, vx_view_from_cstr(path.c_str()), dtype, 32, &error); + require_no_error(error); + REQUIRE(sink != nullptr); + defer { + vx_writer_free(sink); + }; + + constexpr uint64_t NUM_THREADS = 8; + constexpr uint64_t CHUNK_LEN = 1000; + + std::vector pool; + pool.reserve(NUM_THREADS); + for (uint64_t i = 0; i < NUM_THREADS; ++i) { + pool.emplace_back([&, i = i] { + const vx_array *array = pattern_array(i * CHUNK_LEN, CHUNK_LEN); + vx_error *push_error = nullptr; + vx_writer_push(sink, array, &push_error); + vx_array_free(array); + if (push_error != nullptr) { + vx_error_free(push_error); + throw std::runtime_error("push failed"); + } + }); + } + for (auto &thread : pool) { + thread.join(); + } + + vx_writer_close(sink, &error); + require_no_error(error); + + vx_data_source_options opts = {}; + const vx_view ds_path = vx_view_from_cstr(path.c_str()); + opts.paths = &ds_path; + opts.paths_len = 1; + + const vx_data_source *ds = vx_data_source_new(session, &opts, &error); + require_no_error(error); + REQUIRE(ds != nullptr); + defer { + vx_data_source_free(ds); + }; + + vx_estimate row_count = {}; + vx_data_source_get_row_count(ds, &row_count); + REQUIRE(row_count.type == VX_ESTIMATE_EXACT); + REQUIRE(row_count.estimate == NUM_THREADS * CHUNK_LEN); +} + +} // namespace