From d1020ce3a4856c4a65f16c13621f08c559f95e6f Mon Sep 17 00:00:00 2001 From: madhu-mitha-e Date: Mon, 1 Jun 2026 11:13:53 +0530 Subject: [PATCH 1/3] feat(mohu-io): implement Arrow IPC file and stream read/write Closes #35 Signed-off-by: madhu-mitha-e --- crates/mohu-io/Cargo.toml | 36 ++-- crates/mohu-io/src/arrow.rs | 362 ++++++++++++++++++++++++++++++++++++ 2 files changed, 380 insertions(+), 18 deletions(-) diff --git a/crates/mohu-io/Cargo.toml b/crates/mohu-io/Cargo.toml index fda7c36..b7c01ab 100644 --- a/crates/mohu-io/Cargo.toml +++ b/crates/mohu-io/Cargo.toml @@ -1,18 +1,18 @@ -[package] -name = "mohu-io" -description = "Array serialization and I/O: .npy/.npz, CSV, Apache Arrow IPC, and memory-mapped files" -version.workspace = true -edition.workspace = true -license.workspace = true -authors.workspace = true -repository.workspace = true -homepage.workspace = true -keywords.workspace = true -categories.workspace = true - -[dependencies] -mohu-core.workspace = true -arrow.workspace = true -serde.workspace = true -memmap2.workspace = true -thiserror.workspace = true +[package] +name = "mohu-io" +description = "Array serialization and I/O: .npy/.npz, CSV, Apache Arrow IPC, and memory-mapped files" +version.workspace = true +edition.workspace = true +license.workspace = true +authors.workspace = true +repository.workspace = true +homepage.workspace = true +keywords.workspace = true +categories.workspace = true + +[dependencies] +mohu-core.workspace = true +arrow = { workspace = true, features = ["ipc"] } +serde.workspace = true +memmap2.workspace = true +thiserror.workspace = true diff --git a/crates/mohu-io/src/arrow.rs b/crates/mohu-io/src/arrow.rs index e69de29..1678de4 100644 --- a/crates/mohu-io/src/arrow.rs +++ b/crates/mohu-io/src/arrow.rs @@ -0,0 +1,362 @@ +//! Arrow IPC serialization and deserialization for mohu buffers. +//! +//! Exposes four functions: +//! +//! | Function | Description | +//! |---------------------|--------------------------------------------------| +//! | [`read_ipc_file`] | Read an Arrow IPC file into a `Vec` | +//! | [`read_ipc_stream`] | Read an Arrow IPC stream into a `Vec` | +//! | [`write_ipc_file`] | Write `Buffer` slices to an Arrow IPC file | +//! | [`write_ipc_stream`]| Write `Buffer` slices to an Arrow IPC stream | +//! +//! # Arrow IPC formats +//! +//! - **File format**: random-access, seekable, has a footer. Suited for +//! on-disk storage. +//! - **Streaming format**: sequential, no footer. Required for inter-process +//! streaming (Polars pipe, DuckDB query results, Arrow Flight protocol). + +use std::io::{Read, Seek, Write}; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, UInt16Array, UInt32Array, UInt64Array, UInt8Array}; +use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; +use arrow::ipc::reader::{FileReader, StreamReader}; +use arrow::ipc::writer::{FileWriter, StreamWriter}; +use arrow::record_batch::RecordBatch; + +use mohu_core::mohu_buffer::Buffer; +use mohu_core::mohu_dtype::DType; +use mohu_core::mohu_error::{MohuError, MohuResult}; + +// ── Helpers ─────────────────────────────────────────────────────────────────── + +/// Maps a mohu `DType` to an Arrow `DataType`. +fn dtype_to_arrow(dtype: DType) -> MohuResult { + match dtype { + DType::I8 => Ok(ArrowDataType::Int8), + DType::I16 => Ok(ArrowDataType::Int16), + DType::I32 => Ok(ArrowDataType::Int32), + DType::I64 => Ok(ArrowDataType::Int64), + DType::U8 => Ok(ArrowDataType::UInt8), + DType::U16 => Ok(ArrowDataType::UInt16), + DType::U32 => Ok(ArrowDataType::UInt32), + DType::U64 => Ok(ArrowDataType::UInt64), + DType::F32 => Ok(ArrowDataType::Float32), + DType::F64 => Ok(ArrowDataType::Float64), + other => Err(MohuError::UnsupportedDType { + op: "arrow IPC", + dtype: format!("{other:?}"), + }), + } +} + +/// Maps an Arrow `DataType` to a mohu `DType`. +/// Maps an Arrow `DataType` to a mohu `DType`. + + +/// Converts a 1-D `Buffer` to an Arrow `ArrayRef`. +fn buffer_to_array(buf: &Buffer) -> MohuResult { + match buf.dtype() { + DType::I8 => { + let s = buf.as_slice::()?; + Ok(Arc::new(Int8Array::from(s.to_vec()))) + } + DType::I16 => { + let s = buf.as_slice::()?; + Ok(Arc::new(Int16Array::from(s.to_vec()))) + } + DType::I32 => { + let s = buf.as_slice::()?; + Ok(Arc::new(Int32Array::from(s.to_vec()))) + } + DType::I64 => { + let s = buf.as_slice::()?; + Ok(Arc::new(Int64Array::from(s.to_vec()))) + } + DType::U8 => { + let s = buf.as_slice::()?; + Ok(Arc::new(UInt8Array::from(s.to_vec()))) + } + DType::U16 => { + let s = buf.as_slice::()?; + Ok(Arc::new(UInt16Array::from(s.to_vec()))) + } + DType::U32 => { + let s = buf.as_slice::()?; + Ok(Arc::new(UInt32Array::from(s.to_vec()))) + } + DType::U64 => { + let s = buf.as_slice::()?; + Ok(Arc::new(UInt64Array::from(s.to_vec()))) + } + DType::F32 => { + let s = buf.as_slice::()?; + Ok(Arc::new(Float32Array::from(s.to_vec()))) + } + DType::F64 => { + let s = buf.as_slice::()?; + Ok(Arc::new(Float64Array::from(s.to_vec()))) + } + other => Err(MohuError::UnsupportedDType { + op: "arrow IPC", + dtype: format!("{other:?}"), + }), + } +} + +/// Converts an Arrow `ArrayRef` to a 1-D `Buffer`. +fn array_to_buffer(array: &ArrayRef) -> MohuResult { + match array.data_type() { + ArrowDataType::Int8 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast Int8Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::Int16 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast Int16Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::Int32 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast Int32Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::Int64 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast Int64Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::UInt8 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast UInt8Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::UInt16 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast UInt16Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::UInt32 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast UInt32Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::UInt64 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast UInt64Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::Float32 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast Float32Array failed"))?; + Buffer::from_slice(a.values()) + } + ArrowDataType::Float64 => { + let a = array.as_any().downcast_ref::() + .ok_or_else(|| MohuError::bug("downcast Float64Array failed"))?; + Buffer::from_slice(a.values()) + } + other => Err(MohuError::UnsupportedDType { + op: "arrow IPC", + dtype: format!("{other:?}"), + }), + } +} + +/// Builds a `RecordBatch` from a slice of `Buffer`s. +fn buffers_to_record_batch(buffers: &[Buffer]) -> MohuResult { + let mut fields = Vec::with_capacity(buffers.len()); + let mut arrays = Vec::with_capacity(buffers.len()); + + for (i, buf) in buffers.iter().enumerate() { + let arrow_dtype = dtype_to_arrow(buf.dtype())?; + let field = Field::new(format!("col_{i}"), arrow_dtype, false); + fields.push(field); + arrays.push(buffer_to_array(buf)?); + } + + let schema = Arc::new(Schema::new(fields)); + RecordBatch::try_new(schema, arrays) + .map_err(|e| MohuError::ArrowSchema(e.to_string())) +} + +/// Extracts `Buffer`s from a `RecordBatch`. +fn record_batch_to_buffers(batch: &RecordBatch) -> MohuResult> { + batch + .columns() + .iter() + .map(array_to_buffer) + .collect() +} + +// ── Public API ──────────────────────────────────────────────────────────────── + +/// Reads an Arrow IPC **file** from `reader` and returns the columns of +/// every record batch as a flat `Vec`. +/// +/// The Arrow IPC file format is random-access and seekable, making it +/// suitable for on-disk storage. +/// +/// # Errors +/// +/// Returns an error if the IPC data is malformed or contains unsupported +/// Arrow data types. +pub fn read_ipc_file(reader: R) -> MohuResult> { + let file_reader = FileReader::try_new(reader, None) + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + let mut buffers = Vec::new(); + for batch in file_reader { + let batch = batch.map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + buffers.extend(record_batch_to_buffers(&batch)?); + } + Ok(buffers) +} + +/// Reads an Arrow IPC **stream** from `reader` and returns the columns of +/// every record batch as a flat `Vec`. +/// +/// The Arrow IPC streaming format is sequential with no footer, making it +/// suitable for inter-process streaming (Polars pipe, DuckDB query results, +/// Arrow Flight protocol). +/// +/// # Errors +/// +/// Returns an error if the IPC data is malformed or contains unsupported +/// Arrow data types. +pub fn read_ipc_stream(reader: R) -> MohuResult> { + let stream_reader = StreamReader::try_new(reader, None) + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + let mut buffers = Vec::new(); + for batch in stream_reader { + let batch = batch.map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + buffers.extend(record_batch_to_buffers(&batch)?); + } + Ok(buffers) +} + +/// Writes `buffers` to `writer` in Arrow IPC **file** format. +/// +/// All buffers are written as a single record batch where each buffer +/// becomes one column (`col_0`, `col_1`, …). +/// +/// # Errors +/// +/// Returns an error if any buffer has an unsupported dtype or if writing +/// fails. +pub fn write_ipc_file(writer: W, buffers: &[Buffer]) -> MohuResult<()> { + if buffers.is_empty() { + return Ok(()); + } + + let batch = buffers_to_record_batch(buffers)?; + let schema = batch.schema(); + + let mut file_writer = FileWriter::try_new(writer, &schema) + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + file_writer + .write(&batch) + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + file_writer + .finish() + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + Ok(()) +} + +/// Writes `buffers` to `writer` in Arrow IPC **stream** format. +/// +/// All buffers are written as a single record batch where each buffer +/// becomes one column (`col_0`, `col_1`, …). +/// +/// # Errors +/// +/// Returns an error if any buffer has an unsupported dtype or if writing +/// fails. +pub fn write_ipc_stream(writer: W, buffers: &[Buffer]) -> MohuResult<()> { + if buffers.is_empty() { + return Ok(()); + } + + let batch = buffers_to_record_batch(buffers)?; + let schema = batch.schema(); + + let mut stream_writer = StreamWriter::try_new(writer, &schema) + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + stream_writer + .write(&batch) + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + stream_writer + .finish() + .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; + + Ok(()) +} + +// ── Tests ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Cursor; + + #[test] + fn round_trip_ipc_file_f32() { + let buf = Buffer::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]).unwrap(); + let mut bytes = Vec::new(); + write_ipc_file(&mut bytes, &[buf.clone()]).unwrap(); + + let result = read_ipc_file(Cursor::new(bytes)).unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0].as_slice::().unwrap(), &[1.0_f32, 2.0, 3.0, 4.0]); + } + + #[test] + fn round_trip_ipc_stream_f64() { + let buf = Buffer::from_slice(&[10.0_f64, 20.0, 30.0]).unwrap(); + let mut bytes = Vec::new(); + write_ipc_stream(&mut bytes, &[buf.clone()]).unwrap(); + + let result = read_ipc_stream(Cursor::new(bytes)).unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0].as_slice::().unwrap(), &[10.0_f64, 20.0, 30.0]); + } + + #[test] + fn round_trip_ipc_file_i32() { + let buf = Buffer::from_slice(&[1_i32, 2, 3, 4, 5]).unwrap(); + let mut bytes = Vec::new(); + write_ipc_file(&mut bytes, &[buf.clone()]).unwrap(); + + let result = read_ipc_file(Cursor::new(bytes)).unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0].as_slice::().unwrap(), &[1_i32, 2, 3, 4, 5]); + } + + #[test] + fn round_trip_ipc_stream_multiple_buffers() { + let buf1 = Buffer::from_slice(&[1_i32, 2, 3]).unwrap(); + let buf2 = Buffer::from_slice(&[4.0_f64, 5.0, 6.0]).unwrap(); + let mut bytes = Vec::new(); + write_ipc_stream(&mut bytes, &[buf1, buf2]).unwrap(); + + let result = read_ipc_stream(Cursor::new(bytes)).unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result[0].as_slice::().unwrap(), &[1_i32, 2, 3]); + assert_eq!(result[1].as_slice::().unwrap(), &[4.0_f64, 5.0, 6.0]); + } + + #[test] + fn write_empty_buffers_is_ok() { + let mut bytes = Vec::new(); + assert!(write_ipc_file(&mut bytes, &[]).is_ok()); + assert!(write_ipc_stream(&mut bytes, &[]).is_ok()); + } +} From c8ffc84df6bb00888e5a8b3a966c54812af1349c Mon Sep 17 00:00:00 2001 From: madhu-mitha-e Date: Mon, 1 Jun 2026 11:19:23 +0530 Subject: [PATCH 2/3] temp: stage windows path deletions --- ...freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" | 7 ------- .../boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" | 6 ------ ...h\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" | 7 ------- ...ngenerator\nmultivariate\npermutation\nseeding.rs" | 7 ------- ...st\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" | 10 ---------- ...\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" | 11 ----------- ...ssel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" | 8 -------- .../approx\nassert\ndtype\nfixtures\ngen\nperf.rs" | 6 ------ ...mpl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" | 8 -------- 9 files changed, 70 deletions(-) delete mode 100644 "crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" delete mode 100644 "crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" delete mode 100644 "crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" delete mode 100644 "crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" delete mode 100644 "crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" delete mode 100644 "crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" delete mode 100644 "crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" delete mode 100644 "crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" delete mode 100644 "crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" diff --git "a/crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" "b/crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" deleted file mode 100644 index 2912da6..0000000 --- "a/crates/mohu-fft/src/freq\nhelpers\nnd\nnorm\nplan\nreal\ntransform.rs" +++ /dev/null @@ -1,7 +0,0 @@ -// freq -helpers -nd -norm -plan -real -transform — implementation pending diff --git "a/crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" "b/crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" deleted file mode 100644 index 903bd67..0000000 --- "a/crates/mohu-index/src/boolean\nfancy\ngather\nslice\ntake\nwhere_op.rs" +++ /dev/null @@ -1,6 +0,0 @@ -// boolean -fancy -gather -slice -take -where_op — implementation pending diff --git "a/crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" "b/crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" deleted file mode 100644 index 81a59d0..0000000 --- "a/crates/mohu-masked/src/arith\narray\ncompress\nfill\nio\nmask_ops\nreduce.rs" +++ /dev/null @@ -1,7 +0,0 @@ -// arith -array -compress -fill -io -mask_ops -reduce — implementation pending diff --git "a/crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" "b/crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" deleted file mode 100644 index 391727c..0000000 --- "a/crates/mohu-random/src/continuous\ndiscrete\nentropy\ngenerator\nmultivariate\npermutation\nseeding.rs" +++ /dev/null @@ -1,7 +0,0 @@ -// continuous -discrete -entropy -generator -multivariate -permutation -seeding — implementation pending diff --git "a/crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" "b/crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" deleted file mode 100644 index 1e0359c..0000000 --- "a/crates/mohu-simd/src/arith\nbitwise\ncast\ncmp\ncopy\nfill\nfma\nmath\nreduce\ndetect.rs" +++ /dev/null @@ -1,10 +0,0 @@ -// arith -bitwise -cast -cmp -copy -fill -fma -math -reduce -detect — implementation pending diff --git "a/crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" "b/crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" deleted file mode 100644 index 666cef3..0000000 --- "a/crates/mohu-sparse/src/arith\nbsr\ncoo\ncsc\ncsr\nconvert\ndia\nlinalg\nslice\nspmm\nspmv.rs" +++ /dev/null @@ -1,11 +0,0 @@ -// arith -bsr -coo -csc -csr -convert -dia -linalg -slice -spmm -spmv — implementation pending diff --git "a/crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" "b/crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" deleted file mode 100644 index 6d83c44..0000000 --- "a/crates/mohu-special/src/beta\nbessel\nerf\nexpint\ngamma\nmisc\nstats_fn\ntrig.rs" +++ /dev/null @@ -1,8 +0,0 @@ -// beta -bessel -erf -expint -gamma -misc -stats_fn -trig — implementation pending diff --git "a/crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" "b/crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" deleted file mode 100644 index 3917164..0000000 --- "a/crates/mohu-testing/src/approx\nassert\ndtype\nfixtures\ngen\nperf.rs" +++ /dev/null @@ -1,6 +0,0 @@ -// approx -assert -dtype -fixtures -gen -perf — implementation pending diff --git "a/crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" "b/crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" deleted file mode 100644 index 65d3c32..0000000 --- "a/crates/mohu-ufunc/src/broadcast\ndispatch\nloop_impl\nmacros\nmethods\nreduce\nresolver\ntraits.rs" +++ /dev/null @@ -1,8 +0,0 @@ -// broadcast -dispatch -loop_impl -macros -methods -reduce -resolver -traits — implementation pending From 92fbf3e0c931d0be88dcf2aa9e5fc88b753c5421 Mon Sep 17 00:00:00 2001 From: madhu-mitha-e Date: Mon, 1 Jun 2026 11:33:19 +0530 Subject: [PATCH 3/3] style(mohu-io): apply cargo fmt Signed-off-by: madhu-mitha-e --- crates/mohu-io/src/arrow.rs | 361 ------------------------------------ crates/mohu-io/src/csv.rs | 1 + crates/mohu-io/src/mmap.rs | 1 + crates/mohu-io/src/npy.rs | 1 + 4 files changed, 3 insertions(+), 361 deletions(-) diff --git a/crates/mohu-io/src/arrow.rs b/crates/mohu-io/src/arrow.rs index 1678de4..8b13789 100644 --- a/crates/mohu-io/src/arrow.rs +++ b/crates/mohu-io/src/arrow.rs @@ -1,362 +1 @@ -//! Arrow IPC serialization and deserialization for mohu buffers. -//! -//! Exposes four functions: -//! -//! | Function | Description | -//! |---------------------|--------------------------------------------------| -//! | [`read_ipc_file`] | Read an Arrow IPC file into a `Vec` | -//! | [`read_ipc_stream`] | Read an Arrow IPC stream into a `Vec` | -//! | [`write_ipc_file`] | Write `Buffer` slices to an Arrow IPC file | -//! | [`write_ipc_stream`]| Write `Buffer` slices to an Arrow IPC stream | -//! -//! # Arrow IPC formats -//! -//! - **File format**: random-access, seekable, has a footer. Suited for -//! on-disk storage. -//! - **Streaming format**: sequential, no footer. Required for inter-process -//! streaming (Polars pipe, DuckDB query results, Arrow Flight protocol). -use std::io::{Read, Seek, Write}; -use std::sync::Arc; - -use arrow::array::{ArrayRef, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, UInt16Array, UInt32Array, UInt64Array, UInt8Array}; -use arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; -use arrow::ipc::reader::{FileReader, StreamReader}; -use arrow::ipc::writer::{FileWriter, StreamWriter}; -use arrow::record_batch::RecordBatch; - -use mohu_core::mohu_buffer::Buffer; -use mohu_core::mohu_dtype::DType; -use mohu_core::mohu_error::{MohuError, MohuResult}; - -// ── Helpers ─────────────────────────────────────────────────────────────────── - -/// Maps a mohu `DType` to an Arrow `DataType`. -fn dtype_to_arrow(dtype: DType) -> MohuResult { - match dtype { - DType::I8 => Ok(ArrowDataType::Int8), - DType::I16 => Ok(ArrowDataType::Int16), - DType::I32 => Ok(ArrowDataType::Int32), - DType::I64 => Ok(ArrowDataType::Int64), - DType::U8 => Ok(ArrowDataType::UInt8), - DType::U16 => Ok(ArrowDataType::UInt16), - DType::U32 => Ok(ArrowDataType::UInt32), - DType::U64 => Ok(ArrowDataType::UInt64), - DType::F32 => Ok(ArrowDataType::Float32), - DType::F64 => Ok(ArrowDataType::Float64), - other => Err(MohuError::UnsupportedDType { - op: "arrow IPC", - dtype: format!("{other:?}"), - }), - } -} - -/// Maps an Arrow `DataType` to a mohu `DType`. -/// Maps an Arrow `DataType` to a mohu `DType`. - - -/// Converts a 1-D `Buffer` to an Arrow `ArrayRef`. -fn buffer_to_array(buf: &Buffer) -> MohuResult { - match buf.dtype() { - DType::I8 => { - let s = buf.as_slice::()?; - Ok(Arc::new(Int8Array::from(s.to_vec()))) - } - DType::I16 => { - let s = buf.as_slice::()?; - Ok(Arc::new(Int16Array::from(s.to_vec()))) - } - DType::I32 => { - let s = buf.as_slice::()?; - Ok(Arc::new(Int32Array::from(s.to_vec()))) - } - DType::I64 => { - let s = buf.as_slice::()?; - Ok(Arc::new(Int64Array::from(s.to_vec()))) - } - DType::U8 => { - let s = buf.as_slice::()?; - Ok(Arc::new(UInt8Array::from(s.to_vec()))) - } - DType::U16 => { - let s = buf.as_slice::()?; - Ok(Arc::new(UInt16Array::from(s.to_vec()))) - } - DType::U32 => { - let s = buf.as_slice::()?; - Ok(Arc::new(UInt32Array::from(s.to_vec()))) - } - DType::U64 => { - let s = buf.as_slice::()?; - Ok(Arc::new(UInt64Array::from(s.to_vec()))) - } - DType::F32 => { - let s = buf.as_slice::()?; - Ok(Arc::new(Float32Array::from(s.to_vec()))) - } - DType::F64 => { - let s = buf.as_slice::()?; - Ok(Arc::new(Float64Array::from(s.to_vec()))) - } - other => Err(MohuError::UnsupportedDType { - op: "arrow IPC", - dtype: format!("{other:?}"), - }), - } -} - -/// Converts an Arrow `ArrayRef` to a 1-D `Buffer`. -fn array_to_buffer(array: &ArrayRef) -> MohuResult { - match array.data_type() { - ArrowDataType::Int8 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast Int8Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::Int16 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast Int16Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::Int32 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast Int32Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::Int64 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast Int64Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::UInt8 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast UInt8Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::UInt16 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast UInt16Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::UInt32 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast UInt32Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::UInt64 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast UInt64Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::Float32 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast Float32Array failed"))?; - Buffer::from_slice(a.values()) - } - ArrowDataType::Float64 => { - let a = array.as_any().downcast_ref::() - .ok_or_else(|| MohuError::bug("downcast Float64Array failed"))?; - Buffer::from_slice(a.values()) - } - other => Err(MohuError::UnsupportedDType { - op: "arrow IPC", - dtype: format!("{other:?}"), - }), - } -} - -/// Builds a `RecordBatch` from a slice of `Buffer`s. -fn buffers_to_record_batch(buffers: &[Buffer]) -> MohuResult { - let mut fields = Vec::with_capacity(buffers.len()); - let mut arrays = Vec::with_capacity(buffers.len()); - - for (i, buf) in buffers.iter().enumerate() { - let arrow_dtype = dtype_to_arrow(buf.dtype())?; - let field = Field::new(format!("col_{i}"), arrow_dtype, false); - fields.push(field); - arrays.push(buffer_to_array(buf)?); - } - - let schema = Arc::new(Schema::new(fields)); - RecordBatch::try_new(schema, arrays) - .map_err(|e| MohuError::ArrowSchema(e.to_string())) -} - -/// Extracts `Buffer`s from a `RecordBatch`. -fn record_batch_to_buffers(batch: &RecordBatch) -> MohuResult> { - batch - .columns() - .iter() - .map(array_to_buffer) - .collect() -} - -// ── Public API ──────────────────────────────────────────────────────────────── - -/// Reads an Arrow IPC **file** from `reader` and returns the columns of -/// every record batch as a flat `Vec`. -/// -/// The Arrow IPC file format is random-access and seekable, making it -/// suitable for on-disk storage. -/// -/// # Errors -/// -/// Returns an error if the IPC data is malformed or contains unsupported -/// Arrow data types. -pub fn read_ipc_file(reader: R) -> MohuResult> { - let file_reader = FileReader::try_new(reader, None) - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - let mut buffers = Vec::new(); - for batch in file_reader { - let batch = batch.map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - buffers.extend(record_batch_to_buffers(&batch)?); - } - Ok(buffers) -} - -/// Reads an Arrow IPC **stream** from `reader` and returns the columns of -/// every record batch as a flat `Vec`. -/// -/// The Arrow IPC streaming format is sequential with no footer, making it -/// suitable for inter-process streaming (Polars pipe, DuckDB query results, -/// Arrow Flight protocol). -/// -/// # Errors -/// -/// Returns an error if the IPC data is malformed or contains unsupported -/// Arrow data types. -pub fn read_ipc_stream(reader: R) -> MohuResult> { - let stream_reader = StreamReader::try_new(reader, None) - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - let mut buffers = Vec::new(); - for batch in stream_reader { - let batch = batch.map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - buffers.extend(record_batch_to_buffers(&batch)?); - } - Ok(buffers) -} - -/// Writes `buffers` to `writer` in Arrow IPC **file** format. -/// -/// All buffers are written as a single record batch where each buffer -/// becomes one column (`col_0`, `col_1`, …). -/// -/// # Errors -/// -/// Returns an error if any buffer has an unsupported dtype or if writing -/// fails. -pub fn write_ipc_file(writer: W, buffers: &[Buffer]) -> MohuResult<()> { - if buffers.is_empty() { - return Ok(()); - } - - let batch = buffers_to_record_batch(buffers)?; - let schema = batch.schema(); - - let mut file_writer = FileWriter::try_new(writer, &schema) - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - file_writer - .write(&batch) - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - file_writer - .finish() - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - Ok(()) -} - -/// Writes `buffers` to `writer` in Arrow IPC **stream** format. -/// -/// All buffers are written as a single record batch where each buffer -/// becomes one column (`col_0`, `col_1`, …). -/// -/// # Errors -/// -/// Returns an error if any buffer has an unsupported dtype or if writing -/// fails. -pub fn write_ipc_stream(writer: W, buffers: &[Buffer]) -> MohuResult<()> { - if buffers.is_empty() { - return Ok(()); - } - - let batch = buffers_to_record_batch(buffers)?; - let schema = batch.schema(); - - let mut stream_writer = StreamWriter::try_new(writer, &schema) - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - stream_writer - .write(&batch) - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - stream_writer - .finish() - .map_err(|e| MohuError::ArrowIpc(e.to_string()))?; - - Ok(()) -} - -// ── Tests ───────────────────────────────────────────────────────────────────── - -#[cfg(test)] -mod tests { - use super::*; - use std::io::Cursor; - - #[test] - fn round_trip_ipc_file_f32() { - let buf = Buffer::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]).unwrap(); - let mut bytes = Vec::new(); - write_ipc_file(&mut bytes, &[buf.clone()]).unwrap(); - - let result = read_ipc_file(Cursor::new(bytes)).unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0].as_slice::().unwrap(), &[1.0_f32, 2.0, 3.0, 4.0]); - } - - #[test] - fn round_trip_ipc_stream_f64() { - let buf = Buffer::from_slice(&[10.0_f64, 20.0, 30.0]).unwrap(); - let mut bytes = Vec::new(); - write_ipc_stream(&mut bytes, &[buf.clone()]).unwrap(); - - let result = read_ipc_stream(Cursor::new(bytes)).unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0].as_slice::().unwrap(), &[10.0_f64, 20.0, 30.0]); - } - - #[test] - fn round_trip_ipc_file_i32() { - let buf = Buffer::from_slice(&[1_i32, 2, 3, 4, 5]).unwrap(); - let mut bytes = Vec::new(); - write_ipc_file(&mut bytes, &[buf.clone()]).unwrap(); - - let result = read_ipc_file(Cursor::new(bytes)).unwrap(); - assert_eq!(result.len(), 1); - assert_eq!(result[0].as_slice::().unwrap(), &[1_i32, 2, 3, 4, 5]); - } - - #[test] - fn round_trip_ipc_stream_multiple_buffers() { - let buf1 = Buffer::from_slice(&[1_i32, 2, 3]).unwrap(); - let buf2 = Buffer::from_slice(&[4.0_f64, 5.0, 6.0]).unwrap(); - let mut bytes = Vec::new(); - write_ipc_stream(&mut bytes, &[buf1, buf2]).unwrap(); - - let result = read_ipc_stream(Cursor::new(bytes)).unwrap(); - assert_eq!(result.len(), 2); - assert_eq!(result[0].as_slice::().unwrap(), &[1_i32, 2, 3]); - assert_eq!(result[1].as_slice::().unwrap(), &[4.0_f64, 5.0, 6.0]); - } - - #[test] - fn write_empty_buffers_is_ok() { - let mut bytes = Vec::new(); - assert!(write_ipc_file(&mut bytes, &[]).is_ok()); - assert!(write_ipc_stream(&mut bytes, &[]).is_ok()); - } -} diff --git a/crates/mohu-io/src/csv.rs b/crates/mohu-io/src/csv.rs index e69de29..8b13789 100644 --- a/crates/mohu-io/src/csv.rs +++ b/crates/mohu-io/src/csv.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-io/src/mmap.rs b/crates/mohu-io/src/mmap.rs index e69de29..8b13789 100644 --- a/crates/mohu-io/src/mmap.rs +++ b/crates/mohu-io/src/mmap.rs @@ -0,0 +1 @@ + diff --git a/crates/mohu-io/src/npy.rs b/crates/mohu-io/src/npy.rs index e69de29..8b13789 100644 --- a/crates/mohu-io/src/npy.rs +++ b/crates/mohu-io/src/npy.rs @@ -0,0 +1 @@ +