diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 2fad41a..f5693ae 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -725,7 +725,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" dependencies = [ "chrono", - "phf", + "phf 0.12.1", ] [[package]] @@ -2207,6 +2207,8 @@ dependencies = [ "arrow", "arrow-schema", "num-traits", + "phf 0.13.1", + "time", "vec64", ] @@ -2330,6 +2332,15 @@ dependencies = [ "libm", ] +[[package]] +name = "num_threads" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c7398b9c8b70908f6371f47ed36737907c87c52af34c268fed0bf0ceb92ead9" +dependencies = [ + "libc", +] + [[package]] name = "objc2-core-foundation" version = "0.3.2" @@ -2545,7 +2556,41 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" dependencies = [ - "phf_shared", + "phf_shared 0.12.1", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_macros", + "phf_shared 0.13.1", + "serde", +] + +[[package]] +name = "phf_generator" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "135ace3a761e564ec88c03a77317a7c6b80bb7f7135ef2544dbe054243b89737" +dependencies = [ + "fastrand", + "phf_shared 0.13.1", +] + +[[package]] +name = "phf_macros" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "812f032b54b1e759ccd5f8b6677695d5268c588701effba24601f6932f8269ef" +dependencies = [ + "phf_generator", + "phf_shared 0.13.1", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -2557,6 +2602,15 @@ dependencies = [ "siphasher", ] +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -4317,7 +4371,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cdb87b95ec50ddfa440816d227a17b2ccbdda963a316a727fda0fc4334f7d134" dependencies = [ "deranged", + "libc", "num-conv", + "num_threads", "powerfmt", "serde_core", "time-core", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 5f8e094..22e07e4 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -157,7 +157,7 @@ required-features = ["mmap"] default_categorical_8 = ["minarrow/default_categorical_8"] extended_categorical = ["default_categorical_8", "minarrow/extended_categorical"] large_string = ["minarrow/large_string"] -datetime = ["minarrow/datetime"] +datetime = ["minarrow/datetime", "minarrow/datetime_ops"] decimal = ["minarrow/decimal"] extended_numeric_types = ["minarrow/extended_numeric_types"] lbuffer = ["minarrow/lbuffer"] diff --git a/rust/benches/file/chunked_throughput.rs b/rust/benches/file/chunked_throughput.rs index e5abf9a..b884e9c 100644 --- a/rust/benches/file/chunked_throughput.rs +++ b/rust/benches/file/chunked_throughput.rs @@ -124,7 +124,7 @@ fn bench_chunked_arrow(c: &mut Criterion) { // throughput measurements. let read_dir = fresh_dir("arrow_read"); { - let mut w = ChunkedArrowWriter::new(&read_dir, BASE).unwrap(); + let mut w = ChunkedArrowWriter::new(&read_dir, BASE, None).unwrap(); for t in &tables { w.write_chunk(t).unwrap(); } @@ -142,14 +142,14 @@ fn bench_chunked_arrow(c: &mut Criterion) { // // Directory cleanup runs in the next iteration's setup and is not timed. let write_serial_logical = |dir: PathBuf| { - let mut w = ChunkedArrowWriter::new(&dir, BASE).unwrap(); + let mut w = ChunkedArrowWriter::new(&dir, BASE, None).unwrap(); for t in &tables { w.write_chunk(t).unwrap(); } }; let write_serial_physical = |dir: PathBuf| { - let mut w = ChunkedArrowWriter::new(&dir, BASE).unwrap(); + let mut w = ChunkedArrowWriter::new(&dir, BASE, None).unwrap(); let mut paths: Vec = Vec::with_capacity(N_CHUNKS); for t in &tables { @@ -183,12 +183,12 @@ fn bench_chunked_arrow(c: &mut Criterion) { }); let par_write_logical = |dir: PathBuf| { - let w = ChunkedArrowWriter::new(&dir, BASE).unwrap(); + let w = ChunkedArrowWriter::new(&dir, BASE, None).unwrap(); let paths = w.par_write_all(&table_refs, None).unwrap(); assert_eq!(paths.len(), N_CHUNKS); }; let par_write_physical = |dir: PathBuf| { - let w = ChunkedArrowWriter::new(&dir, BASE).unwrap(); + let w = ChunkedArrowWriter::new(&dir, BASE, None).unwrap(); let paths = w.par_write_all(&table_refs, None).unwrap(); assert_eq!(paths.len(), N_CHUNKS); for p in &paths { @@ -269,7 +269,7 @@ fn bench_chunked_parquet(c: &mut Criterion) { // Pre-write the dataset once to learn the real physical byte count. let read_dir = fresh_dir("parquet_read"); { - let mut w = ChunkedParquetWriter::new(&read_dir, BASE, None).unwrap(); + let mut w = ChunkedParquetWriter::new(&read_dir, BASE, None, None).unwrap(); for t in &tables { w.write_chunk(t).unwrap(); } @@ -285,13 +285,13 @@ fn bench_chunked_parquet(c: &mut Criterion) { // dir. Denominator = output file bytes. Cleanup of the chunk dir // is NOT in the timed region. let write_serial_logical = |dir: PathBuf| { - let mut w = ChunkedParquetWriter::new(&dir, BASE, None).unwrap(); + let mut w = ChunkedParquetWriter::new(&dir, BASE, None, None).unwrap(); for t in &tables { w.write_chunk(t).unwrap(); } }; let write_serial_physical = |dir: PathBuf| { - let mut w = ChunkedParquetWriter::new(&dir, BASE, None).unwrap(); + let mut w = ChunkedParquetWriter::new(&dir, BASE, None, None).unwrap(); let mut paths: Vec = Vec::with_capacity(N_CHUNKS); for t in &tables { paths.push(w.write_chunk(t).unwrap()); @@ -321,12 +321,12 @@ fn bench_chunked_parquet(c: &mut Criterion) { }); let par_write_logical = |dir: PathBuf| { - let w = ChunkedParquetWriter::new(&dir, BASE, None).unwrap(); + let w = ChunkedParquetWriter::new(&dir, BASE, None, None).unwrap(); let paths = w.par_write_all(&table_refs, None).unwrap(); assert_eq!(paths.len(), N_CHUNKS); }; let par_write_physical = |dir: PathBuf| { - let w = ChunkedParquetWriter::new(&dir, BASE, None).unwrap(); + let w = ChunkedParquetWriter::new(&dir, BASE, None, None).unwrap(); let paths = w.par_write_all(&table_refs, None).unwrap(); assert_eq!(paths.len(), N_CHUNKS); for p in &paths { @@ -405,7 +405,7 @@ fn bench_chunked_csv(c: &mut Criterion) { // Pre-write the dataset once to learn the real physical byte count. let read_dir = fresh_dir("csv_read"); { - let mut w = ChunkedCsvWriter::new(&read_dir, BASE, CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&read_dir, BASE, CsvEncodeOptions::default(), None).unwrap(); for t in &tables { w.write_chunk(t).unwrap(); } @@ -420,14 +420,14 @@ fn bench_chunked_csv(c: &mut Criterion) { // Physical writes synchronise each chunk and the parent directory, and use the // encoded file size. Directory cleanup runs outside the timed region. let write_serial_logical = |dir: PathBuf| { - let mut w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default(), None).unwrap(); for t in &tables { w.write_chunk(t).unwrap(); } }; let write_serial_physical = |dir: PathBuf| { - let mut w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default(), None).unwrap(); let mut paths: Vec = Vec::with_capacity(N_CHUNKS); for t in &tables { paths.push(w.write_chunk(t).unwrap()); @@ -457,12 +457,12 @@ fn bench_chunked_csv(c: &mut Criterion) { }); let par_write_logical = |dir: PathBuf| { - let w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default()).unwrap(); + let w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default(), None).unwrap(); let paths = w.par_write_all(&table_refs, None).unwrap(); assert_eq!(paths.len(), N_CHUNKS); }; let par_write_physical = |dir: PathBuf| { - let w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default()).unwrap(); + let w = ChunkedCsvWriter::new(&dir, BASE, CsvEncodeOptions::default(), None).unwrap(); let paths = w.par_write_all(&table_refs, None).unwrap(); assert_eq!(paths.len(), N_CHUNKS); for p in &paths { diff --git a/rust/scripts/run-object-storage-s3-tests.sh b/rust/scripts/run-object-storage-s3-tests.sh index b7e5cc8..a78e71f 100755 --- a/rust/scripts/run-object-storage-s3-tests.sh +++ b/rust/scripts/run-object-storage-s3-tests.sh @@ -23,4 +23,4 @@ export AWS_REGION="${AWS_REGION:-$(aws configure get region || true)}" : "${AWS_REGION:?set AWS_REGION to the region of the bucket}" echo "Running object storage tests against ${LIGHTSTREAM_S3_URL} in ${AWS_REGION}" -cargo test --features "aws parquet csv json" --test object_storage_s3 -- --ignored "$@" +cargo test --features "aws datetime parquet csv json" --test object_storage_s3 -- --ignored "$@" diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 4382f06..fa33880 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -250,7 +250,8 @@ pub mod models { pub mod parquet; /// Chunked file readers - glob a directory of `-NNNNN.` - /// files and present them as an ordered iterator of `Table`s. + /// files, including its partition directories, and present them as + /// an ordered iterator of `Table`s. pub mod chunked { /// Chunked CSV reader. #[cfg(feature = "csv")] @@ -360,6 +361,9 @@ pub mod models { /// Chunked Arrow IPC writer. Holds a tokio runtime internally /// to drive the async IPC file writer. pub mod arrow; + + /// Partition paths for the chunked writers. + pub mod partition; } /// TCP table writer diff --git a/rust/src/models/readers/chunked/arrow.rs b/rust/src/models/readers/chunked/arrow.rs index e3fa761..832cc68 100644 --- a/rust/src/models/readers/chunked/arrow.rs +++ b/rust/src/models/readers/chunked/arrow.rs @@ -9,8 +9,10 @@ //! Reads a directory of `-NNNNNNNNNN.arrow` files emitted by //! `ChunkedArrowWriter` and presents them as an ordered iterator of //! `Table`s. Each chunk file is a complete, independently readable Arrow -//! IPC file. The reader sorts files by their numeric index so consumers -//! see batches in write order. +//! IPC file. The reader includes files in nested partition +//! directories and sorts files by partition directory and then by numeric +//! index. Consumers of an unpartitioned directory see batches in write +//! order. //! //! ## Concurrency model //! @@ -36,7 +38,7 @@ use std::sync::Arc; use minarrow::{Concatenate, SuperTable, Table}; use crate::models::readers::ipc::file_table::FileTableReader; -use crate::models::readers::object_storage::list_dir; +use crate::models::readers::object_storage::list_files; use crate::traits::chunked_table_reader::ChunkedTableReader; /// Iterator over chunk files in a directory written by @@ -64,7 +66,7 @@ impl ChunkedTableReader for ChunkedArrowReader { fn list_paths>(dir: P, base: &str) -> io::Result> { let prefix = format!("{base}-"); let mut indexed: Vec<(u64, PathBuf)> = Vec::new(); - for path in list_dir(dir)? { + for path in list_files(dir)? { let Some(name) = path.file_name().and_then(|s| s.to_str()) else { continue; }; @@ -77,7 +79,12 @@ impl ChunkedTableReader for ChunkedArrowReader { }; indexed.push((index, path)); } - indexed.sort_by_key(|(i, _)| *i); + // Order by partition directory, then by index within it, which is + // write order for an unpartitioned directory and time order for + // `PartitionType::Auto` partitions. + indexed.sort_by(|(a_index, a), (b_index, b)| { + a.parent().cmp(&b.parent()).then(a_index.cmp(b_index)) + }); Ok(indexed.into_iter().map(|(_, p)| p).collect()) } @@ -140,16 +147,31 @@ mod tests { use std::fs; use super::*; + use std::sync::Arc; + use crate::models::writers::chunked::arrow::ChunkedArrowWriter; + use crate::models::writers::chunked::partition::PartitionType; use crate::traits::chunked_table_writer::ChunkedTableWriter; - use minarrow::{Table, fa_i32}; + use minarrow::{Table, Vec64, fa_i32}; + + /// 2026-02-09T07:00:00Z and 08:00:00Z in epoch nanoseconds. + #[cfg(feature = "datetime")] + const SEVEN_AM: u64 = 1_770_620_400_000_000_000; + #[cfg(feature = "datetime")] + const EIGHT_AM: u64 = SEVEN_AM + 3_600_000_000_000; + + /// Table of `n_rows` rows, identifying each chunk by its row count. + fn rows(n_rows: i32) -> Table { + let values: Vec64 = (0..n_rows).collect(); + Table::new("b", Some(vec![fa_i32!("n", @vec64 values)])) + } #[test] fn par_load_batched_returns_batches_in_write_order() { let dir = std::env::temp_dir().join("lightstream_chunked_arrow_par_reader"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedArrowWriter::new(&dir, "part").unwrap(); + let mut w = ChunkedArrowWriter::new(&dir, "part", None).unwrap(); for i in 0..16i32 { w.write_chunk(&Table::new( "b", @@ -187,4 +209,72 @@ mod tests { fs::remove_dir_all(&dir).ok(); } + + #[cfg(feature = "datetime")] + #[test] + fn auto_partitioned_chunks_read_back_by_partition_then_index() { + let dir = std::env::temp_dir().join("lightstream_chunked_arrow_auto_partition"); + let _ = fs::remove_dir_all(&dir); + + let mut w = ChunkedArrowWriter::new(&dir, "part", Some(PartitionType::Auto)).unwrap(); + let p0 = w.write_chunk_for_time(&rows(1), EIGHT_AM).unwrap(); + let p1 = w.write_chunk_for_time(&rows(2), SEVEN_AM).unwrap(); + let p2 = w.write_chunk_for_time(&rows(3), EIGHT_AM).unwrap(); + assert_eq!(p0, dir.join("date=2026-02-09/hour=08/part-0000000000.arrow")); + assert_eq!(p1, dir.join("date=2026-02-09/hour=07/part-0000000001.arrow")); + assert_eq!(p2, dir.join("date=2026-02-09/hour=08/part-0000000002.arrow")); + + // The 07:00 chunk precedes both 08:00 chunks, which keep write order. + let all = ChunkedArrowReader::open(&dir, "part", ()).unwrap().load_batched().unwrap(); + let n_rows: Vec = all.batches.iter().map(|batch| batch.n_rows).collect(); + assert_eq!(n_rows, vec![2, 1, 3]); + + // Opening a partition directory reads only that partition. + let eight = ChunkedArrowReader::open(dir.join("date=2026-02-09/hour=08"), "part", ()) + .unwrap() + .load_batched() + .unwrap(); + let n_rows: Vec = eight.batches.iter().map(|batch| batch.n_rows).collect(); + assert_eq!(n_rows, vec![1, 3]); + + fs::remove_dir_all(&dir).ok(); + } + + #[test] + fn custom_partitions_come_from_the_table() { + let dir = std::env::temp_dir().join("lightstream_chunked_arrow_custom_partition"); + let _ = fs::remove_dir_all(&dir); + + let by_size = PartitionType::Custom(Arc::new(|table: &Table, _time: u64| { + format!("rows={}", table.n_rows) + })); + let mut w = ChunkedArrowWriter::new(&dir, "part", Some(by_size)).unwrap(); + let path = w.write_chunk(&rows(4)).unwrap(); + assert_eq!(path, dir.join("rows=4/part-0000000000.arrow")); + + fs::remove_dir_all(&dir).ok(); + } + + #[cfg(feature = "datetime")] + #[test] + fn par_write_all_creates_partition_directories() { + let dir = std::env::temp_dir().join("lightstream_chunked_arrow_par_partition"); + let _ = fs::remove_dir_all(&dir); + + let w = ChunkedArrowWriter::new(&dir, "part", Some(PartitionType::Auto)).unwrap(); + let tables: Vec = (1..=6).map(rows).collect(); + let refs: Vec<&Table> = tables.iter().collect(); + let paths = w.par_write_all(&refs, Some(3)).unwrap(); + + // Batches written in one call share a single partition. + let partition = paths[0].parent().unwrap(); + assert!(partition.starts_with(&dir) && partition != dir); + assert!(paths.iter().all(|path| path.parent() == Some(partition))); + + let read = ChunkedArrowReader::par_load_batched(&dir, "part", (), None).unwrap(); + let n_rows: Vec = read.batches.iter().map(|batch| batch.n_rows).collect(); + assert_eq!(n_rows, vec![1, 2, 3, 4, 5, 6]); + + fs::remove_dir_all(&dir).ok(); + } } diff --git a/rust/src/models/readers/chunked/csv.rs b/rust/src/models/readers/chunked/csv.rs index 7ad2c80..aa14e9a 100644 --- a/rust/src/models/readers/chunked/csv.rs +++ b/rust/src/models/readers/chunked/csv.rs @@ -9,8 +9,10 @@ //! Reads a directory of `-NNNNNNNNNN.csv` files emitted by //! `ChunkedCsvWriter` and presents them as an ordered iterator of `Table`s. //! Each chunk file is a complete, independently parseable CSV (header + -//! rows). The reader sorts files by their numeric index so consumers see -//! batches in write order. +//! rows). The reader includes files in nested partition +//! directories and sorts files by partition directory and then by numeric +//! index. Consumers of an unpartitioned directory see batches in write +//! order. //! //! Inherits [`ChunkedTableReader::par_load_batched`] for sync parallel //! decode across chunk files. @@ -24,7 +26,7 @@ use minarrow::{ColumnSelection, Concatenate, SuperTable, Table}; use crate::models::decoders::csv::CsvDecodeOptions; use crate::models::readers::csv::CsvReader; -use crate::models::readers::object_storage::{StorageReader, list_dir}; +use crate::models::readers::object_storage::{StorageReader, list_files}; use crate::traits::chunked_table_reader::ChunkedTableReader; /// Per-format `Options` for `ChunkedCsvReader::open`. @@ -78,7 +80,7 @@ impl ChunkedTableReader for ChunkedCsvReader { fn list_paths>(dir: P, base: &str) -> io::Result> { let prefix = format!("{base}-"); let mut indexed: Vec<(u64, PathBuf)> = Vec::new(); - for path in list_dir(dir)? { + for path in list_files(dir)? { let Some(name) = path.file_name().and_then(|s| s.to_str()) else { continue; }; @@ -91,7 +93,12 @@ impl ChunkedTableReader for ChunkedCsvReader { }; indexed.push((index, path)); } - indexed.sort_by_key(|(i, _)| *i); + // Order by partition directory, then by index within it, which is + // write order for an unpartitioned directory and time order for + // `PartitionType::Auto` partitions. + indexed.sort_by(|(a_index, a), (b_index, b)| { + a.parent().cmp(&b.parent()).then(a_index.cmp(b_index)) + }); Ok(indexed.into_iter().map(|(_, p)| p).collect()) } @@ -171,7 +178,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_csv_test_reader"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default(), None).unwrap(); w.write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 1, 2, 3)]))) .unwrap(); w.write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 4, 5)]))) @@ -203,7 +210,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_csv_par_reader"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default(), None).unwrap(); for i in 0..12i32 { w.write_chunk(&Table::new( "b", @@ -244,7 +251,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_csv_test_iter"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default(), None).unwrap(); w.write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 10)]))) .unwrap(); w.write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 20, 21)]))) diff --git a/rust/src/models/readers/chunked/parquet.rs b/rust/src/models/readers/chunked/parquet.rs index 5715faa..f5abfee 100644 --- a/rust/src/models/readers/chunked/parquet.rs +++ b/rust/src/models/readers/chunked/parquet.rs @@ -9,8 +9,10 @@ //! Reads a directory of `-NNNNNNNNNN.parquet` files emitted by //! `ChunkedParquetWriter` and presents them as an ordered iterator of //! `Table`s. Each chunk file is a complete, independently parseable -//! Parquet file. The reader sorts files by their numeric index so -//! consumers see batches in write order. +//! Parquet file. The reader includes files in nested partition +//! directories and sorts files by partition directory and then by numeric +//! index. Consumers of an unpartitioned directory see batches in write +//! order. //! //! ## Concurrency model //! @@ -28,7 +30,7 @@ use minarrow::Table; use crate::error::IoError; use crate::models::readers::parquet::{load_parquet_table_cols, load_parquet_table}; -use crate::models::readers::object_storage::{StorageReader, list_dir}; +use crate::models::readers::object_storage::{StorageReader, list_files}; use crate::traits::chunked_table_reader::ChunkedTableReader; use std::sync::Arc; @@ -60,7 +62,7 @@ impl ChunkedTableReader for ChunkedParquetReader { fn list_paths>(dir: P, base: &str) -> Result, IoError> { let prefix = format!("{base}-"); let mut indexed: Vec<(u64, PathBuf)> = Vec::new(); - for path in list_dir(dir)? { + for path in list_files(dir)? { let Some(name) = path.file_name().and_then(|s| s.to_str()) else { continue; }; @@ -73,7 +75,12 @@ impl ChunkedTableReader for ChunkedParquetReader { }; indexed.push((index, path)); } - indexed.sort_by_key(|(i, _)| *i); + // Order by partition directory, then by index within it, which is + // write order for an unpartitioned directory and time order for + // `PartitionType::Auto` partitions. + indexed.sort_by(|(a_index, a), (b_index, b)| { + a.parent().cmp(&b.parent()).then(a_index.cmp(b_index)) + }); Ok(indexed.into_iter().map(|(_, p)| p).collect()) } @@ -126,7 +133,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_parquet_par_reader"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedParquetWriter::new(&dir, "part", None).unwrap(); + let mut w = ChunkedParquetWriter::new(&dir, "part", None, None).unwrap(); for i in 0..8i32 { w.write_chunk(&Table::new( "b", diff --git a/rust/src/models/readers/object_storage.rs b/rust/src/models/readers/object_storage.rs index c6d12fa..e818176 100644 --- a/rust/src/models/readers/object_storage.rs +++ b/rust/src/models/readers/object_storage.rs @@ -38,7 +38,9 @@ use std::path::{Path, PathBuf}; #[cfg(feature = "object_storage")] use bytes::Bytes; #[cfg(feature = "object_storage")] -use object_store::ObjectStoreExt; +use futures_util::TryStreamExt; +#[cfg(feature = "object_storage")] +use object_store::{ObjectMeta, ObjectStoreExt}; #[cfg(feature = "object_storage")] use crate::object_storage::{ObjectLocation, current_handle, io_error, object_url, resolve}; @@ -237,6 +239,44 @@ pub fn list_dir>(dir: P) -> io::Result> { .collect() } +/// Paths of every file beneath the directory or object prefix `dir`, +/// including files in nested directories, in no particular order. +/// +/// Object entries are returned as object URLs beneath `dir`, as for +/// [`list_dir`]. On object storage this is one listing of every object under +/// the prefix. +pub fn list_files>(dir: P) -> io::Result> { + #[cfg(feature = "object_storage")] + if let Some(url) = object_url(dir.as_ref()) { + let location = resolve(&url)?; + let objects: Vec = current_handle()? + .block_on(location.store.list(Some(&location.path)).try_collect()) + .map_err(io_error)?; + let base = url.trim_end_matches('/'); + return Ok(objects + .iter() + .filter_map(|meta| meta.location.prefix_match(&location.path)) + .map(|parts| { + let relative: Vec = parts.map(|part| part.as_ref().to_string()).collect(); + PathBuf::from(format!("{base}/{}", relative.join("/"))) + }) + .collect()); + } + let mut files = Vec::new(); + let mut pending = vec![dir.as_ref().to_path_buf()]; + while let Some(dir) = pending.pop() { + for entry in fs::read_dir(&dir)? { + let path = entry?.path(); + if path.is_dir() { + pending.push(path); + } else { + files.push(path); + } + } + } + Ok(files) +} + /// Read exactly `buf.len()` bytes starting at `offset` without touching /// the file's seek position. Backed by `pread(2)` on Unix and a /// `seek_read` loop on Windows so callers can issue many positional diff --git a/rust/src/models/writers/chunked/arrow.rs b/rust/src/models/writers/chunked/arrow.rs index 9fe5bab..a673844 100644 --- a/rust/src/models/writers/chunked/arrow.rs +++ b/rust/src/models/writers/chunked/arrow.rs @@ -32,6 +32,7 @@ use std::sync::atomic::AtomicU64; use minarrow::{Field, Table}; use crate::enums::IPCMessageProtocol; +use crate::models::writers::chunked::partition::PartitionType; use crate::models::writers::object_storage::{StorageWriter, create_dir_all}; use crate::models::writers::ipc::sync_table::SyncTableWriter; use crate::traits::chunked_table_writer::ChunkedTableWriter; @@ -46,6 +47,7 @@ pub struct ChunkedArrowWriter { dir: PathBuf, base: String, next_index: AtomicU64, + partition: Option, } impl ChunkedArrowWriter { @@ -55,13 +57,22 @@ impl ChunkedArrowWriter { /// has no per-chunk options at this layer (each file is a complete /// Arrow IPC File so external readers like pyarrow and arrow-rs can /// open it). - pub fn new>(dir: P, base: &str) -> io::Result { + /// + /// `partition` places each chunk in a partition directory beneath `dir`, + /// for example `/date=2026-10-14/hour=13/-NNNNNNNNNN.arrow`. + /// `None` writes every chunk into `dir`. + pub fn new>( + dir: P, + base: &str, + partition: Option, + ) -> io::Result { let dir = dir.as_ref().to_path_buf(); create_dir_all(&dir)?; Ok(Self { dir, base: base.to_string(), next_index: AtomicU64::new(0), + partition, }) } } @@ -90,6 +101,11 @@ impl ChunkedTableWriter for ChunkedArrowWriter { &self.next_index } + /// Partitioning applied to chunk paths, if configured. + fn partition(&self) -> Option<&PartitionType> { + self.partition.as_ref() + } + /// Encode `table` as a complete Arrow IPC File at `path`. Shared /// per-format encode path used by both `write_chunk` and /// `par_write_all`. @@ -117,7 +133,7 @@ mod tests { fn writes_indexed_chunk_files_and_round_trips() { let dir = std::env::temp_dir().join("lightstream_chunked_arrow_test_writer"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedArrowWriter::new(&dir, "part").unwrap(); + let mut w = ChunkedArrowWriter::new(&dir, "part", None).unwrap(); let p0 = w .write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 1, 2, 3)]))) .unwrap(); @@ -152,7 +168,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_arrow_par_writer"); let _ = fs::remove_dir_all(&dir); - let w = ChunkedArrowWriter::new(&dir, "part").unwrap(); + let w = ChunkedArrowWriter::new(&dir, "part", None).unwrap(); let tables: Vec
= (0..8i32) .map(|i| Table::new("b", Some(vec![fa_i32!("n", i, i + 100)]))) .collect(); diff --git a/rust/src/models/writers/chunked/csv.rs b/rust/src/models/writers/chunked/csv.rs index 4932726..28d4f85 100644 --- a/rust/src/models/writers/chunked/csv.rs +++ b/rust/src/models/writers/chunked/csv.rs @@ -26,6 +26,7 @@ use std::sync::atomic::AtomicU64; use minarrow::Table; use crate::models::encoders::csv::CsvEncodeOptions; +use crate::models::writers::chunked::partition::PartitionType; use crate::models::writers::object_storage::{StorageWriter, create_dir_all}; use crate::models::writers::csv::CsvWriter; use crate::traits::chunked_table_writer::ChunkedTableWriter; @@ -42,13 +43,23 @@ pub struct ChunkedCsvWriter { base: String, options: CsvEncodeOptions, next_index: AtomicU64, + partition: Option, } impl ChunkedCsvWriter { /// Create a new chunked writer rooted at `dir`. Each chunk filename /// will be `-NNNNNNNNNN.csv` with a 10-digit zero-padded /// index. The directory is created if it does not exist. - pub fn new>(dir: P, base: &str, options: CsvEncodeOptions) -> io::Result { + /// + /// `partition` places each chunk in a partition directory beneath `dir`, + /// for example `/date=2026-10-14/hour=13/-NNNNNNNNNN.csv`. + /// `None` writes every chunk into `dir`. + pub fn new>( + dir: P, + base: &str, + options: CsvEncodeOptions, + partition: Option, + ) -> io::Result { let dir = dir.as_ref().to_path_buf(); create_dir_all(&dir)?; Ok(Self { @@ -56,6 +67,7 @@ impl ChunkedCsvWriter { base: base.to_string(), options, next_index: AtomicU64::new(0), + partition, }) } } @@ -84,6 +96,11 @@ impl ChunkedTableWriter for ChunkedCsvWriter { &self.next_index } + /// Partitioning applied to chunk paths, if configured. + fn partition(&self) -> Option<&PartitionType> { + self.partition.as_ref() + } + /// Encode `table` as a complete CSV file at `path` using the writer's /// configured encode options. Shared per-format encode path used by /// both `write_chunk` and `par_write_all`. @@ -109,7 +126,7 @@ mod tests { fn writes_indexed_chunk_files() { let dir = std::env::temp_dir().join("lightstream_chunked_csv_test_writer"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default()).unwrap(); + let mut w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default(), None).unwrap(); let p0 = w .write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 0, 1, 2)]))) .unwrap(); @@ -143,7 +160,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_csv_par_writer"); let _ = fs::remove_dir_all(&dir); - let w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default()).unwrap(); + let w = ChunkedCsvWriter::new(&dir, "part", CsvEncodeOptions::default(), None).unwrap(); let tables: Vec
= (0..10i32) .map(|i| Table::new("b", Some(vec![fa_i32!("n", i, i + 100)]))) .collect(); diff --git a/rust/src/models/writers/chunked/parquet.rs b/rust/src/models/writers/chunked/parquet.rs index 34bfac1..8e20b33 100644 --- a/rust/src/models/writers/chunked/parquet.rs +++ b/rust/src/models/writers/chunked/parquet.rs @@ -26,6 +26,7 @@ use minarrow::Table; use crate::compression::Compression; use crate::error::IoError; +use crate::models::writers::chunked::partition::PartitionType; use crate::models::writers::object_storage::{StorageWriter, create_dir_all}; use crate::models::writers::parquet::write_parquet_table; use crate::traits::chunked_table_writer::ChunkedTableWriter; @@ -37,16 +38,22 @@ pub struct ChunkedParquetWriter { base: String, compression: Option, next_index: AtomicU64, + partition: Option, } impl ChunkedParquetWriter { /// Create a new chunked writer rooted at `dir`. Each chunk filename /// will be `-NNNNNNNNNN.parquet` with a 10-digit zero-padded /// index. The directory is created if it does not exist. + /// + /// `partition` places each chunk in a partition directory beneath `dir`, + /// for example `/date=2026-10-14/hour=13/-NNNNNNNNNN.parquet`. + /// `None` writes every chunk into `dir`. pub fn new>( dir: P, base: &str, compression: Option, + partition: Option, ) -> Result { let dir = dir.as_ref().to_path_buf(); create_dir_all(&dir)?; @@ -55,6 +62,7 @@ impl ChunkedParquetWriter { base: base.to_string(), compression, next_index: AtomicU64::new(0), + partition, }) } } @@ -83,6 +91,11 @@ impl ChunkedTableWriter for ChunkedParquetWriter { &self.next_index } + /// Partitioning applied to chunk paths, if configured. + fn partition(&self) -> Option<&PartitionType> { + self.partition.as_ref() + } + /// Encode `table` as a complete Parquet file at `path` using the /// writer's configured compression. Shared per-format encode path /// used by both `write_chunk` and `par_write_all`. @@ -108,7 +121,7 @@ mod tests { fn writes_indexed_chunk_files_and_round_trips() { let dir = std::env::temp_dir().join("lightstream_chunked_parquet_test_writer"); let _ = fs::remove_dir_all(&dir); - let mut w = ChunkedParquetWriter::new(&dir, "part", None).unwrap(); + let mut w = ChunkedParquetWriter::new(&dir, "part", None, None).unwrap(); let p0 = w .write_chunk(&Table::new("b", Some(vec![fa_i32!("n", 1, 2, 3)]))) .unwrap(); @@ -143,7 +156,7 @@ mod tests { let dir = std::env::temp_dir().join("lightstream_chunked_parquet_par_writer"); let _ = fs::remove_dir_all(&dir); - let w = ChunkedParquetWriter::new(&dir, "part", None).unwrap(); + let w = ChunkedParquetWriter::new(&dir, "part", None, None).unwrap(); let tables: Vec
= (0..6i32) .map(|i| Table::new("b", Some(vec![fa_i32!("n", i, i + 100)]))) .collect(); diff --git a/rust/src/models/writers/chunked/partition.rs b/rust/src/models/writers/chunked/partition.rs new file mode 100644 index 0000000..c7fec95 --- /dev/null +++ b/rust/src/models/writers/chunked/partition.rs @@ -0,0 +1,132 @@ +// Copyright Peter G. Bower 2025-2026. +// +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. + +//! # Chunk partitioning +//! +//! Partition paths for the chunked writers. A partitioned writer places each +//! chunk in a subdirectory of its root, for example +//! `/date=2026-10-14/hour=13/-NNNNNNNNNN.parquet`. +//! +//! ## Context +//! Readers locate chunks by listing a directory or object prefix, and the +//! partition path determines how much of a dataset a read for one period of +//! time needs to list. This matters most on object storage, where each page +//! of a listing is a separate request. +//! +//! The chunked readers include chunks in nested partition directories. +//! Opening a reader on a partition path, such as `/date=2026-10-14`, +//! reads that partition alone. +//! +//! ## Partition types +//! - [`PartitionType::Auto`] derives the path from the time of each write in +//! the `key=value` directory convention, using minarrow's datetime support +//! under the `datetime` feature. Query engines such as DuckDB, +//! Polars and Spark recognise this convention, using it to prune by path +//! and exposing the values as columns. +//! - [`PartitionType::Custom`] takes the path from a caller-supplied +//! function, for schemes keyed on other columns such as venue or symbol. +//! +//! ## Time +//! Callers supply the time through +//! [`ChunkedTableWriter::write_chunk_for_time`](crate::traits::chunked_table_writer::ChunkedTableWriter::write_chunk_for_time), +//! for example the start of the event-time window that produced the chunk. +//! The other write methods use the current system time. + +use std::fmt; +use std::sync::Arc; + +use minarrow::Table; +#[cfg(feature = "datetime")] +use minarrow::{DatetimeArray, TimeUnit}; + +/// How a chunked writer groups its chunks into partition directories. +#[derive(Clone)] +pub enum PartitionType { + /// `date=YYYY-MM-DD/hour=HH` in UTC, from the time of each write. + /// + /// Hourly partitions bound the listing a reader needs for any period. + #[cfg(feature = "datetime")] + Auto, + /// Partition path returned by the function for each write, given the + /// Table and its time in epoch nanoseconds. + /// + /// The path is `/`-separated and relative to the writer's directory, + /// for example `venue=XNAS/date=2026-10-14`. + Custom(Arc String + Send + Sync>), +} + +impl PartitionType { + /// Partition path for a write of `table` at `time_ns` epoch nanoseconds. + pub fn partition(&self, table: &Table, time_ns: u64) -> String { + match self { + #[cfg(feature = "datetime")] + Self::Auto => auto_partition(time_ns), + Self::Custom(partition) => partition(table, time_ns), + } + } +} + +impl fmt::Debug for PartitionType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + #[cfg(feature = "datetime")] + Self::Auto => f.write_str("Auto"), + Self::Custom(_) => f.write_str("Custom"), + } + } +} + +/// Format an epoch-nanosecond time as `date=YYYY-MM-DD/hour=HH` in UTC. +#[cfg(feature = "datetime")] +fn auto_partition(time_ns: u64) -> String { + let time = DatetimeArray::from_slice(&[time_ns as i64], Some(TimeUnit::Nanoseconds)); + let (year, month, day, hour, ..) = time + .tuple_dt(0) + .expect("nanosecond timestamp within the supported date range"); + format!("date={year:04}-{month:02}-{day:02}/hour={hour:02}") +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 2026-02-09T07:48:36.123456789Z + #[cfg(feature = "datetime")] + const SAMPLE: u64 = 1_770_623_316_123_456_789; + + #[cfg(feature = "datetime")] + #[test] + fn auto_partition_formats_date_and_hour_in_utc() { + assert_eq!(auto_partition(SAMPLE), "date=2026-02-09/hour=07"); + } + + #[cfg(feature = "datetime")] + #[test] + fn auto_partition_at_the_epoch_is_the_first_hour_of_1970() { + assert_eq!(auto_partition(0), "date=1970-01-01/hour=00"); + } + + #[cfg(feature = "datetime")] + #[test] + fn auto_partition_handles_the_leap_day() { + // 2024-02-29T23:59:59Z + let time = 1_709_251_199_000_000_000; + assert_eq!(auto_partition(time), "date=2024-02-29/hour=23"); + assert_eq!( + auto_partition(time + 1_000_000_000), + "date=2024-03-01/hour=00" + ); + } + + #[test] + fn custom_partition_receives_the_table_and_time() { + let partition = PartitionType::Custom(Arc::new(|table: &Table, time: u64| { + format!("table={}/time={time}", table.name) + })); + let table = Table::new("trades".to_string(), None); + assert_eq!(partition.partition(&table, 42), "table=trades/time=42"); + } +} diff --git a/rust/src/traits/chunked_table_reader.rs b/rust/src/traits/chunked_table_reader.rs index 1064b80..db52a98 100644 --- a/rust/src/traits/chunked_table_reader.rs +++ b/rust/src/traits/chunked_table_reader.rs @@ -14,6 +14,11 @@ //! and a `load_table` that further consolidates the chunks into a //! single `Table`. The `*_cols` variants project to the named columns. //! +//! Chunks in nested partition directories, as written by a partitioned +//! chunked writer, are included and ordered by partition directory and then +//! by index. Opening a reader on one partition directory, such as +//! `/date=2026-10-14`, reads that partition alone. +//! //! `load_batched` deliberately does not consolidate; the per-chunk batches //! are preserved in the returned `SuperTable`. Callers wanting a single //! `Table` use `load_table` (which itself calls @@ -36,7 +41,8 @@ use tokio::runtime::Handle; /// Open / load_batched / iterate / par_load_batched shape shared by the /// per-format chunked readers. /// -/// Implementors yield one `Table` per chunk file in ascending index order. +/// Implementors yield one `Table` per chunk file, ordered by partition +/// directory and then by ascending index. /// Per-format extras (CSV decode options, batch size, etc.) flow through /// the associated `Options` type; formats with no extras use `()`. pub trait ChunkedTableReader: Iterator> + Sized { @@ -45,18 +51,19 @@ pub trait ChunkedTableReader: Iterator> + Size /// Per-format constructor options. Use `()` if no extras are needed. type Options; - /// Open every `-*.` file inside `dir`, sorted by numeric - /// suffix so chunks come out in write order. Files whose names don't - /// match the pattern are ignored. + /// Open every `-*.` file beneath `dir`, including files in + /// nested partition directories, sorted by partition directory and then + /// by numeric suffix. Chunks of an unpartitioned directory come out in + /// write order. Files whose names don't match the pattern are ignored. fn open>( dir: P, base: &str, options: Self::Options, ) -> Result; - /// List chunk file paths in `dir` matching this format's naming - /// pattern, sorted by ascending numeric index. Used by `open` to - /// populate the reader's path list. + /// List chunk file paths beneath `dir` matching this format's naming + /// pattern, sorted by partition directory and then by ascending numeric + /// index. Used by `open` to populate the reader's path list. fn list_paths>(dir: P, base: &str) -> Result, Self::Error>; /// Return the chunk paths held by this reader, in write order. diff --git a/rust/src/traits/chunked_table_writer.rs b/rust/src/traits/chunked_table_writer.rs index ae0fc4e..4cc66bf 100644 --- a/rust/src/traits/chunked_table_writer.rs +++ b/rust/src/traits/chunked_table_writer.rs @@ -25,14 +25,32 @@ //! write each batch end-to-end via the writer's `write_chunk_at`. Index //! reservation goes through the same counter as `write_chunk`, so the //! two paths compose without producing duplicate filenames. +//! +//! ## Partitioning +//! +//! Writers configured with a [`PartitionType`] place each chunk in a +//! partition directory beneath their root, for example +//! `/date=2026-10-14/hour=13/-NNNNNNNNNN.`. +//! +//! - [`ChunkedTableWriter::write_chunk_for_time`] takes the time used for +//! the partition path. +//! - `write_chunk`, `write_all` and `par_write_all` use the current system +//! time. The two `_all` methods use one time for every batch in the call. +//! - Partition directories are created on their first write. +//! - Chunk indices come from the writer's single counter in every +//! partition. use std::io; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; use minarrow::{SuperTable, Table}; use tokio::runtime::Handle; +use crate::models::writers::chunked::partition::PartitionType; +use crate::models::writers::object_storage::create_dir_all; + /// extension / dir / base / write_chunk / write_chunk_at / par_write_all /// shape shared by the per-format chunked writers. /// @@ -64,6 +82,10 @@ pub trait ChunkedTableWriter: Sized { /// reserve indices off this atom so the two paths compose. fn counter(&self) -> &AtomicU64; + /// Partitioning applied to chunk paths, or `None` for a single flat + /// directory. + fn partition(&self) -> Option<&PartitionType>; + /// Total number of chunks written so far. fn batches_written(&self) -> u64 { self.counter().load(Ordering::Relaxed) @@ -81,6 +103,21 @@ pub trait ChunkedTableWriter: Sized { )) } + /// Build the path for chunk `index` of `table`, inside the partition + /// directory for `time_ns` epoch nanoseconds when the writer is + /// partitioned. + /// + /// Without partitioning this is [`Self::chunk_path_for`]. + fn partitioned_chunk_path(&self, table: &Table, time_ns: u64, index: u64) -> PathBuf { + match self.partition() { + Some(partition) => self + .dir() + .join(partition.partition(table, time_ns)) + .join(format!("{}-{:010}.{}", self.base(), index, Self::extension())), + None => self.chunk_path_for(index), + } + } + /// Encode `table` into a complete chunk file at `path`, using any /// per-format options stored on the writer. Shared per-format encode /// path used both by `write_chunk` (one chunk at a time) and by @@ -89,9 +126,26 @@ pub trait ChunkedTableWriter: Sized { /// Reserve the next chunk index off the shared counter and write /// `table` to that chunk file. Returns the path of the written file. + /// + /// A partitioned writer takes the partition from the current system + /// time. See [`Self::write_chunk_for_time`] to supply the time. fn write_chunk(&mut self, table: &Table) -> Result { + self.write_chunk_for_time(table, current_time_ns()) + } + + /// Write `table` as the next chunk at `time_ns` epoch nanoseconds and + /// return the path of the written file. + /// + /// The time selects the partition directory, and is typically the start + /// of the event-time window that produced `table`. Replayed or late data + /// is then written to the same partition as the original data. Writers + /// without partitioning ignore the time. + fn write_chunk_for_time(&mut self, table: &Table, time_ns: u64) -> Result { let idx = self.counter().fetch_add(1, Ordering::Relaxed); - let path = self.chunk_path_for(idx); + let path = self.partitioned_chunk_path(table, time_ns, idx); + if self.partition().is_some() { + create_dir_all(path.parent().unwrap_or(self.dir()))?; + } self.write_chunk_at(&path, table)?; Ok(path) } @@ -100,9 +154,10 @@ pub trait ChunkedTableWriter: Sized { /// file, in order. Returns the paths of the written files. Advances /// the shared counter by `supertable.batches.len()`. fn write_all(&mut self, supertable: &SuperTable) -> Result, Self::Error> { + let time_ns = current_time_ns(); let mut paths = Vec::with_capacity(supertable.batches.len()); for batch in supertable.batches.iter() { - paths.push(self.write_chunk(batch.as_ref())?); + paths.push(self.write_chunk_for_time(batch.as_ref(), time_ns)?); } Ok(paths) } @@ -151,6 +206,7 @@ pub trait ChunkedTableWriter: Sized { // any concurrent `write_chunk` callers or with a later resumed // streaming write. let start = self.counter().fetch_add(n_files as u64, Ordering::Relaxed); + let time_ns = current_time_ns(); let cursor = AtomicU64::new(0); // Per-worker capacity hint: ceil(n_files / n_threads). Workers @@ -178,7 +234,10 @@ pub trait ChunkedTableWriter: Sized { break; } let idx = start + i; - let path = me.chunk_path_for(idx); + let path = me.partitioned_chunk_path(tables[i as usize], time_ns, idx); + if me.partition().is_some() { + create_dir_all(path.parent().unwrap_or(me.dir()))?; + } me.write_chunk_at(&path, tables[i as usize])?; local.push((i, path)); } @@ -200,3 +259,10 @@ pub trait ChunkedTableWriter: Sized { Ok(all.into_iter().map(|(_, p)| p).collect()) } } + +/// Current system time in epoch nanoseconds. +fn current_time_ns() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| elapsed.as_nanos() as u64) +} diff --git a/rust/tests/object_storage_checks/mod.rs b/rust/tests/object_storage_checks/mod.rs index 447e3f0..21f41f3 100644 --- a/rust/tests/object_storage_checks/mod.rs +++ b/rust/tests/object_storage_checks/mod.rs @@ -29,7 +29,7 @@ use lightstream::models::streams::object_store::{ObjectByteStream, ObjectWrite}; use lightstream::models::writers::ipc::sync_table::write_tables_to_file_sync; use lightstream::models::writers::ipc::table::{TableWriter, write_tables_to_file}; use lightstream::models::writers::object_storage::StorageWriter; -use minarrow::{Array, Field, FieldArray, IntegerArray, NumericArray, StringArray, Table}; +use minarrow::{Array, Field, NumericArray, Table, Vec64, fa_i64, fa_str32}; use tokio::runtime::Runtime; use tokio::task::spawn_blocking; @@ -41,21 +41,14 @@ pub fn runtime() -> Runtime { /// Table of `n` rows whose `id` column starts at `first`. pub fn sample_table(first: i64, n: usize) -> Table { - let ids: Vec = (first..first + n as i64).collect(); - let venues: Vec<&str> = ids + let ids: Vec64 = (first..first + n as i64).collect(); + let venues: Vec64<&str> = ids .iter() .map(|id| if id % 2 == 0 { "XNAS" } else { "XNYS" }) .collect(); Table::new( "trades".to_string(), - vec![ - FieldArray::from_arr("id", Array::from_int64(IntegerArray::from_slice(&ids))), - FieldArray::from_arr( - "venue", - Array::from_string32(StringArray::from_slice(&venues)), - ), - ] - .into(), + vec![fa_i64!("id", @vec64 ids), fa_str32!("venue", @vec64 venues)].into(), ) } @@ -223,7 +216,7 @@ pub mod parquet { pub fn chunked_parquet_round_trips_beneath_an_object_prefix(base: &str) { let dir = format!("{base}/chunked-parquet/run"); - let mut writer = ChunkedParquetWriter::new(&dir, "trades", None).unwrap(); + let mut writer = ChunkedParquetWriter::new(&dir, "trades", None, None).unwrap(); for first in [0, 10, 20] { writer.write_chunk(&sample_table(first, 10)).unwrap(); } @@ -235,7 +228,7 @@ pub mod parquet { pub fn parallel_chunk_writes_reach_the_object_prefix(base: &str) { let dir = format!("{base}/chunked-parquet/parallel"); - let writer = ChunkedParquetWriter::new(&dir, "trades", None).unwrap(); + let writer = ChunkedParquetWriter::new(&dir, "trades", None, None).unwrap(); let tables: Vec
= (0..6).map(|i| sample_table(i * 5, 5)).collect(); let refs: Vec<&Table> = tables.iter().collect(); writer.par_write_all(&refs, Some(3)).unwrap(); @@ -247,7 +240,7 @@ pub mod parquet { pub fn chunked_parquet_loads_as_a_super_table(base: &str) { let dir = format!("{base}/chunked-parquet/super-table"); - let mut writer = ChunkedParquetWriter::new(&dir, "trades", None).unwrap(); + let mut writer = ChunkedParquetWriter::new(&dir, "trades", None, None).unwrap(); for first in [0, 10, 20, 30] { writer.write_chunk(&sample_table(first, 10)).unwrap(); } @@ -273,12 +266,20 @@ pub mod arrow_chunks { use super::*; use lightstream::models::readers::chunked::arrow::ChunkedArrowReader; use lightstream::models::writers::chunked::arrow::ChunkedArrowWriter; + #[cfg(feature = "datetime")] + use lightstream::models::writers::chunked::partition::PartitionType; use lightstream::traits::chunked_table_reader::ChunkedTableReader; use lightstream::traits::chunked_table_writer::ChunkedTableWriter; + /// 2026-02-09T07:00:00Z in epoch nanoseconds. + #[cfg(feature = "datetime")] + const SEVEN_AM: u64 = 1_770_620_400_000_000_000; + #[cfg(feature = "datetime")] + const ONE_HOUR: u64 = 3_600_000_000_000; + pub fn chunked_arrow_round_trips_beneath_an_object_prefix(base: &str) { let dir = format!("{base}/chunked-arrow/run"); - let mut writer = ChunkedArrowWriter::new(&dir, "trades").unwrap(); + let mut writer = ChunkedArrowWriter::new(&dir, "trades", None).unwrap(); for first in [0, 7, 14] { writer.write_chunk(&sample_table(first, 7)).unwrap(); } @@ -287,6 +288,30 @@ pub mod arrow_chunks { let tables: Vec
= reader.map(|chunk| chunk.unwrap()).collect(); assert_eq!(ids(&tables), (0..21).collect::>()); } + + #[cfg(feature = "datetime")] + pub fn auto_partitioned_chunks_round_trip_beneath_an_object_prefix(base: &str) { + let dir = format!("{base}/chunked-arrow/partitioned"); + let mut writer = ChunkedArrowWriter::new(&dir, "trades", Some(PartitionType::Auto)).unwrap(); + // Written out of time order: ids 10..20 at 08:00, then 0..10 at 07:00. + let eight = writer + .write_chunk_for_time(&sample_table(10, 10), SEVEN_AM + ONE_HOUR) + .unwrap(); + writer.write_chunk_for_time(&sample_table(0, 10), SEVEN_AM).unwrap(); + assert_eq!( + eight.to_str().unwrap(), + format!("{dir}/date=2026-02-09/hour=08/trades-0000000000.arrow") + ); + + let all = ChunkedArrowReader::open(&dir, "trades", ()).unwrap(); + let tables: Vec
= all.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..20).collect::>()); + + let hour = format!("{dir}/date=2026-02-09/hour=07"); + let seven = ChunkedArrowReader::open(&hour, "trades", ()).unwrap(); + let tables: Vec
= seven.map(|chunk| chunk.unwrap()).collect(); + assert_eq!(ids(&tables), (0..10).collect::>()); + } } #[cfg(feature = "csv")] @@ -313,7 +338,7 @@ pub mod csv { pub fn chunked_csv_round_trips_beneath_an_object_prefix(base: &str) { let dir = format!("{base}/chunked-csv/run"); - let mut writer = ChunkedCsvWriter::new(&dir, "trades", Default::default()).unwrap(); + let mut writer = ChunkedCsvWriter::new(&dir, "trades", Default::default(), None).unwrap(); for first in [0, 4] { writer.write_chunk(&sample_table(first, 4)).unwrap(); } diff --git a/rust/tests/object_storage_roundtrip.rs b/rust/tests/object_storage_roundtrip.rs index 8004047..1a2005a 100644 --- a/rust/tests/object_storage_roundtrip.rs +++ b/rust/tests/object_storage_roundtrip.rs @@ -120,6 +120,12 @@ fn chunked_arrow_round_trips_beneath_an_object_prefix() { run(checks::arrow_chunks::chunked_arrow_round_trips_beneath_an_object_prefix); } +#[cfg(feature = "datetime")] +#[test] +fn auto_partitioned_chunks_round_trip_beneath_an_object_prefix() { + run(checks::arrow_chunks::auto_partitioned_chunks_round_trip_beneath_an_object_prefix); +} + #[cfg(feature = "csv")] #[test] fn csv_round_trips_through_the_path_entry_points() { diff --git a/rust/tests/object_storage_s3.rs b/rust/tests/object_storage_s3.rs index 38df89a..abb28c5 100644 --- a/rust/tests/object_storage_s3.rs +++ b/rust/tests/object_storage_s3.rs @@ -242,6 +242,16 @@ fn chunked_arrow_round_trips_beneath_an_object_prefix() { ); } +#[cfg(feature = "datetime")] +#[test] +#[ignore = "requires LIGHTSTREAM_S3_URL and AWS credentials"] +fn auto_partitioned_chunks_round_trip_beneath_an_object_prefix() { + run( + "partitioned-arrow", + checks::arrow_chunks::auto_partitioned_chunks_round_trip_beneath_an_object_prefix, + ); +} + #[cfg(feature = "csv")] #[test] #[ignore = "requires LIGHTSTREAM_S3_URL and AWS credentials"]