From b768a4c7d33a849279c2792722f1400996bb63af Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:37:00 -0700 Subject: [PATCH 01/13] feat(file): cap the bytes of string and binary data blocks, choosing encodings once per row block A string or binary column is written in data blocks of `row_block_size` (8,192) rows whatever its rows weigh: `RepartitionWriterOptions::effective_block_len` shrinks blocks only for types with a fixed element size. A column of ~3 KB JSON documents gets ~12.5 MiB blocks, and since a reader never reads less than one block, a one-row lookup reads and copies all of it. `WriteStrategyBuilder::with_data_block_max_bytes` caps the bytes a block's rows reference. A block over the cap is split into blocks of a power-of-two number of rows, which divides the row block size, so split blocks line up with zone statistics and with every other column's blocks. Blocks are measured by the bytes their views reference, not by their buffers, which slices of a larger array share. The cap is off by default. Splitting alone makes writes 3-4x slower on such a column. Choosing an encoding compresses a sample once per candidate scheme, and a block of at most 1,024 rows is its own sample, so every small block paid several full compressions. `CascadingCompressor::compress_recording` returns the schemes a compression chose as a `CompressionPlan`, and `compress_with_plan` replays them without estimating; a path the plan never decided still searches, and a replayed scheme is still kept only when it shrinks its input. `CompressingStrategy::with_group_rows` compresses consecutive blocks smaller than the row count together through the new `CompressorPlugin::compress_chunks`, which BtrBlocks implements by searching the first block and replaying on the rest, so encodings are still chosen once per row block. A block of the full row count is always compressed alone, so columns whose blocks are not split compress exactly as before. On 98,304 rows of ~3 KB JSON with a 4 MiB cap: 0.82 MB read per single-row lookup instead of 6.1 MB, +7% file size, write CPU within 10% of the default (3.5x without encoding reuse). --- Cargo.lock | 4 + SPICE_PATCHES.md | 1 + vortex-btrblocks/Cargo.toml | 9 + .../benches/small_chunk_compress.rs | 304 ++++++++++++++++++ vortex-btrblocks/src/canonical_compressor.rs | 85 +++++ vortex-btrblocks/src/lib.rs | 1 + vortex-compressor/src/compressor/cascade.rs | 74 ++++- vortex-compressor/src/compressor/mod.rs | 1 + vortex-compressor/src/compressor/plan.rs | 109 +++++++ vortex-compressor/src/lib.rs | 1 + vortex-compressor/src/scheme/ctx.rs | 21 ++ vortex-file/src/strategy.rs | 26 +- vortex-file/tests/test_write_table.rs | 59 ++++ vortex-layout/src/layouts/compressed.rs | 256 ++++++++++++++- vortex-layout/src/layouts/repartition.rs | 159 +++++++++ 15 files changed, 1098 insertions(+), 12 deletions(-) create mode 100644 vortex-btrblocks/benches/small_chunk_compress.rs create mode 100644 vortex-compressor/src/compressor/plan.rs diff --git a/Cargo.lock b/Cargo.lock index 2c573148d20..6971a32c4f0 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", diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index cef260fd871..ba952f79354 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -60,6 +60,7 @@ 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::compress_recording`/`compress_with_plan` choose encodings once per row block and replay them on its split blocks | `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-file --test test_write_table test_data_block_max_bytes_roundtrip` | Proposed — additive; off unless the cap is set | **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-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..187af0fd618 --- /dev/null +++ b/vortex-btrblocks/benches/small_chunk_compress.rs @@ -0,0 +1,304 @@ +// 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). + +#![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); + let mut plan = from_sample.then(|| { + // A 1,024-row stratified sample (16 x 64 rows) of the block. + let stride = (block_end - block) / 16; + let slices: Vec = (0..16) + .map(|slice| { + let start = block + slice * stride; + array.slice(start..(start + 64).min(block_end)).unwrap() + }) + .collect(); + let sample = vortex_array::arrays::ChunkedArray::try_new(slices, array.dtype().clone()) + .unwrap() + .into_array(); + compressor.compress_recording(&sample, &mut ctx).unwrap().1 + }); + 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 +} + +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 P2 8k sample, {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 + ); + } + } +} diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index ed528145d71..ec8de8c0343 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -311,4 +311,89 @@ 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(()) + } } 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..475fa14345f 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,47 @@ 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())) + } + + /// 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 +339,33 @@ 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 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()); }; 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/src/strategy.rs b/vortex-file/src/strategy.rs index 47218ed532a..e934c93dee5 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,10 @@ impl WriteStrategyBuilder { ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - let compressing = CompressingStrategy::new(buffered, data_compressor); + // Choose encodings once per row block, even when a block is split into smaller data + // blocks (see `with_data_block_max_bytes`). + let compressing = CompressingStrategy::new(buffered, data_compressor) + .with_group_rows(self.row_block_size); // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( @@ -238,7 +259,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..b6846d950a7 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,61 @@ 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(()) +} diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 87af23bb469..5e36bd4b75e 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -4,13 +4,16 @@ 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::VortexSessionExecute; 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 +23,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 +33,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 +80,31 @@ impl CompressorPlugin for BtrBlocksCompressor { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.compress(chunk, ctx) } + + /// Searches for encodings on the first chunk and replays them on the rest. + /// + /// 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. Replaying keeps one search per group. + fn compress_chunks( + &self, + chunks: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some((first, rest)) = chunks.split_first() else { + return Ok(Vec::new()); + }; + if rest.is_empty() { + return Ok(vec![self.compress(first, ctx)?]); + } + let (compressed, plan) = self.compress_recording(first, ctx)?; + let mut out = Vec::with_capacity(chunks.len()); + out.push(compressed); + for chunk in rest { + out.push(self.compress_with_plan(chunk, &plan, ctx)?); + } + Ok(out) + } } /// A layout writer that compresses chunks. @@ -59,6 +114,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 +127,20 @@ 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. + 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 +171,36 @@ impl LayoutStrategy for CompressingStrategy { let compute_session = session.clone(); let handle = session.handle(); - let stream = stream - .map(move |chunk| { + let stream = group_chunks(stream, self.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 +213,160 @@ 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::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 super::*; + use crate::sequence::SequentialArrayStreamExt; + + 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)); + } +} diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 7344ee7be51..8a4b1a245f9 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -13,6 +13,8 @@ use vortex_array::Canonical; 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 +82,7 @@ impl RepartitionWriterOptions { pub struct RepartitionStrategy { child: Arc, options: RepartitionWriterOptions, + block_size_maximum: Option, } impl RepartitionStrategy { @@ -87,8 +90,51 @@ 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 number of rows per block that keeps each block of `chunk` at or under `maximum` bytes, +/// as a power of two no larger than `max_rows`, or `None` when `chunk` needs no split. +/// +/// 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_rows(chunk: &ArrayRef, maximum: u64, max_rows: usize) -> Option { + let views = chunk.as_opt::()?; + let len = chunk.len(); + if len <= 1 { + return None; + } + let referenced: u64 = views + .views() + .iter() + .map(|view| u64::from(view.len())) + .sum::() + + (len * size_of::()) as u64; + if referenced <= maximum { + return None; + } + let fit = usize::try_from(u128::from(maximum) * len as u128 / u128::from(referenced)) + .unwrap_or(usize::MAX) + .clamp(1, max_rows.max(1)); + let rows = 1_usize << (usize::BITS - 1 - fit.leading_zeros()); + (rows < len).then_some(rows) } #[async_trait] @@ -128,6 +174,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 +186,26 @@ 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(rows) = block_size_maximum + .and_then(|maximum| split_rows(&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())?; + yield ( + sequence_pointer.advance(), + pending.into_array().execute::(&mut ctx)?.into_array(), + ); + } + let mut offset = 0; + while offset < chunk.len() { + let end = (offset + rows).min(chunk.len()); + yield (sequence_pointer.advance(), chunk.slice(offset..end)?); + offset = end; + } + continue; + } let mut offset = 0; while offset < chunk.len() { let end = (offset + block_len).min(chunk.len()); @@ -252,6 +319,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 +348,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 +601,88 @@ 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 + })?; + (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(()) + } + + #[test] + fn split_rows_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_rows(&slice, ONE_MEG, 8192), + None, + "64 rows reference ~64 KiB" + ); + assert_eq!(split_rows(&wide, ONE_MEG, 8192), Some(1024)); + // A maximum smaller than one row still yields blocks of one row. + assert_eq!(split_rows(&wide, 100, 8192), Some(1)); + // Non-string chunks are never split. + let ints = PrimitiveArray::from_iter(0..8192_i64).into_array(); + assert_eq!(split_rows(&ints, 1, 8192), None); + } } From df11de53682a98a60237be8bc543886a00dc4e47 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:26:06 -0700 Subject: [PATCH 02/13] perf(layout): skip the chunks a known mask excludes when reading a row range `ChunkedReader` evaluated a filter or projection on every chunk overlapping the row range, and the flat reader requests its segment before awaiting the mask, so every chunk in range was read even when the mask selected none of its rows. A point read passes an already-resolved mask over a whole 8,192-row range; once a string column's blocks are split under a byte cap, reading one row of that range fetched every split block, which made the lookup slower than before the split. `MaskFuture::known` returns the mask when the future was created from one (`ready`, `new_true`, or a slice of either). `ChunkedReader` skips a chunk whose known mask slice selects no rows: its projection contributes nothing and its filter result is all false, so neither its segment nor its children are requested. A mask that is still pending is evaluated as before, so scans keep issuing segment requests ahead of the mask. --- SPICE_PATCHES.md | 1 + vortex-array/src/mask_future.rs | 46 ++++++++- vortex-layout/src/layouts/chunked/reader.rs | 109 +++++++++++++++++++- 3 files changed, 153 insertions(+), 3 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index ba952f79354..5d1e31deeff 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -61,6 +61,7 @@ exists: a grep proves presence, a test proves behaviour. | 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::compress_recording`/`compress_with_plan` choose encodings once per row block and replay them on its split blocks | `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-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..c53ee0a2467 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,18 @@ 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. + /// + /// 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 +66,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 +85,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)) }) } @@ -84,6 +102,7 @@ impl MaskFuture { Self { inner: self.inner.inspect(f).boxed().shared(), len, + known: self.known, } } } @@ -124,3 +143,28 @@ mod tests { }) } } + +#[cfg(test)] +mod known_tests { + 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()); + } +} 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); + }) + } } From 1fcbdc98d98d04b39c29b69d2a5e39f37edb539e Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Thu, 22 Oct 2026 03:28:52 -0700 Subject: [PATCH 03/13] fix(layout): choose a row block's encodings on a sample of the whole block, not its first chunk Replaying the schemes chosen on a block's first split chunk encoded the rest to suit that chunk: with incompressible rows ahead of JSON, the JSON got FSST instead of OnPair, a 15% larger file and 2.4x the LIKE CPU. CascadingCompressor::plan_for chooses on the stratified sample compress() would draw from the unsplit block, so each split chunk is encoded as the block would have been. --- SPICE_PATCHES.md | 2 +- vortex-compressor/src/compressor/cascade.rs | 23 +++++ vortex-layout/src/layouts/compressed.rs | 102 +++++++++++++++++--- 3 files changed, 115 insertions(+), 12 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 5d1e31deeff..5b8e38c6fc2 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -60,7 +60,7 @@ 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::compress_recording`/`compress_with_plan` choose encodings once per row block and replay them on its split blocks | `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-file --test test_write_table test_data_block_max_bytes_roundtrip` | Proposed — additive; off unless the cap is set | +| 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 diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 475fa14345f..632f73e1eed 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -88,6 +88,29 @@ impl CascadingCompressor { 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 diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 5e36bd4b75e..c14927b44a7 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -9,7 +9,9 @@ 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::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; @@ -81,29 +83,31 @@ impl CompressorPlugin for BtrBlocksCompressor { self.compress(chunk, ctx) } - /// Searches for encodings on the first chunk and replays them on the rest. + /// 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. Replaying keeps one search per group. + /// 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 Some((first, rest)) = chunks.split_first() else { + let [first, ..] = chunks else { return Ok(Vec::new()); }; - if rest.is_empty() { + if chunks.len() == 1 { return Ok(vec![self.compress(first, ctx)?]); } - let (compressed, plan) = self.compress_recording(first, ctx)?; - let mut out = Vec::with_capacity(chunks.len()); - out.push(compressed); - for chunk in rest { - out.push(self.compress_with_plan(chunk, &plan, ctx)?); - } - Ok(out) + 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() } } @@ -369,4 +373,80 @@ mod tests { ); 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(()) + } } From 877e96db66a657ae29637cbf1b06def2810143d7 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Thu, 22 Oct 2026 03:28:52 -0700 Subject: [PATCH 04/13] bench(file): read-path costs of masked scans and of byte-capped wide string blocks --- Cargo.lock | 1 + vortex-file/Cargo.toml | 11 + vortex-file/benches/scan_mask_paths.rs | 256 +++++++++++ vortex-file/benches/wide_string_cases.rs | 517 +++++++++++++++++++++++ 4 files changed, 785 insertions(+) create mode 100644 vortex-file/benches/scan_mask_paths.rs create mode 100644 vortex-file/benches/wide_string_cases.rs diff --git a/Cargo.lock b/Cargo.lock index 6971a32c4f0..3fdaaa5d45b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11056,6 +11056,7 @@ dependencies = [ "futures", "itertools 0.14.0", "kanal", + "libc", "moka", "object_store", "parking_lot", 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..2b18bd026e3 --- /dev/null +++ b/vortex-file/benches/scan_mask_paths.rs @@ -0,0 +1,256 @@ +// 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). + +#![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)] + +#[path = "../tests/common/mod.rs"] +mod common; + +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); + 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) +} + +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 + ); + } +} diff --git a/vortex-file/benches/wide_string_cases.rs b/vortex-file/benches/wide_string_cases.rs new file mode 100644 index 00000000000..7a991c1fc60 --- /dev/null +++ b/vortex-file/benches/wide_string_cases.rs @@ -0,0 +1,517 @@ +// 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). + +#![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")] + +#[path = "../tests/common/mod.rs"] +mod common; + +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); + 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) + ) +} + +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 + ); + } + } +} From ed8b9e422ddd43a396c3030f9703df6476441bcf Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:46:20 -0700 Subject: [PATCH 05/13] fix(compressor): choose afresh when a replayed scheme cannot encode the data A plan replays the schemes chosen on other data, and some schemes accept a dtype yet reject data they cannot encode: a sequence recorded on 0..n failed the write of random integers with 'cannot sequence encode array'. On that failure the cascade now chooses as if there were no plan; an error from that choice still fails the write. --- vortex-btrblocks/src/canonical_compressor.rs | 28 +++++++++++ vortex-compressor/src/compressor/cascade.rs | 50 ++++++++++++++++---- 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index ec8de8c0343..e95431a904a 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -396,4 +396,32 @@ mod tests { 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-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 632f73e1eed..487cdfb9c0c 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -374,6 +374,7 @@ impl CascadingCompressor { .map(|scheme| (scheme, super::select::WinnerEstimate::AlwaysUse)), None => None, }; + let from_plan = chosen.is_some(); let chosen = match chosen { Some(chosen) => Some(chosen), None => { @@ -396,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); From 7dda1ca4002d34362363844e4631e4cfb6d6e51a Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:46:20 -0700 Subject: [PATCH 06/13] fix(array): a slice of an inspected MaskFuture still runs its callback inspect kept the known mask, so a slice resolved from it without awaiting the inspected future and never ran the callback. An inspected future is no longer known. --- vortex-array/src/mask_future.rs | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/mask_future.rs b/vortex-array/src/mask_future.rs index c53ee0a2467..8fa7a41d21a 100644 --- a/vortex-array/src/mask_future.rs +++ b/vortex-array/src/mask_future.rs @@ -47,7 +47,8 @@ impl MaskFuture { } /// 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. + /// 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> { @@ -99,10 +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: self.known, + known: None, } } } @@ -146,6 +149,11 @@ 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; @@ -167,4 +175,18 @@ mod known_tests { 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(()) + }) + } } From 47f66d56aca95c867b86b3aa2f48188dcf24305a Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:46:20 -0700 Subject: [PATCH 07/13] fix(layout): group only the string and binary blocks a byte maximum splits Grouping applied to every chunk smaller than a row block, with or without a maximum, so the blocks of a wide fixed-size list (an embedding column) compressed serially in one task: same bytes and CPU, 3.5x the write wall time. Only Utf8/Binary chunks group, and only when with_data_block_max_bytes is set. --- vortex-file/src/strategy.rs | 13 ++-- vortex-layout/src/layouts/compressed.rs | 92 ++++++++++++++++++++++++- 2 files changed, 100 insertions(+), 5 deletions(-) diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index e934c93dee5..f087cf6c6e8 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -239,10 +239,15 @@ impl WriteStrategyBuilder { ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - // Choose encodings once per row block, even when a block is split into smaller data - // blocks (see `with_data_block_max_bytes`). - let compressing = CompressingStrategy::new(buffered, data_compressor) - .with_group_rows(self.row_block_size); + // 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( diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index c14927b44a7..d5792c64a04 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -12,6 +12,7 @@ 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; @@ -140,6 +141,11 @@ impl CompressingStrategy { /// /// 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 @@ -175,7 +181,12 @@ impl LayoutStrategy for CompressingStrategy { let compute_session = session.clone(); let handle = session.handle(); - let stream = group_chunks(stream, self.group_rows) + 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); @@ -283,6 +294,7 @@ fn group_chunks( #[cfg(test)] mod tests { + use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; @@ -292,9 +304,15 @@ mod tests { 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 @@ -449,4 +467,76 @@ mod tests { } 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(()) + } } From 66d7dbf358044a7ec565b11fab74f3cd2b6f63d9 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:46:20 -0700 Subject: [PATCH 08/13] fix(layout): halve a split block until every piece fits the byte maximum The rows per block came from the chunk's average row width, so wide rows clustered at one end left a block over the maximum (25.7 MB under 4 MiB). Each block over the maximum is now halved until its pieces fit or hold one row; uniform rows split as before. --- vortex-layout/src/layouts/repartition.rs | 110 +++++++++++++++++------ 1 file changed, 83 insertions(+), 27 deletions(-) diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 8a4b1a245f9..579d11d2f86 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; @@ -109,32 +110,54 @@ impl RepartitionStrategy { } } -/// The number of rows per block that keeps each block of `chunk` at or under `maximum` bytes, -/// as a power of two no larger than `max_rows`, or `None` when `chunk` needs no split. +/// 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_rows(chunk: &ArrayRef, maximum: u64, max_rows: usize) -> Option { +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 referenced: u64 = views - .views() - .iter() - .map(|view| u64::from(view.len())) - .sum::() - + (len * size_of::()) as u64; - if referenced <= maximum { + 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 fit = usize::try_from(u128::from(maximum) * len as u128 / u128::from(referenced)) - .unwrap_or(usize::MAX) - .clamp(1, max_rows.max(1)); - let rows = 1_usize << (usize::BITS - 1 - fit.leading_zeros()); - (rows < len).then_some(rows) + + 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) } #[async_trait] @@ -186,8 +209,8 @@ 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(rows) = block_size_maximum - .and_then(|maximum| split_rows(&chunk, maximum, block_len)) + 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(); @@ -198,11 +221,8 @@ impl LayoutStrategy for RepartitionStrategy { pending.into_array().execute::(&mut ctx)?.into_array(), ); } - let mut offset = 0; - while offset < chunk.len() { - let end = (offset + rows).min(chunk.len()); - yield (sequence_pointer.advance(), chunk.slice(offset..end)?); - offset = end; + for range in ranges { + yield (sequence_pointer.advance(), chunk.slice(range)?); } continue; } @@ -668,21 +688,57 @@ mod tests { Ok(()) } + fn lengths(ranges: Option>>) -> Option> { + ranges.map(|ranges| ranges.iter().map(ExactSizeIterator::len).collect()) + } + #[test] - fn split_rows_measures_referenced_bytes_not_buffers() { + 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_rows(&slice, ONE_MEG, 8192), + split_ranges(&slice, ONE_MEG, 8192), None, "64 rows reference ~64 KiB" ); - assert_eq!(split_rows(&wide, ONE_MEG, 8192), Some(1024)); + 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!(split_rows(&wide, 100, 8192), Some(1)); + 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_rows(&ints, 1, 8192), None); + 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(()) } } From c036875b592a7ffc393bce0c9668e589839373f1 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:46:20 -0700 Subject: [PATCH 09/13] bench: build the CPU-timing benches only on Unix getrusage and clock_gettime are Unix APIs; elsewhere the benches compile to a stub that says so. --- .../benches/small_chunk_compress.rs | 520 ++++++----- vortex-file/benches/scan_mask_paths.rs | 459 ++++----- vortex-file/benches/wide_string_cases.rs | 884 +++++++++--------- 3 files changed, 953 insertions(+), 910 deletions(-) diff --git a/vortex-btrblocks/benches/small_chunk_compress.rs b/vortex-btrblocks/benches/small_chunk_compress.rs index 187af0fd618..8f0875c3064 100644 --- a/vortex-btrblocks/benches/small_chunk_compress.rs +++ b/vortex-btrblocks/benches/small_chunk_compress.rs @@ -8,297 +8,311 @@ //! //! `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. -#![expect(clippy::unwrap_used)] -#![expect(clippy::cast_precision_loss)] +#[cfg(unix)] +mod unix { + #![expect(clippy::unwrap_used)] + #![expect(clippy::cast_precision_loss)] -use std::sync::LazyLock; + 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; + 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); + 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 + /// 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(','); + /// 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 + )); } - 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 } - 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; + 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 } - 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 -} + /// 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; + /// 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; + 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:?}]")); - } + // 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()) }; } - } - 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)); + 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:?}]")); + } } } - 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; + + 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)); } - }); - 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; + } + 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())) + 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); - let mut plan = from_sample.then(|| { - // A 1,024-row stratified sample (16 x 64 rows) of the block. - let stride = (block_end - block) / 16; - let slices: Vec = (0..16) - .map(|slice| { - let start = block + slice * stride; - array.slice(start..(start + 64).min(block_end)).unwrap() - }) - .collect(); - let sample = vortex_array::arrays::ChunkedArray::try_new(slices, array.dtype().clone()) - .unwrap() - .into_array(); - compressor.compress_recording(&sample, &mut ctx).unwrap().1 - }); - 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(); + 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); + let mut plan = from_sample.then(|| { + // A 1,024-row stratified sample (16 x 64 rows) of the block. + let stride = (block_end - block) / 16; + let slices: Vec = (0..16) + .map(|slice| { + let start = block + slice * stride; + array.slice(start..(start + 64).min(block_end)).unwrap() + }) + .collect(); + let sample = + vortex_array::arrays::ChunkedArray::try_new(slices, array.dtype().clone()) + .unwrap() + .into_array(); + compressor.compress_recording(&sample, &mut ctx).unwrap().1 + }); + 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 } - bytes -} -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(); + 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(); + 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 P2 8k sample, {chunk_rows} rows"), - Box::new(move || compress_planned(full, array, rows, chunk_rows, true)), - )); - } + 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 P2 8k sample, {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; + // 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) + "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" ); - 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) { + for (i, (label, _)) in arms.iter().enumerate() { + let t = &mut times[i]; + t.sort_by(f64::total_cmp); println!( - " {:>9.1} ms/rep {:>8} calls/rep {key}", - ms / reps as f64, - *n / reps as u64 + "{:<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-file/benches/scan_mask_paths.rs b/vortex-file/benches/scan_mask_paths.rs index 2b18bd026e3..605c35206b3 100644 --- a/vortex-file/benches/scan_mask_paths.rs +++ b/vortex-file/benches/scan_mask_paths.rs @@ -6,251 +6,266 @@ //! //! `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. -#![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)] - +#[cfg(unix)] #[path = "../tests/common/mod.rs"] mod common; -use std::sync::LazyLock; -use std::time::Instant; +#[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 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; + use std::sync::LazyLock; + use std::time::Instant; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session() - .with::() - .with::(); - vortex_file::register_default_encodings(&session); - common::enable_all_registered_array_encodings(&session); - session -}); + 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; -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) -} + 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 + }); -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 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) } -} - -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()); + 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 + } } - if let Some(s) = selection { - builder = builder.with_selection(s); + + 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() } - 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) + + async fn scan( + file: &VortexFile, + projection: Expression, + filter: Option, + selection: Option, + ) -> (usize, u64) { + let dtype = file.dtype().clone(); + let mut builder = file + .scan() .unwrap() - .into_array(); - rows += chunk.len(); - bytes += chunk.nbytes(); + .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) } - (rows, bytes) -} -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 + 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 digest = bytes.iter().fold(0xcbf2_9ce4_8422_2325_u64, |h, b| { - (h ^ u64::from(*b)).wrapping_mul(0x100_0000_01b3) + 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() }); - 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( + // 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(["s"], root()), + select(["a", "b", "s"], root()), + Some(and( + gt(get_item("k", root()), lit(lo)), + lt(get_item("k", root()), lit(hi)), + )), None, - Some(Selection::IncludeByIndex( - StrictSortedBuffer::try_new(Buffer::from_iter([row])).unwrap(), + )) + }), + ), + ( + "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(), )), - )); - 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); + )) + }), + ), + ( + "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 + "{:<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 index 7a991c1fc60..d85976bb103 100644 --- a/vortex-file/benches/wide_string_cases.rs +++ b/vortex-file/benches/wide_string_cases.rs @@ -20,498 +20,512 @@ //! `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. -#![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")] - +#[cfg(unix)] #[path = "../tests/common/mod.rs"] mod common; -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); - 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) -} +#[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) -} + 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 + 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 - )); + 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 } - s.push_str("]}"); - s -} -struct Dataset { - name: &'static str, - rows: Vec, -} + 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) + 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 low = if name == "shift_lo_hi" { + let odd = if name == "rand_then_json" { in_block < 1_024 } else { in_block >= 7_168 }; - let s = if low { - fixed[i % 8].clone() + let s = if odd { + (0..188) + .map(|_| format!("{:016x}", rng.next())) + .collect::() } 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, -} + .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 } + } -impl CountingReadAt { - fn take(&self) -> (u64, u64) { - ( - self.requests.swap(0, Ordering::Relaxed), - self.bytes.swap(0, Ordering::Relaxed), + 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() } -} -impl VortexReadAt for CountingReadAt { - fn coalesce_config(&self) -> Option { - Some(CoalesceConfig::object_storage()) + /// 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, } - fn concurrency(&self) -> usize { - 16 + impl CountingReadAt { + fn take(&self) -> (u64, u64) { + ( + self.requests.swap(0, Ordering::Relaxed), + self.bytes.swap(0, Ordering::Relaxed), + ) + } } - fn size(&self) -> BoxFuture<'static, VortexResult> { - let len = self.inner.len() as u64; - async move { Ok(len) }.boxed() - } + impl VortexReadAt for CountingReadAt { + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig::object_storage()) + } - 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) - } -} + fn concurrency(&self) -> usize { + 16 + } -struct Arm { - label: &'static str, - file: VortexFile, - reader: CountingReadAt, -} + fn size(&self) -> BoxFuture<'static, VortexResult> { + let len = self.inner.len() as u64; + async move { Ok(len) }.boxed() + } -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 -} + 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) + } + } -/// 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()); + struct Arm { + label: &'static str, + file: VortexFile, + reader: CountingReadAt, } - if let Some(s) = selection { - builder = builder.with_selection(s); + + 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 } - 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); + + /// 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); + 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) } - (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 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) - ) -} + 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) + ) + } -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() { + 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!( - "{: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() + "\n### {} rows={n} raw={:.1}MiB", + ds.name, + raw as f64 / 1048576.0 ); - } - 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); - } + // 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); } - println!(" {label} (rows {})", seen[0].0.0); - for i in 0..2 { + 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} 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 + "{: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() ); } - } - // 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 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 ); - 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 - ); + + // 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"); +} From 6e0a466b7d822f01964cec57dfdfaf6a2231e162 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:37:24 -0700 Subject: [PATCH 10/13] fix(layout): keep coalesced blocks under the byte maximum The maximum was checked only on incoming chunks, so coalescing could merge blocks that each fit into one that did not: a 0.9 MiB block and a 3.9 MiB one became a 5.0 MB block under a 4 MiB maximum. Every block coalescing emits is now split by the same rule. --- vortex-layout/src/layouts/repartition.rs | 74 ++++++++++++++++++++---- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 579d11d2f86..8053e7c8164 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -160,6 +160,19 @@ fn split_ranges(chunk: &ArrayRef, maximum: u64, max_rows: usize) -> Option 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, +) -> VortexResult> { + match maximum.and_then(|maximum| split_ranges(&block, maximum, max_rows)) { + Some(ranges) => ranges.into_iter().map(|range| block.slice(range)).collect(), + None => Ok(vec![block]), + } +} + #[async_trait] impl LayoutStrategy for RepartitionStrategy { async fn write_stream( @@ -216,10 +229,10 @@ impl LayoutStrategy for RepartitionStrategy { let pending = chunks.take_all(); if !pending.is_empty() { let pending = ChunkedArray::try_new(pending, dtype_clone.clone())?; - yield ( - sequence_pointer.advance(), - pending.into_array().execute::(&mut ctx)?.into_array(), - ); + let pending = pending.into_array().execute::(&mut ctx)?.into_array(); + for block in under_maximum(pending, block_size_maximum, block_len)? { + yield (sequence_pointer.advance(), block); + } } for range in ranges { yield (sequence_pointer.advance(), chunk.slice(range)?); @@ -240,10 +253,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)? { + yield (sequence_pointer.advance(), block); + } } } } @@ -254,10 +266,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)? { + yield (sequence_pointer.advance(), block); + } } } } @@ -645,6 +656,10 @@ mod tests { .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 @@ -741,4 +756,39 @@ mod tests { } 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(()) + } } From 00732f6b72f024d43cb426e988ce54097987ab8a Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:37:24 -0700 Subject: [PATCH 11/13] test(file): split blocks of incompressible binary store only their rows --- vortex-file/tests/test_write_table.rs | 58 +++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index b6846d950a7..9b1f7c386f1 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -274,3 +274,61 @@ async fn test_data_block_max_bytes_roundtrip() -> vortex_error::VortexResult<()> ); Ok(()) } + +/// Splitting a block must not copy its whole data buffer into every piece: when no scheme +/// shrinks the bytes (random binary), each piece's segment holds only the rows it covers. +#[tokio::test] +async fn split_blocks_of_incompressible_binary_store_only_their_rows() +-> vortex_error::VortexResult<()> { + let rows = 8192; + let mut state = 0x9e37_79b9_7f4a_7c15_u64; + let values: Vec> = (0..rows) + .map(|_| { + (0..1000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state >> 56) as u8 + }) + .collect() + }) + .collect(); + let data = + vortex_array::arrays::VarBinViewArray::from_iter_bin(values.iter().map(Vec::as_slice)) + .into_array(); + + let mut sizes = Vec::new(); + for max in [None, Some(4 << 20)] { + let strategy = vortex_file::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?; + 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, &mut SESSION.create_execution_ctx()); + sizes.push((file_len, largest)); + } + let ((plain, plain_largest), (split, split_largest)) = (sizes[0], sizes[1]); + assert!(split <= plain + plain / 50, "split file {split} vs {plain}"); + assert!( + split_largest <= 4 << 20, + "largest split segment {split_largest} (unsplit {plain_largest})" + ); + Ok(()) +} From dce56c6ea8b95eebb92633fd65ae6ea271ff1002 Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:12:33 -0700 Subject: [PATCH 12/13] fix(layout): a split piece holds only the bytes its rows reference A split piece was a slice sharing its parent's data buffers, and a buffer is written whole, so with a compressor that stores chunks as they are each piece carried the whole block: random binary split in two wrote 16.5 MB instead of 8.3 MB. Each piece is now compacted at threshold 1.0, which slices a contiguously used buffer rather than copying it. The regression covers BtrBlocks and a pass-through compressor, for a block split as it arrives and one rebuilt by coalescing. --- vortex-file/tests/test_write_table.rs | 115 ++++++++++++++--------- vortex-layout/src/layouts/repartition.rs | 35 ++++++- 2 files changed, 102 insertions(+), 48 deletions(-) diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index 9b1f7c386f1..788bb4d222c 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -275,16 +275,12 @@ async fn test_data_block_max_bytes_roundtrip() -> vortex_error::VortexResult<()> Ok(()) } -/// Splitting a block must not copy its whole data buffer into every piece: when no scheme -/// shrinks the bytes (random binary), each piece's segment holds only the rows it covers. -#[tokio::test] -async fn split_blocks_of_incompressible_binary_store_only_their_rows() --> vortex_error::VortexResult<()> { - let rows = 8192; - let mut state = 0x9e37_79b9_7f4a_7c15_u64; +/// 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..1000) + (0..width) .map(|_| { state ^= state << 13; state ^= state >> 7; @@ -294,41 +290,74 @@ async fn split_blocks_of_incompressible_binary_store_only_their_rows() .collect() }) .collect(); - let data = - vortex_array::arrays::VarBinViewArray::from_iter_bin(values.iter().map(Vec::as_slice)) - .into_array(); - - let mut sizes = Vec::new(); - for max in [None, Some(4 << 20)] { - let strategy = vortex_file::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?; - 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, &mut SESSION.create_execution_ctx()); - sizes.push((file_len, largest)); + 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" + ); + } } - let ((plain, plain_largest), (split, split_largest)) = (sizes[0], sizes[1]); - assert!(split <= plain + plain / 50, "split file {split} vs {plain}"); - assert!( - split_largest <= 4 << 20, - "largest split segment {split_largest} (unsplit {plain_largest})" - ); Ok(()) } diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 8053e7c8164..4147cce9f0d 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -11,6 +11,7 @@ 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; @@ -166,13 +167,37 @@ 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| block.slice(range)).collect(), + 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] impl LayoutStrategy for RepartitionStrategy { async fn write_stream( @@ -230,12 +255,12 @@ impl LayoutStrategy for RepartitionStrategy { 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)? { + 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(), chunk.slice(range)?); + yield (sequence_pointer.advance(), split_piece(&chunk, range, &mut ctx)?); } continue; } @@ -253,7 +278,7 @@ 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(); - for block in under_maximum(canonical, block_size_maximum, block_len)? { + for block in under_maximum(canonical, block_size_maximum, block_len, &mut ctx)? { yield (sequence_pointer.advance(), block); } } @@ -266,7 +291,7 @@ impl LayoutStrategy for RepartitionStrategy { )?; if !to_flush.is_empty() { let canonical = to_flush.into_array().execute::(&mut ctx)?.into_array(); - for block in under_maximum(canonical, block_size_maximum, block_len)? { + for block in under_maximum(canonical, block_size_maximum, block_len, &mut ctx)? { yield (sequence_pointer.advance(), block); } } From 51e7a4c9cfe34128009c4f5d11ee96dbaca884af Mon Sep 17 00:00:00 2001 From: Ben Chambers <35960+bjchambers@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:56:06 -0700 Subject: [PATCH 13/13] bench(btrblocks): plan the sampled arm with plan_for, as the writer does The arm drew its own fixed-position sample, unlike plan_for's seeded stratified one, so its encodings and planning cost could differ from production. --- .../benches/small_chunk_compress.rs | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/vortex-btrblocks/benches/small_chunk_compress.rs b/vortex-btrblocks/benches/small_chunk_compress.rs index 8f0875c3064..9ea2b713489 100644 --- a/vortex-btrblocks/benches/small_chunk_compress.rs +++ b/vortex-btrblocks/benches/small_chunk_compress.rs @@ -179,20 +179,11 @@ mod unix { 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(|| { - // A 1,024-row stratified sample (16 x 64 rows) of the block. - let stride = (block_end - block) / 16; - let slices: Vec = (0..16) - .map(|slice| { - let start = block + slice * stride; - array.slice(start..(start + 64).min(block_end)).unwrap() - }) - .collect(); - let sample = - vortex_array::arrays::ChunkedArray::try_new(slices, array.dtype().clone()) - .unwrap() - .into_array(); - compressor.compress_recording(&sample, &mut ctx).unwrap().1 + 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); @@ -255,7 +246,7 @@ mod unix { Box::new(move || compress_planned(full, array, rows, chunk_rows, false)), )); arms.push(( - format!("plan P2 8k sample, {chunk_rows} rows"), + format!("plan_for on the whole block, {chunk_rows} rows"), Box::new(move || compress_planned(full, array, rows, chunk_rows, true)), )); }