diff --git a/Cargo.lock b/Cargo.lock index 2c573148d20..3fdaaa5d45b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10657,12 +10657,16 @@ dependencies = [ "codspeed-divan-compat", "insta", "itertools 0.14.0", + "libc", + "onpair", "pco", "rand 0.10.2", "rstest", "test-with", "tpchgen", "tpchgen-arrow", + "tracing", + "tracing-subscriber", "vortex-alp", "vortex-array", "vortex-arrow", @@ -11052,6 +11056,7 @@ dependencies = [ "futures", "itertools 0.14.0", "kanal", + "libc", "moka", "object_store", "parking_lot", diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index cef260fd871..5b8e38c6fc2 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -60,6 +60,8 @@ exists: a grep proves presence, a test proves behaviour. | 12 | Restore lint checks on forks | `bb80c537b` | Fork CI not running lints | — (CI config) | No | | 13 | `set_available_parallelism` | this change | Scan and writer fan-out sized from the machine's core count rather than what the host process is entitled to (spiceai/spiceai#12328) | `cargo test -p vortex-utils --test parallelism_declared --test parallelism_declared_too_late` | Proposed — additive, detection unchanged | | 14 | Arrow `Map` alias | `1a6dc54f1` (`lukim/map`) | `Map` columns unwritable: `Array encoding not implemented for Arrow data type Map(...)` on every flush (spiceai/spiceai#13524) | `cargo test -p vortex-arrow --lib executor::map` | Yes — 0.86.1 native `Map` (`#9111`) | +| 17 | Byte-capped data blocks for string/binary columns, with encodings chosen once per row block | this change | A string or binary column is written in blocks of 8,192 rows whatever its row width, because `effective_block_len` sizes blocks only for fixed-width types; a ~3 KB/row JSON column gets ~12.5 MiB blocks, and a one-row lookup reads and copies a whole block (spiceai/spiceai#15323). `WriteStrategyBuilder::with_data_block_max_bytes` splits such a block into power-of-two row counts under the cap. Searching encodings on each small block costs several full compressions per block (a block of at most 1,024 rows is its own sample), so `CompressingStrategy::with_group_rows` and `CascadingCompressor::plan_for`/`compress_with_plan` choose encodings once per row block, on the sample the compressor would draw from the unsplit block, and apply them to its split blocks (a choice made on the first split block alone fits only it: incompressible rows ahead of JSON picked FSST for the JSON, +15% size and 2.4x `LIKE` CPU) | `cargo test -p vortex-layout --lib repartition::tests::a_block_size_maximum_splits_wide_strings_into_power_of_two_blocks` **and** `cargo test -p vortex-btrblocks --lib canonical_compressor::tests::a_replayed_plan_roundtrips_other_data_of_the_same_dtype` **and** `cargo test -p vortex-layout --lib compressed::tests::a_group_is_encoded_as_its_unsplit_data_would_be` **and** `cargo test -p vortex-file --test test_write_table test_data_block_max_bytes_roundtrip` | Proposed — additive; off unless the cap is set | +| 18 | Chunked reads skip chunks a known mask excludes | this change | A read of a row range given an already-resolved mask (`MaskFuture::ready`, as a point read passes) requested every chunk in the range, even those the mask selects no row of, because the flat reader fetches its segment before it looks at the mask. Once string blocks are split under a byte cap (row 17), a point read of one row of an 8,192-row range fetched all of its split blocks, so a lookup got slower rather than faster (spiceai/spiceai#15323). `MaskFuture::known` exposes a mask known up front, and `ChunkedReader` skips chunks it excludes | `cargo test -p vortex-layout --lib chunked::reader::test::a_known_mask_reads_only_the_chunks_it_selects` | Proposed — upstream defect, not Spice behaviour | **Confidence:** rows 1–4, 13 and 14 are verified against the branches. Rows 5–12 are seeded from an audit of non-merge commits on `spiceai-54` authored by Spice engineers; their descriptions diff --git a/vortex-array/src/mask_future.rs b/vortex-array/src/mask_future.rs index a46107623c1..8fa7a41d21a 100644 --- a/vortex-array/src/mask_future.rs +++ b/vortex-array/src/mask_future.rs @@ -20,6 +20,8 @@ use vortex_mask::Mask; pub struct MaskFuture { inner: Shared>>, len: usize, + /// The mask, when it was known before the future was created; see [`Self::known`]. + known: Option, } impl MaskFuture { @@ -40,9 +42,19 @@ impl MaskFuture { .boxed() .shared(), len, + known: None, } } + /// The mask, if it is known without awaiting: this future was created with [`Self::ready`] + /// or [`Self::new_true`], or sliced from one that was (and not given an [`Self::inspect`] + /// callback, which must run). + /// + /// A reader can use it to skip work for rows the mask excludes before requesting any data. + pub fn known(&self) -> Option<&Mask> { + self.known.as_ref() + } + /// Returns the length of the mask. pub fn len(&self) -> usize { self.len @@ -55,7 +67,10 @@ impl MaskFuture { /// Create a MaskFuture from a ready mask. pub fn ready(mask: Mask) -> Self { - Self::new(mask.len(), async move { Ok(mask) }) + let known = mask.clone(); + let mut future = Self::new(mask.len(), async move { Ok(mask) }); + future.known = Some(known); + future } /// Create a MaskFuture that resolves to a mask with all values set to true. @@ -71,6 +86,10 @@ impl MaskFuture { return self.clone(); } + if let Some(known) = &self.known { + return Self::ready(known.slice(range)); + } + let inner = self.inner.clone(); Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }) } @@ -81,9 +100,12 @@ impl MaskFuture { ) -> Self { let len = self.len; + // Not `known`: a slice of a known mask resolves without awaiting this future, and so + // without running `f`. Self { inner: self.inner.inspect(f).boxed().shared(), len, + known: None, } } } @@ -124,3 +146,47 @@ mod tests { }) } } + +#[cfg(test)] +mod known_tests { + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + + use vortex_error::VortexResult; + use vortex_mask::Mask; + + use super::MaskFuture; + + #[test] + fn a_ready_mask_and_its_slices_are_known() { + let ready = MaskFuture::ready(Mask::from_indices(10, vec![2, 7])); + assert_eq!(ready.known(), Some(&Mask::from_indices(10, vec![2, 7]))); + assert_eq!( + ready.slice(5..10).known(), + Some(&Mask::from_indices(5, vec![2])) + ); + assert_eq!(MaskFuture::new_true(4).known(), Some(&Mask::new_true(4))); + } + + #[test] + fn a_pending_mask_is_not_known_until_awaited() { + let pending = MaskFuture::new(4, async { Ok(Mask::new_true(4)) }); + assert!(pending.known().is_none()); + assert!(pending.slice(1..3).known().is_none()); + } + + /// A slice of an inspected future still awaits it, so the callback runs. + #[test] + fn a_slice_of_an_inspected_future_runs_the_callback() -> VortexResult<()> { + futures::executor::block_on(async { + let ran = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&ran); + let fut = MaskFuture::ready(Mask::new_true(4)) + .inspect(move |_| flag.store(true, Ordering::SeqCst)); + assert_eq!(fut.slice(1..3).await?, Mask::new_true(2)); + assert!(ran.load(Ordering::SeqCst)); + Ok(()) + }) + } +} diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 0aa7e9669c3..7337f2bbcbc 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -40,6 +40,10 @@ vortex-zstd = { workspace = true, optional = true } [dev-dependencies] arrow-array = { workspace = true } +libc = "0.2" +onpair = { workspace = true } +tracing = { workspace = true, features = ["std"] } +tracing-subscriber = { workspace = true, features = ["registry"] } divan = { workspace = true } insta = { workspace = true } rstest = { workspace = true } @@ -64,6 +68,11 @@ name = "compress" harness = false test = false +[[bench]] +name = "small_chunk_compress" +harness = false +test = false + [[bench]] name = "compress_listview" harness = false diff --git a/vortex-btrblocks/benches/small_chunk_compress.rs b/vortex-btrblocks/benches/small_chunk_compress.rs new file mode 100644 index 00000000000..9ea2b713489 --- /dev/null +++ b/vortex-btrblocks/benches/small_chunk_compress.rs @@ -0,0 +1,309 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Where compression time goes when a wide string column is written in chunks smaller than the +//! default 8,192 rows. Compresses the same ~3 KB JSON column as 8,192-row chunks and as +//! sub-chunks, with the full scheme search and with only the scheme that wins, and times +//! OnPair's dictionary training and parsing on their own. +//! +//! `cargo bench -p vortex-btrblocks --bench small_chunk_compress` +//! Env: `SCC_ROWS` (default 98,304), `SCC_REPS` (default 5). +//! +//! It measures CPU time with Unix APIs, so on other platforms it only says so. + +#[cfg(unix)] +mod unix { + #![expect(clippy::unwrap_used)] + #![expect(clippy::cast_precision_loss)] + + use std::sync::LazyLock; + + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinViewArray; + use vortex_btrblocks::BtrBlocksCompressor; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_btrblocks::SchemeExt; + use vortex_btrblocks::schemes::string; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + + /// Deterministic xorshift so every run compresses identical data. + struct Rng(u64); + impl Rng { + fn next(&mut self) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x + } + } + + /// A ~3 KB JSON document per row, like a cached HTTP response body. + fn json_row(rng: &mut Rng, row: usize) -> String { + let mut out = String::with_capacity(3200); + out.push_str(&format!("{{\"id\":{row},\"items\":[")); + let mut first = true; + while out.len() < 3000 { + if !first { + out.push(','); + } + first = false; + let value = rng.next(); + out.push_str(&format!( + "{{\"sku\":\"SKU-{:08x}\",\"price\":{}.{:02},\"qty\":{},\"name\":\"item {:x}\",\"tags\":[\"t{}\",\"t{}\"]}}", + value >> 32, + (value >> 32) % 1000, + (value >> 40) % 100, + (value >> 48) % 50, + value >> 20, + (value >> 8) % 17, + (value >> 12) % 23 + )); + } + out.push_str("]}"); + out + } + + fn compress_chunks(compressor: &BtrBlocksCompressor, array: &ArrayRef, rows: usize) -> u64 { + let mut ctx = SESSION.create_execution_ctx(); + let mut bytes = 0; + let mut start = 0; + while start < array.len() { + let end = (start + rows).min(array.len()); + let chunk = array.slice(start..end).unwrap(); + bytes += compressor.compress(&chunk, &mut ctx).unwrap().nbytes(); + start = end; + } + bytes + } + + /// CPU time of the calling thread, in milliseconds. Everything timed here runs on one + /// thread, so this is far less sensitive to other load on the machine than wall time. + fn thread_cpu_ms() -> f64 { + let mut ts = libc::timespec { + tv_sec: 0, + tv_nsec: 0, + }; + // SAFETY: `ts` is a valid out-pointer. + unsafe { libc::clock_gettime(libc::CLOCK_THREAD_CPUTIME_ID, &raw mut ts) }; + ts.tv_sec as f64 * 1e3 + ts.tv_nsec as f64 / 1e6 + } + + /// Exclusive (self) thread-CPU time per span, keyed by span name and its scheme/kind field. + mod self_time { + use std::cell::RefCell; + use std::collections::BTreeMap; + use std::fmt::Debug; + + use tracing::Subscriber; + use tracing::field::Field; + use tracing::field::Visit; + use tracing::span::Attributes; + use tracing::span::Id; + use tracing_subscriber::Layer; + use tracing_subscriber::layer::Context; + use tracing_subscriber::registry::LookupSpan; + + // Everything this bench times runs on one thread. + thread_local! { + static TOTALS: RefCell> = const { RefCell::new(BTreeMap::new()) }; + static STACK: RefCell> = const { RefCell::new(Vec::new()) }; + } + + struct Key(String); + struct KeyVisitor<'a>(&'a mut String); + impl Visit for KeyVisitor<'_> { + fn record_debug(&mut self, field: &Field, value: &dyn Debug) { + if matches!( + field.name(), + "scheme_candidate" | "scheme_chosen" | "stats_kind" + ) { + self.0.push_str(&format!("[{value:?}]")); + } + } + } + + pub struct SelfTime; + impl LookupSpan<'a>> Layer for SelfTime { + fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) { + let mut key = attrs.metadata().name().to_string(); + attrs.record(&mut KeyVisitor(&mut key)); + if let Some(span) = ctx.span(id) { + span.extensions_mut().insert(Key(key)); + } + } + fn on_enter(&self, _id: &Id, _ctx: Context<'_, S>) { + STACK.with(|s| s.borrow_mut().push((super::thread_cpu_ms(), 0.0))); + } + fn on_exit(&self, id: &Id, ctx: Context<'_, S>) { + let Some((start, children)) = STACK.with(|s| s.borrow_mut().pop()) else { + return; + }; + let total = super::thread_cpu_ms() - start; + STACK.with(|s| { + if let Some(parent) = s.borrow_mut().last_mut() { + parent.1 += total; + } + }); + if let Some(span) = ctx.span(id) + && let Some(Key(key)) = span.extensions().get::() + { + TOTALS.with(|totals| { + let mut totals = totals.borrow_mut(); + let entry = totals.entry(key.clone()).or_insert((0.0, 0)); + entry.0 += total - children; + entry.1 += 1; + }); + } + } + } + + pub fn take() -> BTreeMap { + TOTALS.with(|totals| std::mem::take(&mut *totals.borrow_mut())) + } + } + + fn compress_planned( + compressor: &BtrBlocksCompressor, + array: &ArrayRef, + rows: usize, + chunk_rows: usize, + from_sample: bool, + ) -> u64 { + let mut ctx = SESSION.create_execution_ctx(); + let mut bytes = 0; + for block in (0..rows).step_by(8192) { + let block_end = (block + 8192).min(rows); + // The writer's planning: schemes chosen on the sample `compress` would draw from + // the whole block. + let mut plan = from_sample.then(|| { + let whole = array.slice(block..block_end).unwrap(); + compressor.plan_for(&whole, &mut ctx).unwrap() + }); + for start in (block..block_end).step_by(chunk_rows) { + let end = (start + chunk_rows).min(block_end); + let chunk = array.slice(start..end).unwrap(); + let compressed = match &plan { + None => { + let (c, p) = compressor.compress_recording(&chunk, &mut ctx).unwrap(); + plan = Some(p); + c + } + Some(p) => compressor.compress_with_plan(&chunk, p, &mut ctx).unwrap(), + }; + bytes += compressed.nbytes(); + } + } + bytes + } + + pub fn main() { + if std::env::var_os("SCC_SPANS").is_some() { + use tracing_subscriber::layer::SubscriberExt; + let subscriber = tracing_subscriber::registry().with(self_time::SelfTime); + tracing::subscriber::set_global_default(subscriber).unwrap(); + } + let rows: usize = std::env::var("SCC_ROWS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(98_304); + let reps: usize = std::env::var("SCC_REPS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(5); + let mut rng = Rng(0x9e37_79b9_7f4a_7c15); + let strings: Vec = (0..rows).map(|i| json_row(&mut rng, i)).collect(); + let raw_bytes: usize = strings.iter().map(String::len).sum(); + let array = + &VarBinViewArray::from_iter_str(strings.iter().map(String::as_str)).into_array(); + + let full = &BtrBlocksCompressor::default(); + let onpair_only = &BtrBlocksCompressorBuilder::default() + .exclude_schemes([string::FSSTScheme.id(), string::StringDictScheme.id()]) + .build(); + + type Arm<'a> = (String, Box u64 + 'a>); + let mut arms: Vec> = vec![( + "default: full search, 8192 rows".to_string(), + Box::new(move || compress_chunks(full, array, 8192)), + )]; + for chunk_rows in [1024_usize, 339] { + arms.push(( + format!("full search, {chunk_rows} rows"), + Box::new(move || compress_chunks(full, array, chunk_rows)), + )); + arms.push(( + format!("onpair only, {chunk_rows} rows"), + Box::new(move || compress_chunks(onpair_only, array, chunk_rows)), + )); + arms.push(( + format!("plan P1 first sub-chunk, {chunk_rows} rows"), + Box::new(move || compress_planned(full, array, rows, chunk_rows, false)), + )); + arms.push(( + format!("plan_for on the whole block, {chunk_rows} rows"), + Box::new(move || compress_planned(full, array, rows, chunk_rows, true)), + )); + } + + // Interleaved: every rep runs each arm once, so load on the machine falls on all arms alike. + let mut times: Vec> = vec![Vec::new(); arms.len()]; + let mut outs = vec![0_u64; arms.len()]; + let mut spans: Vec> = + vec![Default::default(); arms.len()]; + for _ in 0..reps { + for (i, (_, arm)) in arms.iter_mut().enumerate() { + self_time::take(); + let t = thread_cpu_ms(); + outs[i] = arm(); + times[i].push(thread_cpu_ms() - t); + for (k, (ms, n)) in self_time::take() { + let e = spans[i].entry(k).or_insert((0.0, 0)); + e.0 += ms; + e.1 += n; + } + } + } + println!( + "rows={rows} raw={:.1} MiB reps={reps} interleaved (thread CPU ms)", + raw_bytes as f64 / f64::from(1 << 20) + ); + println!( + "{:<36} {:>9} {:>9} {:>9} {:>10}", + "arm", "median", "min", "max", "out_MiB" + ); + for (i, (label, _)) in arms.iter().enumerate() { + let t = &mut times[i]; + t.sort_by(f64::total_cmp); + println!( + "{:<36} {:>9.1} {:>9.1} {:>9.1} {:>10.2}", + label, + t[t.len() / 2], + t[0], + t[t.len() - 1], + outs[i] as f64 / f64::from(1 << 20) + ); + let mut s: Vec<_> = spans[i].iter().collect(); + s.sort_by(|a, b| b.1.0.total_cmp(&a.1.0)); + for (key, (ms, n)) in s.iter().take(8) { + println!( + " {:>9.1} ms/rep {:>8} calls/rep {key}", + ms / reps as f64, + *n / reps as u64 + ); + } + } + } +} + +fn main() { + #[cfg(unix)] + unix::main(); + #[cfg(not(unix))] + eprintln!("this benchmark measures CPU time with Unix APIs; it does not run on this platform"); +} diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index ed528145d71..e95431a904a 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -311,4 +311,117 @@ mod tests { assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } + + /// Strings shaped like JSON documents, with a seed so two calls draw different data. + fn json_strings(rows: usize, seed: u64) -> VarBinViewArray { + let mut state = seed; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + VarBinViewArray::from_iter_str((0..rows).map(|row| { + let mut s = format!("{{\"id\":{row},\"items\":["); + for _ in 0..(next() % 8 + 1) { + let v = next(); + s.push_str(&format!( + "{{\"sku\":\"SKU-{:08x}\",\"qty\":{}}},", + v >> 32, + (v >> 40) % 50 + )); + } + s.push_str("]}"); + s + })) + } + + #[test] + fn replaying_a_plan_reproduces_the_search_on_the_same_array() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = json_strings(4096, 7).into_array(); + let compressor = BtrBlocksCompressor::default(); + + let searched = compressor.compress(&input, &mut ctx)?; + let (recorded, plan) = compressor.compress_recording(&input, &mut ctx)?; + let replayed = compressor.compress_with_plan(&input, &plan, &mut ctx)?; + + assert!(!plan.is_empty(), "recording decided nothing"); + let tree = searched.display_tree_encodings_only().to_string(); + assert_eq!(recorded.display_tree_encodings_only().to_string(), tree); + assert_eq!(replayed.display_tree_encodings_only().to_string(), tree); + assert_eq!(replayed.nbytes(), searched.nbytes()); + assert_arrays_eq!(replayed, input, &mut ctx); + Ok(()) + } + + #[test] + fn a_replayed_plan_roundtrips_other_data_of_the_same_dtype() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let compressor = BtrBlocksCompressor::default(); + let (recorded, plan) = + compressor.compress_recording(&json_strings(1024, 7).into_array(), &mut ctx)?; + + // Different rows, and short repeated strings the recorded data never had. + let other = json_strings(1024, 11).into_array(); + let replayed = compressor.compress_with_plan(&other, &plan, &mut ctx)?; + assert_arrays_eq!(replayed, other, &mut ctx); + + let repeated = VarBinViewArray::from_iter_str((0..1024).map(|i| ["a", "bb", "ccc"][i % 3])) + .into_array(); + let replayed = compressor.compress_with_plan(&repeated, &plan, &mut ctx)?; + assert_arrays_eq!(replayed, repeated, &mut ctx); + // The replay used the recorded scheme rather than searching again: a search on these + // strings picks a different one. + let searched = compressor.compress(&repeated, &mut ctx)?; + assert_ne!(searched.encoding_id(), recorded.encoding_id()); + assert_eq!(replayed.encoding_id(), recorded.encoding_id()); + Ok(()) + } + + #[test] + fn an_empty_plan_falls_back_to_the_search() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = json_strings(2048, 3).into_array(); + let compressor = BtrBlocksCompressor::default(); + + let searched = compressor.compress(&input, &mut ctx)?; + let replayed = + compressor.compress_with_plan(&input, &crate::CompressionPlan::default(), &mut ctx)?; + + assert_eq!( + replayed.display_tree_encodings_only().to_string(), + searched.display_tree_encodings_only().to_string() + ); + assert_eq!(replayed.nbytes(), searched.nbytes()); + Ok(()) + } + + /// A plan is replayed on other data of the same dtype; a scheme that cannot encode that data + /// must not fail the write. + #[test] + fn a_replayed_plan_survives_data_its_schemes_cannot_encode() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let compressor = BtrBlocksCompressor::default(); + let sequential = vortex_array::arrays::PrimitiveArray::from_iter(0..4096_i64).into_array(); + let (recorded, plan) = compressor.compress_recording(&sequential, &mut ctx)?; + // The premise: the plan holds a scheme that random integers cannot use. + assert_eq!(recorded.encoding_id().to_string(), "vortex.sequence"); + + let mut state = 0x9e37_79b9_u64; + let random = vortex_array::arrays::PrimitiveArray::from_iter((0..4096).map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state % 1_000_000) as i64 + })) + .into_array(); + let replayed = compressor.compress_with_plan(&random, &plan, &mut ctx)?; + assert_eq!( + replayed.encoding_id(), + compressor.compress(&random, &mut ctx)?.encoding_id() + ); + assert_arrays_eq!(replayed, random, &mut ctx); + Ok(()) + } } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 1ca05c86b4e..cbc5e4947ab 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -81,6 +81,7 @@ pub use builder::BtrBlocksCompressorBuilder; pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; pub use vortex_compressor::CascadingCompressor; +pub use vortex_compressor::CompressionPlan; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..487cdfb9c0c 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -32,6 +32,9 @@ use vortex_error::VortexResult; use super::CascadingCompressor; use super::constant; +use super::plan::CompressionPlan; +use super::plan::Decision; +use super::plan::PlanState; use crate::scheme::CompressorContext; use crate::scheme::Scheme; use crate::scheme::SchemeExt; @@ -66,6 +69,70 @@ impl CascadingCompressor { Ok(compressed) } + /// Compresses `array` like [`Self::compress`], and returns the schemes it chose as a plan + /// that [`Self::compress_with_plan`] can replay on other arrays of the same dtype. + /// + /// # Errors + /// + /// Returns an error if canonicalizing `array` or any scheme's compression fails. + pub fn compress_recording( + &self, + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult<(ArrayRef, CompressionPlan)> { + let state = PlanState::recording(); + let canonical = array.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + let ctx = CompressorContext::new().with_plan(std::sync::Arc::clone(&state)); + let compressed = self.compress_canonical(compact, ctx, exec_ctx)?; + Ok((compressed, state.into_plan())) + } + + /// Chooses the schemes for `array` without compressing all of it, returning them as a plan + /// for [`Self::compress_with_plan`] to apply to `array`'s parts. + /// + /// The schemes are chosen on the same stratified sample [`Self::compress`] would estimate + /// them on for the whole of `array`, so applying the plan to slices of `array` encodes them + /// as `array` itself would have been, even when its data differs from one part to another. + /// + /// # Errors + /// + /// Returns an error if canonicalizing the sample or any scheme's compression fails. + pub fn plan_for( + &self, + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let sample = super::sample::sample( + array, + super::sample::SAMPLE_SIZE, + super::sample::sample_count_approx_one_percent(array.len()), + ); + Ok(self.compress_recording(&sample, exec_ctx)?.1) + } + + /// Compresses `array` with the schemes `plan` recorded, without estimating any scheme. + /// + /// Where the plan decided nothing, which happens when `array`'s data takes a path through + /// the cascade the recorded array did not, the full search runs as in [`Self::compress`]. + /// A replayed scheme is still only kept when its output is smaller than its input. + /// + /// # Errors + /// + /// Returns an error if canonicalizing `array` or any scheme's compression fails. + pub fn compress_with_plan( + &self, + array: &ArrayRef, + plan: &CompressionPlan, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let state = PlanState::replaying(plan.clone()); + let canonical = array.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + let ctx = CompressorContext::new().with_plan(state); + self.compress_canonical(compact, ctx, exec_ctx) + } + /// Compresses a child array produced by a cascading scheme. /// /// If the cascade budget is exhausted, the canonical array is returned as-is. Otherwise, the @@ -295,9 +362,34 @@ impl CascadingCompressor { return Ok(data.into_array()); } - let Some((winner, winner_estimate)) = - self.choose_best_scheme(&eligible_schemes, &data, compress_ctx.clone(), exec_ctx)? - else { + let replayed = compress_ctx + .plan() + .and_then(|plan| plan.replayed(compress_ctx.cascade_history())); + let chosen = match replayed { + Some(Decision::Canonical) => return Ok(data.into_array()), + Some(Decision::Scheme(id)) => eligible_schemes + .iter() + .copied() + .find(|scheme| scheme.id() == id) + .map(|scheme| (scheme, super::select::WinnerEstimate::AlwaysUse)), + None => None, + }; + let from_plan = chosen.is_some(); + let chosen = match chosen { + Some(chosen) => Some(chosen), + None => { + self.choose_best_scheme(&eligible_schemes, &data, compress_ctx.clone(), exec_ctx)? + } + }; + if let Some(plan) = compress_ctx.plan() { + plan.record( + compress_ctx.cascade_history(), + chosen.map_or(Decision::Canonical, |(scheme, _)| { + Decision::Scheme(scheme.id()) + }), + ); + } + let Some((winner, winner_estimate)) = chosen else { return Ok(data.into_array()); }; @@ -305,14 +397,47 @@ impl CascadingCompressor { // scheme name and cascade history before propagating. let error_ctx = trace::enabled_error_context(&compress_ctx); let _winner_span = trace::winner_compress_span(winner.id(), before_nbytes).entered(); - let compressed = winner - .compress(self, &data, compress_ctx, exec_ctx) - .inspect_err(|err| { - // NB: this is the only way we can tell which scheme panicked / bailed on their - // data, especially for third-party schemes where the error site may not carry any - // compressor context. - trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err); - })?; + let (winner_estimate, compressed) = + match winner.compress(self, &data, compress_ctx.clone(), exec_ctx) { + Ok(compressed) => (winner_estimate, compressed), + // A replayed scheme was chosen on other data, and some schemes accept a dtype yet + // reject data they cannot encode (a sequence needs evenly spaced values). Choose + // as if there were no plan; an error from that choice is a real failure. + Err(_) if from_plan => { + let Some((winner, winner_estimate)) = self.choose_best_scheme( + &eligible_schemes, + &data, + compress_ctx.clone(), + exec_ctx, + )? + else { + return Ok(data.into_array()); + }; + let compressed = winner + .compress(self, &data, compress_ctx, exec_ctx) + .inspect_err(|err| { + trace::scheme_compress_failed( + winner.id(), + before_nbytes, + error_ctx.as_ref(), + err, + ); + })?; + (winner_estimate, compressed) + } + Err(err) => { + // NB: this is the only way we can tell which scheme panicked / bailed on their + // data, especially for third-party schemes where the error site may not carry + // any compressor context. + trace::scheme_compress_failed( + winner.id(), + before_nbytes, + error_ctx.as_ref(), + &err, + ); + return Err(err); + } + }; let after_nbytes = compressed.nbytes(); let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..5057ed7e7ac 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -5,6 +5,7 @@ mod cascade; mod constant; +pub(crate) mod plan; mod sample; mod select; mod structural; diff --git a/vortex-compressor/src/compressor/plan.rs b/vortex-compressor/src/compressor/plan.rs new file mode 100644 index 00000000000..b266bbdb010 --- /dev/null +++ b/vortex-compressor/src/compressor/plan.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Recording the schemes a compression chose, and replaying them on other arrays. +//! +//! Choosing a scheme compresses a sample of the input once per candidate, and every child a +//! winning scheme produces is chosen for the same way. For an input of at most +//! `SAMPLE_SIZE * SAMPLE_COUNT` rows the sample is the whole input, so the search costs one full +//! compression per candidate. When many small arrays are drawn from the same data, such as the +//! row blocks of one column, the choices made for one of them hold for the rest, and replaying +//! them skips the search. + +use std::sync::Arc; + +use parking_lot::Mutex; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::scheme::SchemeId; + +/// What [`CascadingCompressor`](crate::CascadingCompressor) decided at one point in a cascade. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Decision { + /// Compress with this scheme. + Scheme(SchemeId), + /// No scheme was chosen: keep the canonical array. + Canonical, +} + +/// Whether a [`PlanState`] collects decisions or hands them out. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PlanMode { + /// Decisions are stored as the compression makes them. + Record, + /// Stored decisions replace the search. + Replay, +} + +/// The decisions of one compression, keyed by cascade path (the `(scheme, child index)` chain +/// from the root to the array being compressed). +/// +/// A path that was never decided falls back to the full search when replayed, so a plan +/// recorded on one array can be replayed on any other array of the same dtype. +#[derive(Debug, Clone, Default)] +pub struct CompressionPlan { + /// The decision at each cascade path. + decisions: HashMap, Decision>, +} + +impl CompressionPlan { + /// The number of cascade paths this plan decides. + pub fn len(&self) -> usize { + self.decisions.len() + } + + /// Whether this plan decides nothing. + pub fn is_empty(&self) -> bool { + self.decisions.is_empty() + } +} + +/// A plan being recorded or replayed, shared by every context in one compression. +#[derive(Debug)] +pub(crate) struct PlanState { + /// Whether this state records or replays. + mode: PlanMode, + /// The decisions recorded so far, or the ones to replay. + plan: Mutex, +} + +impl PlanState { + /// A state that records the decisions of the compression it is passed to. + pub(crate) fn recording() -> Arc { + Arc::new(Self { + mode: PlanMode::Record, + plan: Mutex::new(CompressionPlan::default()), + }) + } + + /// A state that replays `plan`. + pub(crate) fn replaying(plan: CompressionPlan) -> Arc { + Arc::new(Self { + mode: PlanMode::Replay, + plan: Mutex::new(plan), + }) + } + + /// The recorded decision for `path`, when replaying. + pub(crate) fn replayed(&self, path: &[(SchemeId, usize)]) -> Option { + if self.mode != PlanMode::Replay { + return None; + } + self.plan.lock().decisions.get(path).copied() + } + + /// Records `decision` for `path`, when recording. + pub(crate) fn record(&self, path: &[(SchemeId, usize)], decision: Decision) { + if self.mode == PlanMode::Record { + self.plan.lock().decisions.insert(path.to_vec(), decision); + } + } + + /// The decisions this state holds. + pub(crate) fn into_plan(self: Arc) -> CompressionPlan { + match Arc::try_unwrap(self) { + Ok(state) => state.plan.into_inner(), + Err(shared) => shared.plan.lock().clone(), + } + } +} diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 55bb9b188f6..2f170ef7254 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -68,5 +68,6 @@ pub mod stats; mod compressor; pub use compressor::CascadingCompressor; +pub use compressor::plan::CompressionPlan; mod trace; diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..ba932b13ca5 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,10 +4,12 @@ //! Compression context for recursive compression. use std::fmt; +use std::sync::Arc; use vortex_error::VortexExpect; use crate::compressor::ROOT_SCHEME_ID; +use crate::compressor::plan::PlanState; use crate::scheme::SchemeId; use crate::stats::GenerateStatsOptions; @@ -38,6 +40,9 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, + + /// The plan this compression records its decisions into or replays them from. + plan: Option>, } impl CompressorContext { @@ -50,6 +55,7 @@ impl CompressorContext { allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), + plan: None, } } } @@ -107,6 +113,21 @@ impl CompressorContext { self } + /// Returns a context that records its decisions into, or replays them from, `plan`. + pub(crate) fn with_plan(mut self, plan: Arc) -> Self { + self.plan = Some(plan); + self + } + + /// The plan to record into or replay from, outside of sampling: a decision made while + /// estimating a scheme on a sample is not the decision for the array itself. + pub(crate) fn plan(&self) -> Option<&PlanState> { + if self.is_sample { + return None; + } + self.plan.as_deref() + } + /// Returns a context marked as sample compression. pub(crate) fn with_sampling(mut self) -> Self { self.is_sample = true; diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 78ff4d1d5b6..596011541a2 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -61,6 +61,7 @@ vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } [dev-dependencies] +libc = "0.2" allocator-api2 = { workspace = true } divan = { workspace = true } rstest = { workspace = true } @@ -89,3 +90,13 @@ tokio = [ wasm-bindgen = ["vortex-layout/wasm-bindgen"] zstd = ["dep:vortex-zstd", "vortex-btrblocks/zstd", "vortex-btrblocks/pco"] tensor = ["dep:vortex-tensor"] + +[[bench]] +name = "scan_mask_paths" +harness = false +test = false + +[[bench]] +name = "wide_string_cases" +harness = false +test = false diff --git a/vortex-file/benches/scan_mask_paths.rs b/vortex-file/benches/scan_mask_paths.rs new file mode 100644 index 00000000000..605c35206b3 --- /dev/null +++ b/vortex-file/benches/scan_mask_paths.rs @@ -0,0 +1,271 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Read-path cost of scans whose row masks are known up front (a row selection) or computed (a +//! filter), over a multi-chunk file of narrow columns. Run against two builds to compare them. +//! +//! `cargo bench -p vortex-file --bench scan_mask_paths`; env `SMP_ROWS` (default 4,194,304), +//! `SMP_REPS` (default 9). +//! +//! It measures CPU time with Unix APIs, so on other platforms it only says so. + +#[cfg(unix)] +#[path = "../tests/common/mod.rs"] +mod common; + +#[cfg(unix)] +mod unix { + #![expect(clippy::unwrap_used)] + #![expect(clippy::cast_precision_loss)] + #![expect(clippy::cast_possible_truncation)] + #![expect(clippy::cast_possible_wrap)] + #![expect(clippy::many_single_char_names)] + + use std::sync::LazyLock; + use std::time::Instant; + + use futures::StreamExt; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::FieldNames; + use vortex_array::expr::Expression; + use vortex_array::expr::and; + use vortex_array::expr::get_item; + use vortex_array::expr::gt; + use vortex_array::expr::lit; + use vortex_array::expr::lt; + use vortex_array::expr::root; + use vortex_array::expr::select; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBuffer; + use vortex_file::OpenOptionsSessionExt; + use vortex_file::VortexFile; + use vortex_file::WriteOptionsSessionExt; + use vortex_io::session::RuntimeSession; + use vortex_layout::session::LayoutSession; + use vortex_scan::selection::Selection; + use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + super::common::enable_all_registered_array_encodings(&session); + session + }); + + fn cpu_ms() -> f64 { + // SAFETY: getrusage fills the zeroed struct. + let mut ru: libc::rusage = unsafe { std::mem::zeroed() }; + // SAFETY: valid pointer to a rusage. + unsafe { libc::getrusage(libc::RUSAGE_SELF, &raw mut ru) }; + let tv = |t: libc::timeval| t.tv_sec as f64 * 1e3 + t.tv_usec as f64 / 1e3; + tv(ru.ru_utime) + tv(ru.ru_stime) + } + + struct Rng(u64); + impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + } + + fn table(rows: usize) -> ArrayRef { + let mut rng = Rng(0x9e37_79b9_7f4a_7c15); + let k = PrimitiveArray::from_iter(0..rows as i64).into_array(); + let a = + PrimitiveArray::from_iter((0..rows).map(|_| (rng.next() % 1_000_000) as i64 - 500_000)) + .into_array(); + let b = PrimitiveArray::from_iter((0..rows).map(|_| (rng.next() % 100_000) as f64 / 100.0)) + .into_array(); + let d = PrimitiveArray::from_iter((0..rows).map(|i| (i / 3000) as i32)).into_array(); + let s = VarBinViewArray::from_iter_str( + (0..rows).map(|i| format!("dist-info-{:014}", i % 97_003)), + ) + .into_array(); + StructArray::new( + FieldNames::from(["k", "a", "b", "d", "s"]), + vec![k, a, b, d, s], + rows, + Validity::NonNullable, + ) + .into_array() + } + + async fn scan( + file: &VortexFile, + projection: Expression, + filter: Option, + selection: Option, + ) -> (usize, u64) { + let dtype = file.dtype().clone(); + let mut builder = file + .scan() + .unwrap() + .with_projection(projection.bind(&dtype).unwrap()); + if let Some(f) = filter { + builder = builder.with_filter(f.bind(&dtype).unwrap()); + } + if let Some(s) = selection { + builder = builder.with_selection(s); + } + let mut stream = builder.into_array_stream().unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let (mut rows, mut bytes) = (0, 0); + while let Some(chunk) = stream.next().await { + let chunk = chunk + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + rows += chunk.len(); + bytes += chunk.nbytes(); + } + (rows, bytes) + } + + pub fn main() { + let rows: usize = std::env::var("SMP_ROWS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4_194_304); + let reps: usize = std::env::var("SMP_REPS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(9); + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + let file = rt.block_on(async { + let mut bytes = Vec::new(); + SESSION + .write_options() + .write(&mut bytes, table(rows).to_array_stream()) + .await + .unwrap(); + let digest = bytes.iter().fold(0xcbf2_9ce4_8422_2325_u64, |h, b| { + (h ^ u64::from(*b)).wrapping_mul(0x100_0000_01b3) + }); + println!( + "rows={rows} file_bytes={} fnv64={digest:016x} reps={reps}", + bytes.len() + ); + SESSION + .open_options() + .open_buffer(ByteBuffer::from(bytes)) + .unwrap() + }); + + // A deletion-style selection: exclude 1% of rows, spread across every chunk. + let excluded: Buffer = (0..rows as u64).step_by(100).collect(); + let point_rows: Vec = { + let mut rng = Rng(42); + (0..200).map(|_| rng.next() % rows as u64).collect() + }; + let lo = (rows as i64) / 2; + let hi = lo + (rows as i64) / 10; + type Case<'a> = (&'a str, Box (usize, u64) + 'a>); + let cases: Vec> = vec![ + ( + "full projection a,b,s", + Box::new(|| rt.block_on(scan(&file, select(["a", "b", "s"], root()), None, None))), + ), + ( + "filter 10% range, project a,b,s", + Box::new(|| { + rt.block_on(scan( + &file, + select(["a", "b", "s"], root()), + Some(and( + gt(get_item("k", root()), lit(lo)), + lt(get_item("k", root()), lit(hi)), + )), + None, + )) + }), + ), + ( + "exclude 1% + filter a>0, project b,s", + Box::new(|| { + rt.block_on(scan( + &file, + select(["b", "s"], root()), + Some(gt(get_item("a", root()), lit(0_i64))), + Some(Selection::ExcludeByIndex( + StrictSortedBuffer::try_new(excluded.clone()).unwrap(), + )), + )) + }), + ), + ( + "200 one-row point reads, project s", + Box::new(|| { + let mut total = (0, 0); + for &row in &point_rows { + let (r, b) = rt.block_on(scan( + &file, + select(["s"], root()), + None, + Some(Selection::IncludeByIndex( + StrictSortedBuffer::try_new(Buffer::from_iter([row])).unwrap(), + )), + )); + total = (total.0 + r, total.1 + b); + } + total + }), + ), + ]; + for (_, case) in &cases { + case(); // warm-up + } + let mut cpu: Vec> = vec![Vec::new(); cases.len()]; + let mut wall: Vec> = vec![Vec::new(); cases.len()]; + let mut outs = vec![(0, 0); cases.len()]; + for _ in 0..reps { + for (i, (_, case)) in cases.iter().enumerate() { + let (c0, w0) = (cpu_ms(), Instant::now()); + outs[i] = case(); + cpu[i].push(cpu_ms() - c0); + wall[i].push(w0.elapsed().as_secs_f64() * 1e3); + } + } + println!( + "{:<40} {:>10} {:>10} {:>10} {:>10} {:>9} {:>12}", + "case", "cpu_med", "cpu_min", "cpu_max", "wall_med", "rows", "bytes" + ); + for (i, (name, _)) in cases.iter().enumerate() { + cpu[i].sort_by(f64::total_cmp); + wall[i].sort_by(f64::total_cmp); + let m = reps / 2; + println!( + "{name:<40} {:>10.1} {:>10.1} {:>10.1} {:>10.1} {:>9} {:>12}", + cpu[i][m], + cpu[i][0], + cpu[i][reps - 1], + wall[i][m], + outs[i].0, + outs[i].1 + ); + } + } +} + +fn main() { + #[cfg(unix)] + unix::main(); + #[cfg(not(unix))] + eprintln!("this benchmark measures CPU time with Unix APIs; it does not run on this platform"); +} diff --git a/vortex-file/benches/wide_string_cases.rs b/vortex-file/benches/wide_string_cases.rs new file mode 100644 index 00000000000..d85976bb103 --- /dev/null +++ b/vortex-file/benches/wide_string_cases.rs @@ -0,0 +1,531 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Cases where a data-block byte maximum could cost more than it saves, each written with and +//! without the maximum in one process and read with the arms interleaved: +//! +//! - `json3k`: ~3 KB JSON rows; full scans, a `LIKE` filter over the wide column, a range filter +//! projecting it, and one-row point reads. +//! - `shift_lo_hi` / `shift_hi_lo`: each 8,192-row block starts (or ends) with 1,024 rows drawn +//! from 8 distinct values and is otherwise unique JSON. The writer dictionary-encodes this +//! column, so the maximum does not apply; kept as a control. +//! - `rand_then_json` / `json_then_rand`: each 8,192-row block starts (or ends) with 1,024 +//! incompressible rows and is otherwise JSON, so the encodings chosen on a block's first +//! chunk do not fit the rest of it. +//! - `wide100k` / `wide1m`: 100 KB and 1 MB rows. +//! +//! Reads go through a counting reader that coalesces like an object store (1 MiB gap, 16 MiB +//! max), so the request counts are what an object-store tier would issue. +//! +//! `cargo bench -p vortex-file --bench wide_string_cases`; env `WSC_CASES` (comma-separated +//! dataset names, default all), `WSC_REPS` (scan repetitions per arm, default 11), +//! `WSC_POINTS` (point reads per repetition, default 1,000), `WSC_MAX_MB` (default 4). +//! +//! It measures CPU time with Unix APIs, so on other platforms it only says so. + +#[cfg(unix)] +#[path = "../tests/common/mod.rs"] +mod common; + +#[cfg(unix)] +mod unix { + #![expect(clippy::unwrap_used)] + #![expect(clippy::cast_precision_loss)] + #![expect(clippy::cast_possible_truncation)] + #![expect(clippy::cast_possible_wrap)] + #![expect(clippy::many_single_char_names)] + #![expect(clippy::panic, reason = "an unknown dataset name is a usage error")] + + use std::sync::Arc; + use std::sync::LazyLock; + use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; + use std::time::Instant; + + use futures::FutureExt; + use futures::StreamExt; + use futures::future::BoxFuture; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::buffer::BufferHandle; + use vortex_array::dtype::FieldNames; + use vortex_array::expr::Expression; + use vortex_array::expr::and; + use vortex_array::expr::get_item; + use vortex_array::expr::gt; + use vortex_array::expr::like; + use vortex_array::expr::lit; + use vortex_array::expr::lt; + use vortex_array::expr::root; + use vortex_array::validity::Validity; + use vortex_buffer::Alignment; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + use vortex_file::OpenOptionsSessionExt; + use vortex_file::VortexFile; + use vortex_file::WriteOptionsSessionExt; + use vortex_file::WriteStrategyBuilder; + use vortex_io::CoalesceConfig; + use vortex_io::VortexReadAt; + use vortex_io::session::RuntimeSession; + use vortex_layout::session::LayoutSession; + use vortex_scan::selection::Selection; + use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + super::common::enable_all_registered_array_encodings(&session); + session + }); + + const NEEDLE: &str = "needle-7f3a9c"; + + fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + } + + fn cpu_ms() -> f64 { + // SAFETY: getrusage fills the zeroed struct. + let mut ru: libc::rusage = unsafe { std::mem::zeroed() }; + // SAFETY: valid pointer to a rusage. + unsafe { libc::getrusage(libc::RUSAGE_SELF, &raw mut ru) }; + let tv = |t: libc::timeval| t.tv_sec as f64 * 1e3 + t.tv_usec as f64 / 1e3; + tv(ru.ru_utime) + tv(ru.ru_stime) + } + + struct Rng(u64); + impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + self.0 + } + } + + const WORDS: [&str; 16] = [ + "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", "india", + "juliet", "kilo", "lima", "mike", "november", "oscar", "papa", + ]; + + /// A JSON document of about `target` bytes: repeated objects of words and numbers, so it + /// compresses like real documents rather than like random bytes or a single repeated string. + fn json(rng: &mut Rng, i: usize, target: usize) -> String { + let mut s = format!("{{\"id\":{i},\"items\":["); + while s.len() < target { + let w = WORDS[(rng.next() % 16) as usize]; + let v = WORDS[(rng.next() % 16) as usize]; + s.push_str(&format!( + "{{\"k\":\"{w}-{}\",\"n\":{},\"t\":\"{v} {w} {}\"}},", + rng.next() % 10_000, + rng.next() % 1_000_000, + rng.next() % 97 + )); + } + s.push_str("]}"); + s + } + + struct Dataset { + name: &'static str, + rows: Vec, + } + + fn dataset(name: &'static str) -> Dataset { + let mut rng = Rng(0x9e37_79b9_7f4a_7c15 ^ name.len() as u64); + let needle = |i: usize, s: String| { + if i.is_multiple_of(997) { + format!("{s}{NEEDLE}") + } else { + s + } + }; + let rows: Vec = match name { + "json3k" => (0..98_304) + .map(|i| needle(i, json(&mut rng, i, 3_000))) + .collect(), + "shift_lo_hi" | "shift_hi_lo" => { + let fixed: Vec = (0..8).map(|i| json(&mut Rng(7 + i), 0, 3_000)).collect(); + (0..65_536) + .map(|i| { + let in_block = i % 8_192; + let low = if name == "shift_lo_hi" { + in_block < 1_024 + } else { + in_block >= 7_168 + }; + let s = if low { + fixed[i % 8].clone() + } else { + json(&mut rng, i, 3_000) + }; + needle(i, s) + }) + .collect() + } + // Incompressible rows (hex of random words) for 1,024 rows of each 8,192-row block + // and JSON for the rest, in either order. 3 KB rows under a 4 MiB maximum split into + // 1,024-row chunks, so the encodings are chosen on exactly the odd-one-out part (or on + // the JSON, and replayed on it). + "rand_then_json" | "json_then_rand" => (0..65_536) + .map(|i| { + let in_block = i % 8_192; + let odd = if name == "rand_then_json" { + in_block < 1_024 + } else { + in_block >= 7_168 + }; + let s = if odd { + (0..188) + .map(|_| format!("{:016x}", rng.next())) + .collect::() + } else { + json(&mut rng, i, 3_000) + }; + needle(i, s) + }) + .collect(), + "wide100k" => (0..2_048) + .map(|i| needle(i, json(&mut rng, i, 100_000))) + .collect(), + "wide1m" => (0..256) + .map(|i| needle(i, json(&mut rng, i, 1_000_000))) + .collect(), + other => panic!("unknown dataset {other}"), + }; + Dataset { name, rows } + } + + fn table(rows: &[String]) -> ArrayRef { + let ids = PrimitiveArray::from_iter(0..rows.len() as i64).into_array(); + let wide = VarBinViewArray::from_iter_str(rows.iter().map(String::as_str)).into_array(); + StructArray::new( + FieldNames::from(["id", "wide"]), + vec![ids, wide], + rows.len(), + Validity::NonNullable, + ) + .into_array() + } + + /// A reader over an in-memory file that coalesces like an object store and counts what it is + /// asked for. + #[derive(Clone)] + struct CountingReadAt { + inner: ByteBuffer, + requests: Arc, + bytes: Arc, + } + + impl CountingReadAt { + fn take(&self) -> (u64, u64) { + ( + self.requests.swap(0, Ordering::Relaxed), + self.bytes.swap(0, Ordering::Relaxed), + ) + } + } + + impl VortexReadAt for CountingReadAt { + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig::object_storage()) + } + + fn concurrency(&self) -> usize { + 16 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let len = self.inner.len() as u64; + async move { Ok(len) }.boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + self.requests.fetch_add(1, Ordering::Relaxed); + self.bytes.fetch_add(length as u64, Ordering::Relaxed); + self.inner.read_at(offset, length, alignment) + } + } + + struct Arm { + label: &'static str, + file: VortexFile, + reader: CountingReadAt, + } + + async fn write(data: &ArrayRef, max: Option) -> Vec { + let strategy = WriteStrategyBuilder::default() + .with_data_block_max_bytes(max) + .build(); + let mut bytes = Vec::new(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut bytes, data.to_array_stream()) + .await + .unwrap(); + bytes + } + + /// Runs a scan; returns (rows, a hash of every output value) so the arms can be checked equal. + async fn scan( + file: &VortexFile, + projection: &Expression, + filter: Option<&Expression>, + selection: Option, + hash: bool, + ) -> (usize, u64) { + let dtype = file.dtype().clone(); + let mut builder = file + .scan() + .unwrap() + .with_projection(projection.clone().bind(&dtype).unwrap()); + if let Some(f) = filter { + builder = builder.with_filter(f.clone().bind(&dtype).unwrap()); + } + if let Some(s) = selection { + builder = builder.with_selection(s); + } + let mut stream = builder.into_array_stream().unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let (mut rows, mut h) = (0, 0xcbf2_9ce4_8422_2325_u64); + while let Some(chunk) = stream.next().await { + let canonical = chunk.unwrap().execute::(&mut ctx).unwrap(); + rows += canonical.len(); + if hash { + match &canonical { + Canonical::VarBinView(v) => { + for i in 0..v.len() { + for b in v.bytes_at(i).as_slice() { + h = (h ^ u64::from(*b)).wrapping_mul(0x100_0000_01b3); + } + } + } + other => { + let ids = other.clone().into_primitive(); + for v in ids.as_slice::() { + h = (h ^ *v as u64).wrapping_mul(0x100_0000_01b3); + } + } + } + } + } + (rows, h) + } + + fn summarize(samples: &mut [f64]) -> String { + samples.sort_by(f64::total_cmp); + let n = samples.len(); + format!( + "med {:8.2} min {:8.2} max {:8.2}", + samples[n / 2], + samples[0], + samples[n - 1] + ) + } + + fn tail(samples: &mut [f64]) -> String { + samples.sort_by(f64::total_cmp); + let n = samples.len(); + let p = |q: f64| samples[((n as f64 * q) as usize).min(n - 1)]; + let p999 = if n >= 10_000 { + format!("{:7.3}", p(0.999)) + } else { + " n/a".to_owned() + }; + format!( + "n={n} p50 {:7.3} p99 {:7.3} p99.9 {p999} ms", + p(0.5), + p(0.99) + ) + } + + pub fn main() { + let reps = env_usize("WSC_REPS", 11); + let points = env_usize("WSC_POINTS", 1_000); + let max = env_usize("WSC_MAX_MB", 4) as u64 * 1024 * 1024; + let cases = std::env::var("WSC_CASES") + .unwrap_or_else(|_| "json3k,shift_lo_hi,shift_hi_lo,wide100k,wide1m".to_owned()); + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(8) + .enable_all() + .build() + .unwrap(); + println!("reps={reps} points={points} max_bytes={max}"); + + for name in cases.split(',') { + let ds = dataset(match name { + "json3k" => "json3k", + "shift_lo_hi" => "shift_lo_hi", + "shift_hi_lo" => "shift_hi_lo", + "rand_then_json" => "rand_then_json", + "json_then_rand" => "json_then_rand", + "wide100k" => "wide100k", + "wide1m" => "wide1m", + other => panic!("unknown dataset {other}"), + }); + let data = table(&ds.rows); + let n = ds.rows.len(); + let raw: usize = ds.rows.iter().map(String::len).sum(); + drop(ds.rows); + println!( + "\n### {} rows={n} raw={:.1}MiB", + ds.name, + raw as f64 / 1048576.0 + ); + + // Writes, interleaved A B B A, three of each. + let mut write_cpu = [Vec::new(), Vec::new()]; + let mut files: [Option>; 2] = [None, None]; + for arm in [0, 1, 1, 0, 0, 1, 1, 0, 0, 1] { + let c0 = cpu_ms(); + let bytes = rt.block_on(write(&data, (arm == 1).then_some(max))); + write_cpu[arm].push(cpu_ms() - c0); + files[arm] = Some(bytes); + } + drop(data); + let arms: Vec = files + .into_iter() + .enumerate() + .map(|(i, bytes)| { + let bytes = bytes.unwrap(); + let reader = CountingReadAt { + inner: ByteBuffer::from(bytes), + requests: Arc::default(), + bytes: Arc::default(), + }; + let file = rt + .block_on(SESSION.open_options().open_read(reader.clone())) + .unwrap(); + Arm { + label: if i == 0 { "no max" } else { "max" }, + file, + reader, + } + }) + .collect(); + for (i, arm) in arms.iter().enumerate() { + println!( + "{:7} write_cpu_ms {} file {:.2}MiB segments {}", + arm.label, + summarize(&mut write_cpu[i].clone()), + arm.reader.inner.len() as f64 / 1048576.0, + arm.file.footer().segment_map().len() + ); + } + + let wide = get_item("wide", root()); + let id = get_item("id", root()); + let lo = (n as i64) / 2; + let hi = lo + (n as i64) / 10; + let range = and(gt(id.clone(), lit(lo)), lt(id.clone(), lit(hi))); + let pattern = like(wide.clone(), lit(format!("%{NEEDLE}%"))); + let scans: [(&str, &Expression, Option<&Expression>); 3] = [ + ("full scan of wide", &wide, None), + ("LIKE on wide, project id", &id, Some(&pattern)), + ("10% id range, project wide", &wide, Some(&range)), + ]; + + for (label, projection, filter) in scans { + // Untimed check that both arms return the same values, and the request counts. + let mut seen = Vec::new(); + for arm in &arms { + arm.reader.take(); + let out = rt.block_on(scan(&arm.file, projection, filter, None, true)); + seen.push((out, arm.reader.take())); + } + assert_eq!( + seen[0].0, seen[1].0, + "{label}: the arms returned different values" + ); + let mut cpu = [Vec::new(), Vec::new()]; + let mut wall = [Vec::new(), Vec::new()]; + for _ in 0..reps { + for i in [0, 1, 1, 0] { + let (c0, w0) = (cpu_ms(), Instant::now()); + rt.block_on(scan(&arms[i].file, projection, filter, None, false)); + cpu[i].push(cpu_ms() - c0); + wall[i].push(w0.elapsed().as_secs_f64() * 1e3); + } + } + println!(" {label} (rows {})", seen[0].0.0); + for i in 0..2 { + println!( + " {:7} cpu_ms {} wall_ms {} requests {:6} bytes {:.2}MiB", + arms[i].label, + summarize(&mut cpu[i]), + summarize(&mut wall[i]), + seen[i].1.0, + seen[i].1.1 as f64 / 1048576.0 + ); + } + } + + // One-row point reads, arms interleaved per batch. + let rows_to_read: Vec = { + let mut rng = Rng(42); + (0..points).map(|_| rng.next() % n as u64).collect() + }; + let mut lat = [Vec::new(), Vec::new()]; + let mut cpu = [0.0, 0.0]; + let mut reqs = [(0, 0), (0, 0)]; + let batches = 10; + for b in 0..batches { + for i in if b % 2 == 0 { [0, 1] } else { [1, 0] } { + arms[i].reader.take(); + let c0 = cpu_ms(); + for &row in &rows_to_read { + let w0 = Instant::now(); + let sel = Selection::IncludeByIndex( + StrictSortedBuffer::try_new(Buffer::from_iter([row])).unwrap(), + ); + let (r, _) = + rt.block_on(scan(&arms[i].file, &wide, None, Some(sel), false)); + assert_eq!(r, 1); + lat[i].push(w0.elapsed().as_secs_f64() * 1e3); + } + cpu[i] += cpu_ms() - c0; + let (q, by) = arms[i].reader.take(); + reqs[i] = (reqs[i].0 + q, reqs[i].1 + by); + } + } + println!(" one-row point reads of wide"); + let total = (points * batches) as f64; + for i in 0..2 { + println!( + " {:7} cpu_ms/read {:.3} wall {} requests/read {:.2} bytes/read {:.0}", + arms[i].label, + cpu[i] / total, + tail(&mut lat[i]), + reqs[i].0 as f64 / total, + reqs[i].1 as f64 / total + ); + } + } + } +} + +fn main() { + #[cfg(unix)] + unix::main(); + #[cfg(not(unix))] + eprintln!("this benchmark measures CPU time with Unix APIs; it does not run on this platform"); +} diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 47218ed532a..f087cf6c6e8 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -68,6 +68,7 @@ pub struct WriteStrategyBuilder { compressor: CompressorConfig, row_block_size: usize, data_block_target_bytes: Option, + data_block_max_bytes: Option, field_writers: HashMap>, flat_strategy: Option>, probe_compressor: Option>, @@ -108,6 +109,7 @@ impl Default for WriteStrategyBuilder { compressor: CompressorConfig::BtrBlocks(BtrBlocksCompressorBuilder::default()), row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), + data_block_max_bytes: None, field_writers: HashMap::new(), flat_strategy: None, probe_compressor: None, @@ -135,6 +137,22 @@ impl WriteStrategyBuilder { self } + /// Cap the uncompressed bytes a data block of a string or binary column references. + /// + /// A data block holds a multiple of [`Self::with_row_block_size`] rows. For a fixed-width + /// type the row count shrinks to stay near [`Self::with_data_block_target_bytes`]; a string + /// or binary column has no fixed width, so its blocks hold the full row count whatever its + /// rows weigh, and a reader that needs one row reads the whole block. With a maximum, a + /// block over it is split into blocks of a power-of-two number of rows, each at or under + /// about the maximum. Zone statistics keep [`Self::with_row_block_size`] rows. Encodings + /// are chosen once per [`Self::with_row_block_size`] rows and reused for the split blocks. + /// + /// `None`, the default, leaves blocks at the row block size. + pub fn with_data_block_max_bytes(mut self, max_bytes: Option) -> Self { + self.data_block_max_bytes = max_bytes; + self + } + /// Enable writing list fields with [`ListLayoutStrategy`]. /// /// **Note**: this is an unstable and experimental layout that is expected to change. @@ -221,7 +239,15 @@ impl WriteStrategyBuilder { ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - let compressing = CompressingStrategy::new(buffered, data_compressor); + // Choose encodings once per row block when a block is split into smaller data blocks + // (see `with_data_block_max_bytes`); without a maximum no block is split. + let compressing = CompressingStrategy::new(buffered, data_compressor).with_group_rows( + if self.data_block_max_bytes.is_some() { + self.row_block_size + } else { + 0 + }, + ); // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( @@ -238,7 +264,8 @@ impl WriteStrategyBuilder { block_size_target: self.data_block_target_bytes, canonicalize: true, }, - ); + ) + .with_block_size_maximum(self.data_block_max_bytes); // 2.1. | 3.1. compress stats tables and dict values. let stats_compressor: Arc = match compressor { diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index 3f69de67394..788bb4d222c 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -18,6 +18,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::FieldNames; use vortex_array::field_path; +use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_buffer::ByteBuffer; @@ -215,3 +216,148 @@ async fn test_write_empty_nullable_struct_column() { } assert_eq!(rows, 0); } + +/// Writes `data` with the default strategy and an optional data block maximum, and returns the +/// file read back with its segment count. +async fn roundtrip_with_block_max( + data: vortex_array::ArrayRef, + max_bytes: Option, +) -> vortex_error::VortexResult<(vortex_array::ArrayRef, usize)> { + let strategy = vortex_file::WriteStrategyBuilder::default() + .with_data_block_max_bytes(max_bytes) + .build(); + let mut bytes = Vec::new(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut bytes, data.to_array_stream()) + .await?; + let vxf = SESSION + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let segments = vxf.footer().segment_map().len(); + let read = vxf.scan()?.into_array_stream()?.read_all().await?; + Ok((read, segments)) +} + +/// A data block maximum splits the blocks of a wide string column and changes no value: every +/// row, including nulls and rows of other columns, reads back as written. +#[tokio::test] +async fn test_data_block_max_bytes_roundtrip() -> vortex_error::VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rows = 20_000; + let wide = vortex_array::arrays::VarBinViewArray::from_iter_nullable_str( + (0..rows) + .map(|i| (i % 97 != 0).then(|| format!("{i:08}-{}", "payload ".repeat(70 + i % 13)))), + ) + .into_array(); + let ids = PrimitiveArray::from_iter(0..rows as i64).into_array(); + let data = StructArray::new( + FieldNames::from(["id", "wide"]), + vec![ids, wide], + rows, + Validity::NonNullable, + ) + .into_array(); + + let (default_read, default_segments) = roundtrip_with_block_max(data.clone(), None).await?; + let (split_read, split_segments) = + roundtrip_with_block_max(data.clone(), Some(256 * 1024)).await?; + + vortex_array::assert_arrays_eq!(default_read, data, &mut ctx); + vortex_array::assert_arrays_eq!(split_read, data, &mut ctx); + // ~600-byte rows: a 256 KiB maximum splits each 8,192-row block of `wide` into 256-row + // blocks, so the file holds many more segments than without it. + assert!( + split_segments > default_segments * 4, + "expected the maximum to split `wide`: {split_segments} segments vs {default_segments}" + ); + Ok(()) +} + +/// Random bytes, `width` per row: no compression scheme shrinks them. +fn random_binary(rows: usize, width: usize, seed: u64) -> vortex_array::ArrayRef { + let mut state = seed; + let values: Vec> = (0..rows) + .map(|_| { + (0..width) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state >> 56) as u8 + }) + .collect() + }) + .collect(); + vortex_array::arrays::VarBinViewArray::from_iter_bin(values.iter().map(Vec::as_slice)) + .into_array() +} + +/// Writes `data` with `max`, through BtrBlocks or a compressor that stores chunks as they are, +/// and returns the file size and its largest segment, checking the values read back. +async fn file_and_largest_segment( + data: &vortex_array::ArrayRef, + max: Option, + pass_through: bool, +) -> vortex_error::VortexResult<(usize, u32)> { + let mut builder = vortex_file::WriteStrategyBuilder::default().with_data_block_max_bytes(max); + if pass_through { + builder = builder.with_compressor( + |chunk: &vortex_array::ArrayRef, _: &mut vortex_array::ExecutionCtx| Ok(chunk.clone()), + ); + } + let mut bytes = Vec::new(); + SESSION + .write_options() + .with_strategy(builder.build()) + .write(&mut bytes, data.to_array_stream()) + .await?; + let file_len = bytes.len(); + let vxf = SESSION + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let largest = vxf + .footer() + .segment_map() + .iter() + .map(|s| s.length) + .max() + .unwrap_or(0); + let read = vxf.scan()?.into_array_stream()?.read_all().await?; + vortex_array::assert_arrays_eq!(read, data.clone(), &mut SESSION.create_execution_ctx()); + Ok((file_len, largest)) +} + +/// Splitting a block must not store its whole data buffer in every piece: when nothing shrinks +/// the bytes (random binary, or a compressor that stores chunks as they are), each piece's +/// segment holds only its own rows. Covers a block split as it arrives and one rebuilt by +/// coalescing two separately allocated chunks. +#[tokio::test] +async fn split_blocks_of_incompressible_binary_store_only_their_rows() +-> vortex_error::VortexResult<()> { + let one_block = random_binary(8192, 1000, 7); + // 100-byte rows (~0.9 MiB, under the coalescing minimum), then 480-byte rows (~3.9 MiB). + let coalesced = vortex_array::arrays::ChunkedArray::try_new( + vec![random_binary(8192, 100, 11), random_binary(8192, 480, 13)], + one_block.dtype().clone(), + )? + .into_array(); + let maximum: u64 = 4 << 20; + for (name, data) in [("one block", &one_block), ("coalesced", &coalesced)] { + for pass_through in [false, true] { + let (plain, _) = file_and_largest_segment(data, None, pass_through).await?; + let (split, largest) = + file_and_largest_segment(data, Some(maximum), pass_through).await?; + assert!( + split <= plain + plain / 50, + "{name}, pass-through {pass_through}: split file {split} bytes vs {plain}" + ); + assert!( + u64::from(largest) <= maximum, + "{name}, pass-through {pass_through}: largest split segment {largest} bytes" + ); + } + } + Ok(()) +} diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index b5f4f5b438d..d51cd646b08 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -351,9 +351,15 @@ impl LayoutReader for ChunkedReader { let mut chunk_evals = vec![]; for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) { + let chunk_mask = mask.slice(mask_range); + // A chunk the mask already excludes passes nothing, so it need not be read. + if chunk_mask.known().is_some_and(Mask::all_false) { + chunk_evals.push(MaskFuture::ready(Mask::new_false(chunk_mask.len()))); + continue; + } let chunk_reader = self.chunk_reader(chunk_idx)?; let chunk_eval = chunk_reader - .filter_evaluation(&chunk_range, expr, mask.slice(mask_range)) + .filter_evaluation(&chunk_range, expr, chunk_mask) .map_err(|err| { err.with_context(format!("While evaluating filter on chunk {chunk_idx}")) })?; @@ -390,15 +396,25 @@ impl LayoutReader for ChunkedReader { let mut chunk_evals = vec![]; for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) { + let chunk_mask = mask.slice(mask_range); + // A chunk the mask already excludes contributes no rows, so its data is never + // requested: a point read that selects one row of a range reads one chunk. + if chunk_mask.known().is_some_and(Mask::all_false) { + continue; + } let chunk_reader = self.chunk_reader(chunk_idx)?; let chunk_eval = chunk_reader - .projection_evaluation(&chunk_range, expr, mask.slice(mask_range)) + .projection_evaluation(&chunk_range, expr, chunk_mask) .map_err(|err| { err.with_context(format!("While evaluating projection on chunk {chunk_idx}")) })?; chunk_evals.push(chunk_eval); } + if chunk_evals.is_empty() { + return Ok(future::ready(Ok(Canonical::empty(expr.dtype()).into_array())).boxed()); + } + Ok(async move { // Split the mask over each chunk. let chunks: Vec<_> = FuturesOrdered::from_iter(chunk_evals).try_collect().await?; @@ -425,6 +441,8 @@ impl LayoutReader for ChunkedReader { #[cfg(test)] mod test { use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use futures::stream; use rstest::fixture; @@ -433,6 +451,7 @@ mod test { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; @@ -442,6 +461,7 @@ mod test { use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; + use vortex_mask::Mask; use vortex_session::registry::ReadContext; use crate::LayoutRef; @@ -452,6 +472,7 @@ mod test { use crate::layouts::flat::FlatLayout; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::scan::split_by::SplitBy; + use crate::segments::SegmentFuture; use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::TestSegments; @@ -570,4 +591,88 @@ mod test { assert_arrays_eq!(result, expected, &mut ctx); }) } + + /// Counts the segments a reader requests. + struct CountingSegmentSource { + inner: Arc, + requests: Arc, + } + + impl SegmentSource for CountingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.requests.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + } + + #[rstest] + #[case::one_row_of_the_middle_chunk(vec![4], vec![5], 1)] + #[case::rows_of_the_first_and_last_chunks(vec![0, 8], vec![1, 9], 2)] + #[case::no_rows(vec![], vec![], 0)] + fn a_known_mask_reads_only_the_chunks_it_selects( + #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), + #[case] selected: Vec, + #[case] expected: Vec, + #[case] expected_requests: usize, + ) { + block_on(|_h| async { + let mut ctx = SESSION.create_execution_ctx(); + let requests = Arc::new(AtomicUsize::new(0)); + let source = Arc::new(CountingSegmentSource { + inner: segments, + requests: Arc::clone(&requests), + }); + let reader = layout + .new_reader("".into(), source, &SESSION, &Default::default()) + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let rows = usize::try_from(layout.row_count()).unwrap(); + let mask = Mask::from_indices(rows, selected); + + let result = reader + .projection_evaluation(&(0..layout.row_count()), &expr, MaskFuture::ready(mask)) + .unwrap() + .await + .unwrap(); + + let expected = PrimitiveArray::from_iter(expected).into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!( + requests.load(Ordering::Relaxed), + expected_requests, + "segments requested" + ); + }) + } + + /// Without a known mask every chunk in range is read, as before. + #[rstest] + fn an_unresolved_mask_reads_every_chunk( + #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), + ) { + block_on(|_h| async { + let mut ctx = SESSION.create_execution_ctx(); + let requests = Arc::new(AtomicUsize::new(0)); + let source = Arc::new(CountingSegmentSource { + inner: segments, + requests: Arc::clone(&requests), + }); + let reader = layout + .new_reader("".into(), source, &SESSION, &Default::default()) + .unwrap(); + let expr = root().bind(reader.dtype()).unwrap(); + let rows = usize::try_from(layout.row_count()).unwrap(); + let mask = Mask::from_indices(rows, vec![4]); + let pending = MaskFuture::new(rows, async move { Ok(mask) }); + assert!(pending.known().is_none()); + + let result = reader + .projection_evaluation(&(0..layout.row_count()), &expr, pending) + .unwrap() + .await + .unwrap(); + assert_arrays_eq!(result, buffer![5i32].into_array(), &mut ctx); + assert_eq!(requests.load(Ordering::Relaxed), 3); + }) + } } diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 87af23bb469..d5792c64a04 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -4,13 +4,19 @@ use std::sync::Arc; use async_trait::async_trait; +use futures::Stream; use futures::StreamExt as _; +use futures::stream; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::dtype::DType; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; @@ -20,6 +26,7 @@ use crate::LayoutStrategy; use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; use crate::sequence::SequencePointer; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; @@ -29,12 +36,38 @@ use crate::sequence::SequentialStreamExt; /// API consumers are free to implement this trait to provide new plugin compressors. pub trait CompressorPlugin: Send + Sync + 'static { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Compresses consecutive chunks of one column, returning one compressed chunk per input. + /// + /// The chunks are drawn from the same data, so a compressor may choose its encodings once + /// and reuse them across the group. The default compresses each chunk on its own. + /// + /// A wrapper must forward this method to the compressor it wraps: inheriting the default + /// silently drops whatever reuse the inner compressor does. + fn compress_chunks( + &self, + chunks: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + chunks + .iter() + .map(|chunk| self.compress_chunk(chunk, ctx)) + .collect() + } } impl CompressorPlugin for Arc { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.as_ref().compress_chunk(chunk, ctx) } + + fn compress_chunks( + &self, + chunks: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.as_ref().compress_chunks(chunks, ctx) + } } impl CompressorPlugin for F @@ -50,6 +83,33 @@ impl CompressorPlugin for BtrBlocksCompressor { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.compress(chunk, ctx) } + + /// Searches for encodings once for the whole group and applies them to every chunk. + /// + /// Choosing an encoding compresses a sample once per candidate, and a chunk of at most + /// 1,024 rows is its own sample, so searching every small chunk costs several full + /// compressions each. The search runs on the sample the compressor would draw from the + /// group as one array, so each chunk is encoded as the unsplit group would have been. A + /// search on the first chunk alone fits only that chunk: when it differs from the rest + /// (incompressible rows ahead of JSON, say), the rest are encoded to suit it. + fn compress_chunks( + &self, + chunks: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let [first, ..] = chunks else { + return Ok(Vec::new()); + }; + if chunks.len() == 1 { + return Ok(vec![self.compress(first, ctx)?]); + } + let group = ChunkedArray::try_new(chunks.to_vec(), first.dtype().clone())?.into_array(); + let plan = self.plan_for(&group, ctx)?; + chunks + .iter() + .map(|chunk| self.compress_with_plan(chunk, &plan, ctx)) + .collect() + } } /// A layout writer that compresses chunks. @@ -59,6 +119,9 @@ pub struct CompressingStrategy { compressor: Arc, stats: Arc<[Stat]>, concurrency: usize, + /// Consecutive chunks are compressed together until they hold this many rows; see + /// [`Self::with_group_rows`]. + group_rows: usize, } impl CompressingStrategy { @@ -69,9 +132,25 @@ impl CompressingStrategy { compressor: Arc::new(compressor), stats: Stat::all().collect(), concurrency: get_available_parallelism().unwrap_or(1), + group_rows: 0, } } + /// Compress consecutive chunks together, through [`CompressorPlugin::compress_chunks`], + /// until a group holds at least `rows` rows. + /// + /// A chunk that holds `rows` or more on its own is compressed alone, as without grouping. + /// The default, 0, compresses every chunk alone. + /// + /// Only string and binary columns are grouped: those are the chunks a byte maximum splits + /// (see `RepartitionStrategy::with_block_size_maximum`). Other chunks smaller than `rows`, + /// such as the blocks of a wide fixed-size list, stay separate tasks that compress in + /// parallel. + pub fn with_group_rows(mut self, rows: usize) -> Self { + self.group_rows = rows; + self + } + pub fn with_concurrency(mut self, concurrency: usize) -> Self { self.concurrency = concurrency; self @@ -102,20 +181,41 @@ impl LayoutStrategy for CompressingStrategy { let compute_session = session.clone(); let handle = session.handle(); - let stream = stream - .map(move |chunk| { + let group_rows = if matches!(dtype, DType::Utf8(_) | DType::Binary(_)) { + self.group_rows + } else { + 0 + }; + let stream = group_chunks(stream, group_rows) + .map(move |group| { let compressor = Arc::clone(&compressor); let stats = Arc::clone(&stats); let session = compute_session.clone(); handle.spawn_cpu(move || { - let (sequence_id, chunk) = chunk?; + let group = group?; let mut ctx = session.create_execution_ctx(); - // Compute the stats for the chunk prior to compression - chunk.statistics().compute_all(&stats, &mut ctx)?; - Ok((sequence_id, compressor.compress_chunk(&chunk, &mut ctx)?)) + let (sequence_ids, chunks): (Vec<_>, Vec<_>) = group.into_iter().unzip(); + // Compute the stats for each chunk prior to compression + for chunk in &chunks { + chunk.statistics().compute_all(&stats, &mut ctx)?; + } + let compressed = compressor.compress_chunks(&chunks, &mut ctx)?; + vortex_ensure!( + compressed.len() == chunks.len(), + "compressor returned {} chunks for {} inputs", + compressed.len(), + chunks.len() + ); + Ok(sequence_ids.into_iter().zip(compressed).collect::>()) }) }) - .buffered(self.concurrency); + .buffered(self.concurrency) + .flat_map(|group: VortexResult>| { + stream::iter(match group { + Ok(chunks) => chunks.into_iter().map(Ok).collect::>(), + Err(err) => vec![Err(err)], + }) + }); self.child .write_stream( @@ -128,3 +228,315 @@ impl LayoutStrategy for CompressingStrategy { .await } } + +/// Groups consecutive chunks of `stream` that each hold fewer than `rows` rows, until a group +/// holds at least `rows` rows. A chunk of `rows` or more rows is always a group of its own, so +/// grouping never replays encodings onto a chunk large enough to search on its own. With `rows` +/// of 0 or 1 every chunk is its own group. An error ends the group in progress, which is +/// emitted first. +fn group_chunks( + stream: SendableSequentialStream, + rows: usize, +) -> impl Stream>> + Send { + type Item = VortexResult<(SequenceId, ArrayRef)>; + struct State { + stream: SendableSequentialStream, + group: Vec<(SequenceId, ArrayRef)>, + group_rows: usize, + /// An item read past the end of the group that was emitted, to start the next one. + carry: Option, + } + stream::unfold( + State { + stream, + group: Vec::new(), + group_rows: 0, + carry: None, + }, + move |mut state| async move { + loop { + let next = match state.carry.take() { + Some(item) => Some(item), + None => state.stream.next().await, + }; + match next { + Some(Ok((sequence_id, chunk))) => { + if chunk.len() >= rows && !state.group.is_empty() { + state.carry = Some(Ok((sequence_id, chunk))); + state.group_rows = 0; + return Some((Ok(std::mem::take(&mut state.group)), state)); + } + state.group_rows += chunk.len(); + state.group.push((sequence_id, chunk)); + if state.group_rows >= rows { + state.group_rows = 0; + return Some((Ok(std::mem::take(&mut state.group)), state)); + } + } + Some(Err(err)) => { + if state.group.is_empty() { + return Some((Err(err), state)); + } + state.carry = Some(Err(err)); + state.group_rows = 0; + return Some((Ok(std::mem::take(&mut state.group)), state)); + } + None if state.group.is_empty() => return None, + None => { + state.group_rows = 0; + return Some((Ok(std::mem::take(&mut state.group)), state)); + } + } + } + }, + ) +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; + use vortex_array::dtype::PType; + use vortex_error::VortexExpect; + use vortex_error::vortex_err; + use vortex_io::runtime::single::block_on; + use vortex_io::session::RuntimeSessionExt; + + use super::*; + use crate::layouts::chunked::writer::ChunkedLayoutStrategy; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::segments::TestSegments; + use crate::sequence::SequenceId; + use crate::sequence::SequentialArrayStreamExt; + use crate::test::new_session; + + fn group_lengths(lengths: &[usize], rows: usize) -> Vec> { + let chunks = lengths + .iter() + .map(|&len| PrimitiveArray::from_iter(0..len as i64).into_array()) + .collect::>(); + let array = ChunkedArray::try_new(chunks, DType::Primitive(PType::I64, NonNullable)) + .vortex_expect("chunked") + .into_array(); + let (ptr, _eof) = SequenceId::root().split(); + let stream = array.to_array_stream().sequenced(ptr); + block_on(|_| async move { + group_chunks(stream, rows) + .map(|group| { + group + .vortex_expect("group") + .iter() + .map(|(_, chunk)| chunk.len()) + .collect::>() + }) + .collect::>() + .await + }) + } + + #[test] + fn small_chunks_group_until_the_row_count_and_large_chunks_stay_alone() { + assert_eq!( + group_lengths(&[100, 100, 9000, 50, 8000, 8192, 4096, 4096, 10], 8192), + vec![ + vec![100, 100], + vec![9000], + vec![50, 8000], + vec![8192], + vec![4096, 4096], + vec![10] + ], + ); + } + + #[test] + fn without_a_row_count_every_chunk_is_its_own_group() { + assert_eq!( + group_lengths(&[100, 100, 9000], 0), + vec![vec![100], vec![100], vec![9000]] + ); + } + + #[test] + fn an_error_ends_the_group_and_is_emitted_after_it() { + let dtype = DType::Primitive(PType::I64, NonNullable); + let (mut ptr, _eof) = SequenceId::root().split(); + let items: Vec> = vec![ + Ok(( + ptr.advance(), + PrimitiveArray::from_iter(0..10_i64).into_array(), + )), + Err(vortex_err!("read failed")), + Ok(( + ptr.advance(), + PrimitiveArray::from_iter(0..10_i64).into_array(), + )), + ]; + let stream = SequentialStreamAdapter::new(dtype, stream::iter(items)).sendable(); + let groups = block_on(|_| async move { + group_chunks(stream, 8192) + .map(|group| group.map(|g| g.len()).map_err(|e| e.to_string())) + .collect::>() + .await + }); + assert_eq!(groups.len(), 3, "{groups:?}"); + assert_eq!(groups[0], Ok(1)); + assert!( + groups[1].as_ref().is_err_and(|e| e.contains("read failed")), + "{groups:?}" + ); + assert_eq!(groups[2], Ok(1)); + } + + fn next(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + *state + } + + /// ~1.5 KB rows of hex digits of random words: no dictionary or symbol table built on + /// them fits other data. + fn hex_rows(rows: usize, seed: u64) -> ArrayRef { + let mut state = seed; + vortex_array::arrays::VarBinViewArray::from_iter_str((0..rows).map(|_| { + (0..96) + .map(|_| format!("{:016x}", next(&mut state))) + .collect::() + })) + .into_array() + } + + /// ~1.5 KB JSON documents of repeated keys and words with varying numbers. + fn documents(rows: usize, seed: u64) -> ArrayRef { + const WORDS: [&str; 8] = [ + "alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "hotel", + ]; + let mut state = seed; + vortex_array::arrays::VarBinViewArray::from_iter_str((0..rows).map(|i| { + let mut s = format!("{{\"id\":{i},\"items\":["); + while s.len() < 1_500 { + let v = next(&mut state); + s.push_str(&format!( + "{{\"k\":\"{}-{}\",\"n\":{},\"t\":\"{} {}\"}},", + WORDS[(v % 8) as usize], + (v >> 8) % 10_000, + (v >> 24) % 1_000_000, + WORDS[((v >> 44) % 8) as usize], + (v >> 50) % 97 + )); + } + s.push_str("]}"); + s + })) + .into_array() + } + + /// A group whose first chunk is unlike the rest is encoded as the unsplit group would be: + /// the schemes are chosen on a sample of the whole group, not on its first chunk. Here the + /// first chunk alone picks FSST, and replaying that on the documents instead of the + /// group's own choice leaves them larger and slower to filter. + #[test] + fn a_group_is_encoded_as_its_unsplit_data_would_be() -> VortexResult<()> { + let mut ctx = crate::test::SESSION.create_execution_ctx(); + let mut chunks = vec![hex_rows(1024, 7)]; + chunks.extend((1..8).map(|i| documents(1024, i))); + let compressor = BtrBlocksCompressor::default(); + + let unsplit = + ChunkedArray::try_new(chunks.clone(), chunks[0].dtype().clone())?.into_array(); + let expected = compressor.compress(&unsplit, &mut ctx)?.encoding_id(); + // The premise: the first chunk on its own chooses differently. + let (first, plan) = compressor.compress_recording(&chunks[0], &mut ctx)?; + assert_ne!(first.encoding_id(), expected); + assert_eq!( + compressor + .compress_with_plan(&chunks[1], &plan, &mut ctx)? + .encoding_id(), + first.encoding_id() + ); + + let grouped = compressor.compress_chunks(&chunks, &mut ctx)?; + for (compressed, chunk) in grouped.iter().zip(&chunks) { + assert_eq!(compressed.encoding_id(), expected); + vortex_array::assert_arrays_eq!(compressed.clone(), chunk.clone(), &mut ctx); + } + Ok(()) + } + + /// Passes chunks through uncompressed, recording how many it is handed at a time. + #[derive(Clone, Default)] + struct GroupSizes(Arc>>); + + impl CompressorPlugin for GroupSizes { + fn compress_chunk(&self, chunk: &ArrayRef, _: &mut ExecutionCtx) -> VortexResult { + Ok(chunk.clone()) + } + + fn compress_chunks( + &self, + chunks: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.0.lock().push(chunks.len()); + chunks.iter().map(|c| self.compress_chunk(c, ctx)).collect() + } + } + + /// Writes `chunks` through a grouping compressor and returns the size of each group. + fn group_sizes(chunks: Vec) -> VortexResult> { + let sizes = GroupSizes::default(); + let strategy = CompressingStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + sizes.clone(), + ) + .with_group_rows(8192); + let dtype = chunks[0].dtype().clone(); + let (ptr, eof) = SequenceId::root().split(); + let stream = ChunkedArray::try_new(chunks, dtype)? + .into_array() + .to_array_stream() + .sequenced(ptr); + let segments = Arc::new(TestSegments::default()); + block_on(|handle| async move { + let session = new_session().with_handle(handle); + strategy + .write_stream( + ArrayContext::empty().into(), + segments, + stream, + eof, + &session, + ) + .await + })?; + let sizes = sizes.0.lock().clone(); + Ok(sizes) + } + + /// Small string chunks, the pieces a byte maximum splits a block into, are compressed as + /// one group; small chunks of other types, such as the blocks of a wide fixed-size list, + /// stay separate tasks so they compress in parallel. + #[test] + fn only_string_and_binary_chunks_are_grouped() -> VortexResult<()> { + let strings = (0..8) + .map(|i| { + vortex_array::arrays::VarBinViewArray::from_iter_str( + (0..1024).map(|row| format!("{i}-{row}")), + ) + .into_array() + }) + .collect(); + assert_eq!(group_sizes(strings)?, vec![8]); + + let ints = (0..8) + .map(|i| PrimitiveArray::from_iter((0..1024).map(|row| i * 1024 + row)).into_array()) + .collect(); + assert_eq!(group_sizes(ints)?, vec![1; 8]); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 7344ee7be51..4147cce9f0d 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::collections::VecDeque; +use std::ops::Range; use std::sync::Arc; use async_stream::try_stream; @@ -10,9 +11,12 @@ use futures::StreamExt as _; use futures::pin_mut; use vortex_array::ArrayRef; use vortex_array::Canonical; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::varbinview::BinaryView; use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -80,6 +84,7 @@ impl RepartitionWriterOptions { pub struct RepartitionStrategy { child: Arc, options: RepartitionWriterOptions, + block_size_maximum: Option, } impl RepartitionStrategy { @@ -87,8 +92,110 @@ impl RepartitionStrategy { Self { child: Arc::new(child), options, + block_size_maximum: None, } } + + /// Split a canonical string or binary chunk whose rows reference more than `maximum` + /// bytes into blocks of at most about `maximum` bytes. + /// + /// [`RepartitionWriterOptions::block_size_target`] sizes blocks of fixed-width types by + /// their element size; a variable-width type has none, so a block of a wide string column + /// holds `block_len_multiple` rows whatever their width, and a reader that wants one row + /// reads all of them. A split block holds a power-of-two number of rows, which divides + /// any power-of-two `block_len_multiple`, so its boundaries line up with the blocks of + /// every other column. Chunks at or under `maximum` are repartitioned as without it. + pub fn with_block_size_maximum(mut self, maximum: Option) -> Self { + self.block_size_maximum = maximum; + self + } +} + +/// The row ranges that split `chunk` into blocks each referencing at most `maximum` bytes, or +/// `None` when `chunk` needs no split. +/// +/// `chunk` is cut into blocks of `max_rows`, and a block over `maximum` is halved, repeatedly, +/// until each piece fits or holds one row. A block's average row width does not bound its parts: +/// wide rows clustered at one end of it would leave that part over `maximum`. Each piece starts +/// at a multiple of its power-of-two length within its block, so pieces line up with the blocks +/// of other columns wherever their splits agree. +/// +/// Only a canonical [`VarBinView`] chunk is measured, by the bytes its rows reference (each +/// view plus the bytes it points at), not by its buffers: rows sliced from a larger array share +/// that array's buffers. +fn split_ranges(chunk: &ArrayRef, maximum: u64, max_rows: usize) -> Option>> { + let views = chunk.as_opt::()?; + let len = chunk.len(); + if len <= 1 { + return None; + } + let view_bytes = size_of::() as u64; + // `prefix[i]` is the bytes rows `0..i` reference. + let mut prefix = Vec::with_capacity(len + 1); + prefix.push(0_u64); + let mut total = 0_u64; + for view in views.views().iter() { + total += u64::from(view.len()) + view_bytes; + prefix.push(total); + } + if total <= maximum { + return None; + } + + let mut ranges = Vec::new(); + let mut pending = Vec::new(); + let block = max_rows.max(1); + for start in (0..len).step_by(block) { + pending.push(start..(start + block).min(len)); + while let Some(range) = pending.pop() { + if range.len() <= 1 || prefix[range.end] - prefix[range.start] <= maximum { + ranges.push(range); + continue; + } + // The first half is a power of two, so the second starts at a multiple of it. + let mid = range.start + range.len().next_power_of_two() / 2; + pending.push(mid..range.end); + pending.push(range.start..mid); + } + } + (ranges.len() > 1).then_some(ranges) +} + +/// `block` as it is, or its pieces from [`split_ranges`] when it is over `maximum`. Coalescing +/// merges blocks that are each under the maximum into one that need not be. +fn under_maximum( + block: ArrayRef, + maximum: Option, + max_rows: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + match maximum.and_then(|maximum| split_ranges(&block, maximum, max_rows)) { + Some(ranges) => ranges + .into_iter() + .map(|range| split_piece(&block, range, ctx)) + .collect(), + None => Ok(vec![block]), + } +} + +/// Rows `range` of the [`VarBinView`] `block`, holding only the bytes those rows reference. +/// +/// A slice shares its parent's data buffers, and a buffer is written whole, so a piece passed +/// on as a slice would store every other piece's bytes too unless the compressor compacts it. +/// Compaction slices a buffer the piece uses contiguously, so this rarely copies. +fn split_piece( + block: &ArrayRef, + range: Range, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let piece = block.slice(range)?; + match piece.as_opt::() { + Some(views) => Ok(views + .into_owned() + .compact_with_threshold(1.0, ctx)? + .into_array()), + None => Ok(piece), + } } #[async_trait] @@ -128,6 +235,7 @@ impl LayoutStrategy for RepartitionStrategy { // segments. let block_len = options.effective_block_len(&dtype); let block_size_minimum = options.block_size_minimum; + let block_size_maximum = self.block_size_maximum; let repartition_session = session.clone(); let repartitioned_stream = try_stream! { @@ -139,6 +247,23 @@ impl LayoutStrategy for RepartitionStrategy { while let Some(chunk) = canonical_stream.as_mut().next().await { let (sequence_id, chunk) = chunk?; let mut sequence_pointer = sequence_id.descend(); + if let Some(ranges) = block_size_maximum + .and_then(|maximum| split_ranges(&chunk, maximum, block_len)) + { + // Blocks buffered from earlier chunks come first, as they are. + let pending = chunks.take_all(); + if !pending.is_empty() { + let pending = ChunkedArray::try_new(pending, dtype_clone.clone())?; + let pending = pending.into_array().execute::(&mut ctx)?.into_array(); + for block in under_maximum(pending, block_size_maximum, block_len, &mut ctx)? { + yield (sequence_pointer.advance(), block); + } + } + for range in ranges { + yield (sequence_pointer.advance(), split_piece(&chunk, range, &mut ctx)?); + } + continue; + } let mut offset = 0; while offset < chunk.len() { let end = (offset + block_len).min(chunk.len()); @@ -153,10 +278,9 @@ impl LayoutStrategy for RepartitionStrategy { ChunkedArray::try_new(output_chunks, dtype_clone.clone())?; if !chunked.is_empty() { let canonical = chunked.into_array().execute::(&mut ctx)?.into_array(); - yield ( - sequence_pointer.advance(), - canonical, - ) + for block in under_maximum(canonical, block_size_maximum, block_len, &mut ctx)? { + yield (sequence_pointer.advance(), block); + } } } } @@ -167,10 +291,9 @@ impl LayoutStrategy for RepartitionStrategy { )?; if !to_flush.is_empty() { let canonical = to_flush.into_array().execute::(&mut ctx)?.into_array(); - yield ( - sequence_pointer.advance(), - canonical, - ) + for block in under_maximum(canonical, block_size_maximum, block_len, &mut ctx)? { + yield (sequence_pointer.advance(), block); + } } } } @@ -252,6 +375,13 @@ impl ChunksBuffer { self.data.push_front((chunk, nb)); } + /// Removes every buffered chunk, in order. + fn take_all(&mut self) -> Vec { + self.row_count = 0; + self.nbytes = 0; + self.data.drain(..).map(|(chunk, _)| chunk).collect() + } + fn pop_front(&mut self) -> Option<(ArrayRef, u64)> { let res = self.data.pop_front(); if let Some((chunk, nb)) = res.as_ref() { @@ -274,6 +404,7 @@ mod tests { use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SharedArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; @@ -526,4 +657,163 @@ mod tests { Ok(()) } + + /// Writes `array` through a coalescing repartition with a block size maximum and returns + /// the row count of each block written. + fn block_rows_with_maximum(array: ArrayRef, maximum: Option) -> VortexResult> { + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let strategy = RepartitionStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + RepartitionWriterOptions { + block_size_minimum: ONE_MEG, + block_len_multiple: 8192, + block_size_target: Some(ONE_MEG), + canonicalize: true, + }, + ) + .with_block_size_maximum(maximum); + let stream = array.to_array_stream().sequenced(ptr); + let layout = block_on(|handle| async move { + let session = new_session().with_handle(handle); + strategy + .write_stream(ctx.into(), segments, stream, eof, &session) + .await + })?; + if layout.nchildren() == 0 { + // A single block is written as that block's own layout. + return Ok(vec![layout.row_count()]); + } + (0..layout.nchildren()) + .map(|i| { + Ok(layout + .slot(i)? + .vortex_expect("chunk slot present") + .row_count()) + }) + .collect() + } + + fn strings(rows: usize, width: usize) -> ArrayRef { + VarBinViewArray::from_iter_str((0..rows).map(|row| format!("{row:0width$}"))).into_array() + } + + #[test] + fn a_block_size_maximum_splits_wide_strings_into_power_of_two_blocks() -> VortexResult<()> { + // 16,384 rows of 1,000 bytes reference 1,016 bytes each (with their views): an 8,192-row + // block is ~8 MiB, so a 1 MiB maximum fits 1,032 rows, rounded down to 1,024. + let rows = block_rows_with_maximum(strings(16_384, 1000), Some(ONE_MEG))?; + assert_eq!(rows, vec![1024; 16]); + Ok(()) + } + + #[test] + fn without_a_maximum_wide_strings_keep_the_row_multiple() -> VortexResult<()> { + let rows = block_rows_with_maximum(strings(16_384, 1000), None)?; + assert_eq!(rows, vec![8192, 8192]); + Ok(()) + } + + #[test] + fn narrow_strings_under_the_maximum_coalesce_as_without_it() -> VortexResult<()> { + // 20-byte strings reference 8192 * 36 bytes, about 288 KiB, per 8,192-row block, far + // under the maximum, so the blocks are whatever coalescing makes of them without it. + // (Coalescing weighs a block by `nbytes()`, which counts the buffer it shares with the + // other slices of the input, so here no two blocks merge.) + let without = block_rows_with_maximum(strings(65_536, 20), None)?; + let with = block_rows_with_maximum(strings(65_536, 20), Some(4 * ONE_MEG))?; + assert_eq!(with, without); + assert_eq!(with.iter().sum::(), 65_536); + Ok(()) + } + + fn lengths(ranges: Option>>) -> Option> { + ranges.map(|ranges| ranges.iter().map(ExactSizeIterator::len).collect()) + } + + #[test] + fn split_ranges_measures_referenced_bytes_not_buffers() { + // Rows sliced from a larger array share its buffers; only the bytes they point at count. + let wide = strings(8192, 1000); + let slice = wide.slice(0..64).vortex_expect("slice"); + assert_eq!( + split_ranges(&slice, ONE_MEG, 8192), + None, + "64 rows reference ~64 KiB" + ); + assert_eq!( + lengths(split_ranges(&wide, ONE_MEG, 8192)), + Some(vec![1024; 8]) + ); + // A maximum smaller than one row still yields blocks of one row. + assert_eq!(lengths(split_ranges(&wide, 100, 8192)), Some(vec![1; 8192])); + // Non-string chunks are never split. + let ints = PrimitiveArray::from_iter(0..8192_i64).into_array(); + assert_eq!(split_ranges(&ints, 1, 8192), None); + } + + #[test] + fn clustered_wide_rows_are_split_until_every_block_fits() -> VortexResult<()> { + // 256 rows of 100,000 bytes, then 7,936 rows of 100: the block's average row fits a + // 1,024-row block under 4 MiB, but its first 1,024 rows reference ~25 MiB. + let widths: Vec = (0..8192) + .map(|i| if i < 256 { 100_000 } else { 100 }) + .collect(); + let array = VarBinViewArray::from_iter_str(widths.iter().map(|&width| "x".repeat(width))) + .into_array(); + let maximum = 4 * ONE_MEG; + let rows = block_rows_with_maximum(array, Some(maximum))?; + + assert_eq!(rows.iter().sum::(), 8192); + let mut start = 0; + for &n in &rows { + let n = usize::try_from(n).vortex_expect("rows fit usize"); + let referenced: usize = widths[start..start + n] + .iter() + .map(|w| w + size_of::()) + .sum(); + assert!( + n == 1 || referenced as u64 <= maximum, + "block of {n} rows at {start} references {referenced} bytes: {rows:?}" + ); + start += n; + } + Ok(()) + } + + #[test] + fn coalesced_blocks_stay_under_the_maximum() -> VortexResult<()> { + // Two separately allocated 8,192-row chunks, each under 4 MiB: 100-byte rows (~0.9 MiB, + // under the 1 MiB coalescing minimum) and 480-byte rows (~3.9 MiB). Coalescing the first + // with the second would make a ~4.8 MiB block. + let widths: Vec = (0..16_384) + .map(|i| if i < 8192 { 100 } else { 480 }) + .collect(); + let chunk = |range: Range| { + VarBinViewArray::from_iter_str(widths[range].iter().map(|&w| "x".repeat(w))) + .into_array() + }; + let dtype = chunk(0..1).dtype().clone(); + let array = + ChunkedArray::try_new(vec![chunk(0..8192), chunk(8192..16_384)], dtype)?.into_array(); + let maximum = 4 * ONE_MEG; + let rows = block_rows_with_maximum(array, Some(maximum))?; + + assert_eq!(rows.iter().sum::(), 16_384); + let mut start = 0; + for &n in &rows { + let n = usize::try_from(n).vortex_expect("rows fit usize"); + let referenced: usize = widths[start..start + n] + .iter() + .map(|w| w + size_of::()) + .sum(); + assert!( + n == 1 || referenced as u64 <= maximum, + "block of {n} rows at {start} references {referenced} bytes: {rows:?}" + ); + start += n; + } + Ok(()) + } }