diff --git a/Cargo.lock b/Cargo.lock index 16b0a83f9ac..f90b6ef27b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9785,6 +9785,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 8826423e995..6a275217cbc 100644 --- a/lang/cpp/include/vortex/writer.hpp +++ b/lang/cpp/include/vortex/writer.hpp @@ -10,20 +10,38 @@ #include #include +#include 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 handle_; + std::unique_ptr handle_; }; } // namespace vortex diff --git a/lang/cpp/src/writer.cpp b/lang/cpp/src/writer.cpp index 57035171a14..cbd11aa06f1 100644 --- a/lang/cpp/src/writer.cpp +++ b/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 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 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 diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index d83974921ec..c441499f792 100644 --- a/vortex-cuda/ffi/README.md +++ b/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 diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 3c83b9b9ffd..06c563438cd 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/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, diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 623e017d05d..e5fa0ef2192 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/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 /// 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 5476c6e494d..bba833e928a 100644 --- a/vortex-ffi/cinclude/vortex.h +++ b/vortex-ffi/cinclude/vortex.h @@ -365,23 +365,6 @@ typedef enum { */ typedef struct vx_array vx_array; -/** - * 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. * @@ -456,6 +439,11 @@ typedef struct vx_struct_fields vx_struct_fields; */ typedef struct vx_struct_fields_builder vx_struct_fields_builder; +/** + * vx_writer allows concurrently writing vx_array's into a .vortex file + */ +typedef struct vx_writer vx_writer; + /** * Array validity descriptor used by C FFI constructors. */ @@ -564,6 +552,20 @@ typedef struct { bool ordered; } vx_scan_options; +/** + * Summary of a written .vortex file + */ +typedef struct { + /** + * Number of rows + */ + uint64_t row_count; + /** + * File size in bytes + */ + uint64_t file_size; +} vx_write_summary; + #ifdef __cplusplus extern "C" { #endif // __cplusplus @@ -1471,34 +1473,6 @@ vx_session *vx_session_new(void); */ vx_session *vx_session_clone(const vx_session *session); -/** - * 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 */ @@ -1602,6 +1576,67 @@ 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 an upper limit on how many pushed vx_array's + * may be buffered before vx_writer_push blocks. + */ +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. + * + * You need to call vx_writer_free after this function even if it returned an + * error. + * + * "summary_out" may be NULL. If it's non-NULL, it is filled with file's row + * count and size. + * + * If this function is called concurrently with vx_writer_push, it will block + * until vx_writer_push finishes. + * + * Errors if writer was already closed. + * + * Thread unsafe. Must be called exactly once. + */ +void vx_writer_close(vx_writer *writer, vx_write_summary *summary_out, vx_error **error_out); + +/** + * Free 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..2bf2856d81e 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) + add_executable(${name} ${name}.c) + target_link_libraries(${name} PRIVATE vortex_ffi_shared nanoarrow_shared) +endforeach() diff --git a/vortex-ffi/examples/write_sample.c b/vortex-ffi/examples/write_sample.c index 456d3bba166..5e2b8bc0726 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, NULL, &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, NULL, &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, NULL, &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 1ef423bb383..b70bed67331 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -18,10 +18,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; @@ -39,14 +39,14 @@ pub use session::vx_session; pub use session::vx_session_free; pub use session::vx_session_new_with; pub use session::vx_session_ref; -pub use sink::vx_array_sink; -pub use sink::vx_array_sink_open_file_with_strategy; pub use string::vx_view; use vortex::dtype::FieldName; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::io::runtime::current::CurrentThreadWorkerPool; +pub use writer::vx_writer; +pub use writer::vx_writer_open_with_strategy; #[cfg(all(feature = "mimalloc", not(miri)))] #[global_allocator] @@ -197,12 +197,13 @@ 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::vx_runtime_set_worker_threads; use crate::vx_runtime_worker_count; + use crate::writer::vx_writer_close; + use crate::writer::vx_writer_free; + use crate::writer::vx_writer_open; + use crate::writer::vx_writer_push; #[test] #[cfg_attr(miri, ignore)] @@ -273,10 +274,11 @@ mod tests { unsafe { let vx_dtype_ptr = vx_dtype::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(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, ptr::null_mut(), &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/sink.rs deleted file mode 100644 index 12d5bfb1909..00000000000 --- a/vortex-ffi/src/sink.rs +++ /dev/null @@ -1,372 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use futures::SinkExt; -use futures::TryStreamExt; -use futures::channel::mpsc; -use futures::channel::mpsc::Sender; -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::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; -use vortex::file::WriteSummary; -use vortex::io::runtime::BlockingRuntime; -use vortex::io::runtime::Task; -use vortex::io::session::RuntimeSessionExt; -use vortex::layout::LayoutStrategy; - -use crate::RUNTIME; -use crate::array::vx_array; -use crate::dtype::vx_dtype; -use crate::error::try_or_default; -use crate::error::vx_error; -use crate::session::vx_session; -use crate::string::vx_view; - -/// 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>, - dtype: DType, -} - -/// Open a file sink with an explicit layout strategy. -/// -/// This is a Rust API for FFI crates layered on top of `vortex-ffi`; the base C API continues to -/// use the default write strategy. File creation and write errors are reported when the returned -/// sink is closed. -/// -/// # Safety -/// -/// `session`, `path`, and `dtype` must satisfy the same requirements as -/// `vx_array_sink_open_file`. -pub unsafe fn vx_array_sink_open_file_with_strategy( - session: *const vx_session, - path: vx_view, - dtype: *const vx_dtype, - strategy: Arc, -) -> VortexResult<*mut vx_array_sink> { - let session = vx_session::as_ref(session).clone(); - - if path.ptr.is_null() { - vortex_bail!("null path"); - } - 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 dtype = file_dtype.clone(); - let array_stream = ArrayStreamAdapter::new(dtype.clone(), rx.into_stream()); - - let writer_session = session.clone(); - let writer = session.handle().spawn(async move { - let mut file = async_fs::File::create(path).await?; - writer_session - .write_options() - .with_strategy(strategy) - .write(&mut file, array_stream) - .await - }); - - Ok(Box::into_raw(Box::new(vx_array_sink { - sink, - writer, - 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. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_open_file( - session: *const vx_session, - path: vx_view, - dtype: *const vx_dtype, - error_out: *mut *mut vx_error, -) -> *mut vx_array_sink { - try_or_default(error_out, || { - let strategy = WriteStrategyBuilder::default().build(); - unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } - }) -} - -/// Push an array into a file sink. -/// Does not take ownership of array. -/// -/// Errors if array's DType doesn't match sink's DType. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_push( - sink: *mut vx_array_sink, - 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()); - - let array = vx_array::as_ref(array); - let sink = unsafe { &mut *sink }; - - vortex_ensure!( - *array.dtype() == sink.dtype, - "array dtype {} does not match sink dtype {}", - array.dtype(), - sink.dtype - ); - RUNTIME - .block_on(sink.sink.send(Ok(array.clone()))) - .map_err(|e| vortex_err!("Send error: {e}")) - }) -} - -/// Closes an array sink, must be called to ensure all the values pushed to the sink are written -/// to the external resource. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_close( - sink: *mut vx_array_sink, - 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); - - 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. -#[unsafe(no_mangle)] -pub unsafe extern "C-unwind" fn vx_array_sink_abort(sink: *mut vx_array_sink) { - if sink.is_null() { - return; - } - drop(unsafe { Box::from_raw(sink) }); -} - -#[cfg(test)] -mod tests { - - use tempfile::NamedTempFile; - use vortex::array::IntoArray; - use vortex::array::arrays::PrimitiveArray; - use vortex::array::validity::Validity; - use vortex::buffer::buffer; - use vortex::dtype::DType; - - use super::*; - use crate::array::vx_array; - use crate::array::vx_array_free; - use crate::data_source::vx_data_source_new; - use crate::data_source::vx_data_source_options; - use crate::dtype::vx_dtype; - use crate::dtype::vx_dtype_free; - use crate::error::vx_error_free; - use crate::session::vx_session_free; - use crate::session::vx_session_new; - - #[test] - #[cfg_attr(miri, ignore)] - fn test_sink_basic_workflow() { - 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(dtype); - - let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(session, path, vx_dtype_ptr, &raw mut error); - assert!(error.is_null()); - assert!(!sink.is_null()); - - // Create and push an array - let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); - let vx_array_ptr = vx_array::new(array.into_array()); - - vx_array_sink_push(sink, vx_array_ptr, &raw mut error); - assert!(error.is_null()); - - // Close the sink - vx_array_sink_close(sink, &raw mut error); - assert!(error.is_null()); - - // Cleanup - vx_array_free(vx_array_ptr); - vx_dtype_free(vx_dtype_ptr); - vx_session_free(session); - } - } - - #[test] - #[cfg_attr(miri, ignore)] - fn test_sink_multiple_arrays() { - 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::U64, false.into()); - let vx_dtype_ptr = vx_dtype::new(dtype); - - let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(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( - buffer![start as u64, (start + 1) as u64, (start + 2) as u64], - Validity::NonNullable, - ); - let vx_array_ptr = vx_array::new(array.into_array()); - - vx_array_sink_push(sink, vx_array_ptr, &raw mut error); - assert!(error.is_null()); - - vx_array_free(vx_array_ptr); - } - - vx_array_sink_close(sink, &raw mut error); - assert!(error.is_null()); - - vx_dtype_free(vx_dtype_ptr); - vx_session_free(session); - } - } - - #[test] - #[cfg_attr(miri, ignore)] - fn test_sink_invalid_path() { - unsafe { - let session = vx_session_new(); - - // Use a path that will fail during file creation (read-only directory on most systems) - let invalid_path = vx_view::from_str("/dev/null/invalid.vortex"); - let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); - let vx_dtype_ptr = vx_dtype::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); - - // The sink creation may succeed but close should fail due to invalid path - if !sink.is_null() { - // Push an array - let array = PrimitiveArray::new(buffer![1i32], Validity::NonNullable); - let vx_array_ptr = vx_array::new(array.into_array()); - vx_array_sink_push(sink, 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); - // 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_dtype_free(vx_dtype_ptr); - vx_session_free(session); - } - } - - #[test] - #[cfg_attr(miri, ignore)] - fn test_sink_abort() { - let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); - let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); - - let file = NamedTempFile::new().unwrap(); - let path = vx_view::from_str(file.path().to_str().unwrap()); - - unsafe { - let session = vx_session_new(); - let dtype = vx_dtype::new(dtype); - - let mut error = std::ptr::null_mut(); - let sink = vx_array_sink_open_file(session, path, dtype, &raw mut error); - assert!(error.is_null()); - - let array = vx_array::new(array.into_array()); - vx_array_sink_push(sink, array, &raw mut error); - assert!(error.is_null()); - vx_array_free(array); - - vx_array_sink_abort(sink); - - let opts = vx_data_source_options { - paths: &raw const path, - paths_len: 1, - }; - let ds = vx_data_source_new(session, &raw const opts, &raw mut error); - assert!(ds.is_null()); - assert!(!error.is_null()); - vx_error_free(error); - - vx_dtype_free(dtype); - vx_session_free(session); - } - } - - #[test] - #[cfg_attr(miri, ignore)] - fn test_sink_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(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); - - assert!(sink.is_null()); - assert!(!error.is_null()); - - vx_error_free(error); - vx_dtype_free(vx_dtype_ptr); - vx_session_free(session); - } - } -} diff --git a/vortex-ffi/src/writer.rs b/vortex-ffi/src/writer.rs new file mode 100644 index 00000000000..36f06b5a561 --- /dev/null +++ b/vortex-ffi/src/writer.rs @@ -0,0 +1,451 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use futures::SinkExt; +use futures::TryStreamExt; +use futures::channel::mpsc; +use futures::channel::mpsc::Sender; +use parking_lot::RwLock; +use vortex::array::ArrayRef; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::dtype::DType; +use vortex::error::VortexError; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::WriteStrategyBuilder; +use vortex::file::WriteSummary; +use vortex::io::runtime::BlockingRuntime; +use vortex::io::runtime::Task; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::LayoutStrategy; +use vortex::session::VortexSession; + +use crate::RUNTIME; +use crate::array::vx_array; +use crate::dtype::vx_dtype; +use crate::error::try_or_default; +use crate::error::vx_error; +use crate::session::vx_session; +use crate::string::vx_view; + +struct Inner { + sender: Option>>, + task: Option>>, +} + +/// vx_writer allows concurrently writing vx_array's into a .vortex file +pub struct vx_writer { + session: VortexSession, + inner: RwLock, + dtype: DType, +} + +impl vx_writer { + fn take_error(&self) -> VortexError { + let task = self.inner.write().task.take(); + match task { + Some(task) => match RUNTIME.block_on(task) { + Ok(_) => vortex_err!("writer is closed"), + Err(e) => e, + }, + None => vortex_err!("writer is closed"), + } + } +} + +/// Summary of a written .vortex file +#[repr(C)] +#[cfg_attr(test, derive(Debug, Clone, Copy))] +pub struct vx_write_summary { + /// Number of rows + pub row_count: u64, + /// File size in bytes + pub file_size: u64, +} + +/// Open a writer for a file at "path" with explicit write strategy. "path" +/// is copied. +/// +/// "dtype" is used to validate pushed arrays so they would all have the same +/// schema. +/// +/// "concurrent_array_limit" bounds how many pushed arrays may be buffered in +/// flight before `vx_writer_push` blocks (channel capacity / backpressure). It +/// caps RAM used for buffering; the actual encoding parallelism is governed by +/// the write strategy. +/// +/// # Safety +/// +/// session and dtype must be non-null pointers to valid objects. +/// path's pointer must be NULL only on len = 0. +pub unsafe fn vx_writer_open_with_strategy( + session: *const vx_session, + path: vx_view, + dtype: *const vx_dtype, + concurrent_array_limit: usize, + strategy: Arc, +) -> VortexResult<*mut vx_writer> { + let session = vx_session::as_ref(session).clone(); + vortex_ensure!(!path.ptr.is_null()); + let path = unsafe { path.as_str() }?.to_string(); + + let file_dtype = vx_dtype::as_ref(dtype); + let (sender, receiver) = mpsc::channel(concurrent_array_limit); + let dtype = file_dtype.clone(); + let array_stream = ArrayStreamAdapter::new(dtype.clone(), receiver.into_stream()); + + let writer = Box::new(vx_writer { + session, + inner: RwLock::new(Inner { + sender: Some(sender), + task: None, + }), + dtype, + }); + + // Spawn the write task on the writer's own session so the task and the + // driving block_on (in push/close) share one executor. + let task_session = writer.session.clone(); + let task = writer.session.handle().spawn(async move { + let mut file = async_fs::File::create(path).await?; + task_session + .write_options() + .with_strategy(strategy) + .write(&mut file, array_stream) + .await + }); + writer.inner.write().task = Some(task); + + Ok(Box::into_raw(writer)) +} + +/// 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 an upper limit on how many pushed vx_array's +/// may be buffered before vx_writer_push blocks. +#[unsafe(no_mangle)] +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_writer { + let strategy = WriteStrategyBuilder::default().build(); + try_or_default(error_out, || unsafe { + vx_writer_open_with_strategy(session, path, dtype, concurrent_array_limit, strategy) + }) +} + +/// 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. +#[unsafe(no_mangle)] +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!(!writer.is_null()); + + let array = vx_array::as_ref(array); + let writer = unsafe { &*writer }; + + vortex_ensure!( + *array.dtype() == writer.dtype, + "array dtype {} does not match writer dtype {}", + array.dtype(), + writer.dtype + ); + + let send_result = { + let inner = writer.inner.read(); + let mut sender = inner + .sender + .clone() + .ok_or_else(|| vortex_err!("writer is closed"))?; + RUNTIME.block_on(sender.send(Ok(array.clone()))) + }; + + match send_result { + Ok(_) => Ok(()), + Err(_) => Err(writer.take_error()), + } + }) +} + +/// 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. +/// +/// You need to call vx_writer_free after this function even if it returned an +/// error. +/// +/// "summary_out" may be NULL. If it's non-NULL, it is filled with file's row +/// count and size. +/// +/// If this function is called concurrently with vx_writer_push, it will block +/// until vx_writer_push finishes. +/// +/// Errors if writer was already closed. +/// +/// Thread unsafe. Must be called exactly once. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_writer_close( + writer: *mut vx_writer, + summary_out: *mut vx_write_summary, + error_out: *mut *mut vx_error, +) { + try_or_default(error_out, || { + vortex_ensure!(!writer.is_null()); + let writer = unsafe { &*writer }; + + let task = { + let mut inner = writer.inner.write(); + drop(inner.sender.take()); + inner.task.take() + }; + let task = task.ok_or_else(|| vortex_err!("writer is closed"))?; + + let summary = RUNTIME.block_on(task)?; + if !summary_out.is_null() { + unsafe { + *summary_out = vx_write_summary { + row_count: summary.row_count(), + file_size: summary.size(), + }; + } + } + VortexResult::Ok(()) + }) +} + +/// Free 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_writer_free(writer: *mut vx_writer) { + if writer.is_null() { + return; + } + drop(unsafe { Box::from_raw(writer) }); +} + +#[cfg(test)] +mod tests { + use std::ptr; + use std::thread::spawn; + + use tempfile::NamedTempFile; + use vortex::array::IntoArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::validity::Validity; + use vortex::buffer::buffer; + use vortex::dtype::DType; + + use super::*; + use crate::array::vx_array; + use crate::array::vx_array_free; + use crate::dtype::vx_dtype; + use crate::dtype::vx_dtype_free; + use crate::error::vx_error_free; + use crate::session::vx_session_free; + use crate::session::vx_session_new; + + #[test] + #[cfg_attr(miri, ignore)] + fn test_writer() { + let temp_file = NamedTempFile::new().unwrap(); + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); + + unsafe { + let session = vx_session_new(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); + let dtype = vx_dtype::new(dtype); + + let mut error = ptr::null_mut(); + let writer = vx_writer_open(session, path, dtype, 1, &raw mut error); + assert!(error.is_null()); + assert!(!writer.is_null()); + + let array = vx_array::new(array.into_array()); + vx_writer_push(writer, array, &raw mut error); + assert!(error.is_null()); + + let mut summary = vx_write_summary { + row_count: 0, + file_size: 0, + }; + vx_writer_close(writer, &raw mut summary, &raw mut error); + assert!(error.is_null()); + assert_eq!(summary.row_count, 3); + assert!(summary.file_size > 0); + + vx_writer_free(writer); + vx_array_free(array); + vx_dtype_free(dtype); + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_writer_holds_session() { + let temp_file = NamedTempFile::new().unwrap(); + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable); + + unsafe { + let session = vx_session_new(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); + let dtype = vx_dtype::new(dtype); + + let mut error = ptr::null_mut(); + let writer = vx_writer_open(session, path, dtype, 1, &raw mut error); + assert!(error.is_null()); + + vx_session_free(session); + + let array = vx_array::new(array.into_array()); + vx_writer_push(writer, array, &raw mut error); + assert!(error.is_null()); + + vx_writer_close(writer, ptr::null_mut(), &raw mut error); + assert!(error.is_null()); + + vx_writer_free(writer); + vx_array_free(array); + vx_dtype_free(dtype); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_writer_concurrent() { + let temp_file = NamedTempFile::new().unwrap(); + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + unsafe { + let session = vx_session_new(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); + let dtype = vx_dtype::new(dtype); + + let mut error = ptr::null_mut(); + let writer = vx_writer_open(session, path, dtype, 4, &raw mut error); + assert!(error.is_null()); + + let addr = writer as usize; + let pool: Vec<_> = (0..4) + .map(|t| { + spawn(move || { + let writer = addr as *mut vx_writer; + let mut err = ptr::null_mut(); + for i in 0..4 { + let v = t * 4 + i; + let array = PrimitiveArray::new(buffer![v], Validity::NonNullable); + let array = vx_array::new(array.into_array()); + vx_writer_push(writer, array, &raw mut err); + assert!(err.is_null()); + vx_array_free(array); + } + }) + }) + .collect(); + for handle in pool { + handle.join().unwrap(); + } + + let mut summary = vx_write_summary { + row_count: 0, + file_size: 0, + }; + vx_writer_close(writer, &raw mut summary, &raw mut error); + assert!(error.is_null()); + assert_eq!(summary.row_count, u64::try_from(16).unwrap()); + + vx_writer_free(writer); + vx_dtype_free(dtype); + vx_session_free(session); + } + } + + #[test] + #[cfg_attr(miri, ignore)] + fn test_writer_null_path() { + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + unsafe { + let session = vx_session_new(); + let vx_dtype_ptr = vx_dtype::new(dtype); + + let mut error = ptr::null_mut(); + let writer = vx_writer_open(session, vx_view::null(), vx_dtype_ptr, 1, &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() { + let temp_file = NamedTempFile::new().unwrap(); + let dtype = DType::Primitive(vortex::dtype::PType::I32, false.into()); + let array = PrimitiveArray::new(buffer![1i32], Validity::NonNullable); + unsafe { + let session = vx_session_new(); + let path = vx_view::from_str(temp_file.path().to_str().unwrap()); + let dtype = vx_dtype::new(dtype); + + let mut error = ptr::null_mut(); + let writer = vx_writer_open(session, path, dtype, 1, &raw mut error); + assert!(error.is_null()); + + let array = vx_array::new(array.into_array()); + vx_writer_push(writer, array, &raw mut error); + assert!(error.is_null()); + + vx_writer_close(writer, ptr::null_mut(), &raw mut error); + assert!(error.is_null()); + + vx_writer_push(writer, array, &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, ptr::null_mut(), &raw mut error); + assert!(!error.is_null()); + vx_error_free(error); + + vx_writer_free(writer); + + vx_array_free(array); + vx_dtype_free(dtype); + vx_session_free(session); + } + } +} diff --git a/vortex-ffi/test/scan.cpp b/vortex-ffi/test/scan.cpp index e1f629dfde0..ef0f4f620c0 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, nullptr, &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..3528903b2a6 --- /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_path() { + std::string name = "test-" + std::to_string(std::random_device {}()) + ".vortex"; + return TempPath {fs::temp_directory_path() / fs::path(name)}; +} diff --git a/vortex-ffi/test/writer.cpp b/vortex-ffi/test/writer.cpp new file mode 100644 index 00000000000..1502a4ae8a8 --- /dev/null +++ b/vortex-ffi/test/writer.cpp @@ -0,0 +1,137 @@ +// 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 *make_array(size_t start, size_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_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 = make_array(0, 1); + defer { + vx_array_free(array); + }; + + vx_writer_push(sink, array, &error); + require_no_error(error); + + vx_writer_close(sink, nullptr, &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, nullptr, &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_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 size_t threads = 8; + constexpr size_t len = 1000; + + std::vector pool; + pool.reserve(threads); + for (uint64_t i = 0; i < threads; ++i) { + pool.emplace_back([&, i = i] { + const vx_array *array = make_array(i * len, len); + vx_error *worker_error = nullptr; + vx_writer_push(sink, array, &worker_error); + vx_array_free(array); + if (worker_error != nullptr) { + vx_error_free(worker_error); + throw std::runtime_error("push failed"); + } + }); + } + for (auto &thread : pool) { + thread.join(); + } + + vx_write_summary summary = {}; + vx_writer_close(sink, &summary, &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 == threads * len); + REQUIRE(row_count.estimate == summary.row_count); +} + +} // namespace