From 0d08927d0cc8e4262da2386017685a135968bef4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 19:48:30 +0000 Subject: [PATCH 01/28] Spread filtered V2 scans across threads and stop concatenating projection pieces Filtered V2 scans used fixed 1,048,576-row chunks. A DataFusion byte-range partition of a ~1M-row file then became a single chunk, so when the surviving rows sat in one partition (PolarSignals time ranges) one thread did all the work. Size filtered chunks by the serialized bytes the projection reads per row and by the partition's rows per thread, keeping 1M-row chunks for cheap projections over large ranges. ProjectionMorsel also wrapped its pieces in a ChunkedArray, which the Arrow export concatenated, re-merging nested dictionaries. Emit pieces as separate batches, grouping only pieces under 4,096 rows. PolarSignals SF1 on DataFusion, V2 vs V1 geomean: 1.79x -> 0.93x. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- vortex-layout/src/scan/execution/chunk.rs | 39 ++++++++-- vortex-layout/src/scan/execution/mod.rs | 3 +- .../src/scan/execution/projection.rs | 75 +++++++++++++------ vortex-layout/src/scan/v2/cache.rs | 19 +++++ vortex-layout/src/scan/v2/mod.rs | 38 ++++++++-- 5 files changed, 141 insertions(+), 33 deletions(-) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 6808c707c48..ca4166cfb3b 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/vortex-layout/src/scan/execution/chunk.rs @@ -7,6 +7,7 @@ use std::ops::Range; use std::sync::Arc; use std::vec::IntoIter; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::request::IoConsumer; use vortex_io::request::IoRequestId; @@ -35,13 +36,41 @@ pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> u if row_count <= MAX_ROWS_PER_CHUNK as u64 { return MAX_ROWS_PER_CHUNK; } - let bytes: u64 = locations + let blocks = blocks_per_target_bytes(total_bytes(locations), row_count); + blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE +} + +/// Rows per chunk of a filtered scan over `rows` rows of the file. +/// +/// Large chunks amortize filter evaluation over cheap projections. A projection that reads many +/// bytes per row, or a row range that would otherwise give fewer chunks than threads, gets smaller +/// chunks so the surviving rows still spread across the runtime's threads. +pub(crate) fn filtered_rows_per_chunk( + rows: u64, + row_count: u64, + projected_bytes: u64, + parallelism: usize, +) -> usize { + let max_blocks = (FILTERED_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64; + let min_blocks = (MAX_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64; + let by_bytes = blocks_per_target_bytes(projected_bytes, row_count).max(1); + let by_threads = (rows / parallelism.max(1) as u64 / ROW_BLOCK_SIZE as u64).max(min_blocks); + let blocks = by_bytes.min(by_threads).min(max_blocks); + usize::try_from(blocks).vortex_expect("at most FILTERED_ROWS_PER_CHUNK blocks") * ROW_BLOCK_SIZE +} + +/// The row blocks that hold [`TARGET_BYTES_PER_CHUNK`] when `bytes` spread over `row_count` rows. +fn blocks_per_target_bytes(bytes: u64, row_count: u64) -> u64 { + let bytes_per_row = bytes.div_ceil(row_count.max(1)).max(1); + TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64 +} + +/// The serialized bytes of every segment in `locations`. +pub(crate) fn total_bytes(locations: &[SegmentLocation]) -> u64 { + locations .iter() .map(|location| u64::from(location.length)) - .sum(); - let bytes_per_row = bytes.div_ceil(row_count).max(1); - let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64; - blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE + .sum() } /// Fixed boundaries in the file's row domain, clipped to the requested range. diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index e371fdb539c..80b6c38ba34 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -30,12 +30,13 @@ use vortex_scan::selection::Selection; use vortex_session::VortexSession; pub(in crate::scan) use self::chunk::ChunkPlanner; -pub(crate) use self::chunk::FILTERED_ROWS_PER_CHUNK; pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK; pub(crate) use self::chunk::ROW_BLOCK_SIZE; +pub(crate) use self::chunk::filtered_rows_per_chunk; #[cfg(test)] pub(crate) use self::chunk::row_chunks; pub(crate) use self::chunk::rows_per_chunk; +pub(crate) use self::chunk::total_bytes; use self::filter::FilterPlanner; pub(super) use self::filter::FilterPlans; use self::projection::ProjectionPlanner; diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 141b7377c26..50eb62d773c 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; use std::sync::Arc; use std::vec::IntoIter; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; @@ -30,6 +32,8 @@ use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; const MAX_OUTPUT_ROWS: usize = 65_536; +/// Pieces smaller than this are grouped with their neighbours into one batch. +const MIN_OUTPUT_ROWS: usize = 4_096; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when /// converting chunked projection results to Arrow. @@ -97,15 +101,19 @@ impl Planner for ProjectionPlanner { } } -/// Runs the projection plan over a split's selected rows and returns them as one batch, in row -/// order. A selection that projects to no rows finishes without a batch. +/// Runs the projection plan over a split's selected rows and returns them in row order, one +/// batch per piece the scan produced. Pieces under [`MIN_OUTPUT_ROWS`] are grouped with their +/// neighbours. Emitting pieces separately avoids concatenating them on conversion to Arrow, +/// which rebuilds nested dictionaries. A selection that projects to no rows finishes without a +/// batch. struct ProjectionMorsel { scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, pieces: Vec, - done: bool, + /// Batches ready to emit once the scan has finished, in row order. + batches: Option>, } impl ProjectionMorsel { @@ -117,24 +125,39 @@ impl ProjectionMorsel { selected, adaptor: None, pieces: Vec::new(), - done: false, + batches: None, } } - fn finish(&mut self) -> VortexResult { - self.done = true; + /// Groups the finished pieces into batches of at least [`MIN_OUTPUT_ROWS`] rows where + /// possible. + fn group_pieces(&mut self) -> VortexResult> { self.pieces.sort_by_key(|piece| piece.rows.start); - let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) - .into_iter() - .map(|piece| piece.array) - .filter(|array| !array.is_empty()) - .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) + let mut batches = VecDeque::new(); + let mut group = Vec::new(); + let mut group_rows = 0; + for piece in std::mem::take(&mut self.pieces) { + if piece.array.is_empty() { + continue; + } + group_rows += piece.array.len(); + group.push(piece.array); + if group_rows >= MIN_OUTPUT_ROWS { + batches.push_back(self.concat(std::mem::take(&mut group))?); + group_rows = 0; + } + } + if !group.is_empty() { + batches.push_back(self.concat(group)?); + } + Ok(batches) + } + + fn concat(&self, mut arrays: Vec) -> VortexResult { + if arrays.len() == 1 { + return Ok(arrays.remove(0)); + } + Ok(ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array()) } } @@ -148,8 +171,14 @@ impl IoConsumer for ProjectionMorsel { impl Morsel for ProjectionMorsel { fn state(&self) -> State { + if let Some(batches) = &self.batches { + return if batches.is_empty() { + State::Done + } else { + State::NeedsCompute + }; + } match &self.adaptor { - _ if self.done => State::Done, None => State::NeedsCompute, Some(adaptor) => match adaptor.state() { State::Done => State::NeedsCompute, @@ -159,8 +188,11 @@ impl Morsel for ProjectionMorsel { } fn compute(&mut self) -> VortexResult { - if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); + if let Some(batches) = &mut self.batches { + return match batches.pop_front() { + Some(batch) => Ok(MorselOutput::Batch(batch)), + None => vortex_bail!("ProjectionMorsel: compute called after Done"), + }; } let Some(adaptor) = self.adaptor.as_mut() else { let scan = Scan::try_new( @@ -181,7 +213,8 @@ impl Morsel for ProjectionMorsel { return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { - return self.finish(); + self.batches = Some(self.group_pieces()?); + return Ok(MorselOutput::Continue); } Ok(match adaptor.compute()? { AdaptorStep::Yield => MorselOutput::Continue, diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 5a26e28ec97..21efeb3ac3f 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -10,6 +10,7 @@ use parking_lot::Mutex; use rustc_hash::FxHashMap; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_error::VortexResult; use super::io::SegmentRanges; use crate::plan::physical::PlanRef; @@ -43,6 +44,8 @@ struct Inner { /// Plans of filters, keyed by the exact filter expression. A proof plan remembers the zones it /// excluded, so sharing the plans proves a file's zones once for every scan of that filter. filters: Mutex>>, + /// Serialized bytes the whole-file projection reads, by projection. + projection_bytes: Mutex>, values: ValueCache, } @@ -112,6 +115,22 @@ impl ScanCache { ) } + /// The serialized bytes a whole-file scan of `projection` reads, computing them with + /// `compute` on first use. + pub(super) fn projection_bytes( + &self, + projection: &BoundExpression, + compute: impl FnOnce() -> VortexResult, + ) -> VortexResult { + let key = ExactBoundExpr(projection.clone()); + if let Some(bytes) = self.0.projection_bytes.lock().get(&key) { + return Ok(*bytes); + } + let bytes = compute()?; + self.0.projection_bytes.lock().insert(key, bytes); + Ok(bytes) + } + pub(super) fn values(&self) -> ValueCache { self.0.values.clone() } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 5f8384fc36b..3cfc947abf4 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -11,8 +11,9 @@ //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`] //! directly and supports projection, filters, row ranges, and selections. Unsupported layouts return an //! error; limits are not part of this builder's API. Work is partitioned into fixed -//! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored -//! layout's chunk boundaries. +//! chunks of up to 65,536 rows, or up to 1,048,576 rows for filtered scans, independently of the +//! stored layout's chunk boundaries. Filtered scans use smaller chunks when the projection reads +//! many bytes per row or the row range would give fewer chunks than threads. //! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. //! //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and proved @@ -51,6 +52,7 @@ use vortex_io::request::IoService; use vortex_io::request::IoSource; use vortex_io::runtime::Handle; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_scan::planning::driver::Progress; use vortex_scan::planning::driver::Run; use vortex_scan::planning::planner::Planner; @@ -65,15 +67,17 @@ pub use self::cache::WeakScanCache; use crate::LayoutRef; use crate::plan::physical::lower; use crate::plan::physical::optimize; +use crate::plan::physical::pipeline::Scan; use crate::plan::physical::plan_row_idx_expression; use crate::scan::arrow::to_record_batch; -use crate::scan::execution::FILTERED_ROWS_PER_CHUNK; use crate::scan::execution::FilterPlans; use crate::scan::execution::MAX_ROWS_PER_CHUNK; use crate::scan::execution::PreparedScan; use crate::scan::execution::ROW_BLOCK_SIZE; pub use crate::scan::execution::SegmentLocation; +use crate::scan::execution::filtered_rows_per_chunk; use crate::scan::execution::rows_per_chunk; +use crate::scan::execution::total_bytes; use crate::scan::v2::io::SegmentScanIo; use crate::scan::v2::io::segment_ranges; use crate::segments::SegmentSource; @@ -296,11 +300,33 @@ impl ScanRun { plan_row_idx_expression(builder.projection.clone(), root)?, &builder.session, )?; - cache.insert_projection(builder.projection, plan) + cache.insert_projection(builder.projection.clone(), plan) } }; + let rows = builder.row_range.unwrap_or(0..row_count); let chunk_rows = if filter.is_some() { - FILTERED_ROWS_PER_CHUNK + let projected_bytes = if identity { + total_bytes(&file.locations) + } else { + cache.projection_bytes(&builder.projection, || { + let segments = Scan::segments( + &builder.session, + &projection, + 0..row_count, + &Mask::new_true(usize::try_from(row_count)?), + )?; + Ok(segments + .iter() + .map(|segment| u64::from(file.locations[**segment as usize].length)) + .sum()) + })? + }; + filtered_rows_per_chunk( + rows.end - rows.start, + row_count, + projected_bytes, + get_available_parallelism().unwrap_or(1), + ) } else if identity { rows_per_chunk(row_count, &file.locations) } else { @@ -323,7 +349,7 @@ impl ScanRun { chunk_rows, }), io, - rows: builder.row_range.unwrap_or(0..row_count), + rows, selection: builder.selection, map_fn: builder.map_fn, ordered: builder.ordered, From d0ac505678195f51ba186c1c7724ad9a367f8c74 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 19:48:30 +0000 Subject: [PATCH 02/28] Dictionary-encode primitive Arrow exports without SipHash Exporting a non-dictionary Vortex array to an Arrow dictionary type went through arrow_cast, which hashes each primitive value with SipHash and builds an ArrowError per value. PolarSignals requests its nested `locations` fields as dictionaries, so this was ~90% of scan CPU for both V1 and V2. Encode primitive values directly, hashing their bits with FxHash and keeping nulls in the codes, and fall back to arrow_cast for other types or when the codes overflow. PolarSignals SF1 on DataFusion (V2): q00-q03 ~3.6x faster, total 4.28s -> 1.90s. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- Cargo.lock | 1 + vortex-arrow/Cargo.toml | 1 + vortex-arrow/src/executor/dictionary.rs | 140 +++++++++++++++++++++++- 3 files changed, 141 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 98826bda1d7..4a4e93d18c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11463,6 +11463,7 @@ dependencies = [ "mimalloc", "num-traits", "rstest", + "rustc-hash", "simdutf8", "tracing", "vortex-array", diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 7ee40444e84..3900814ccdc 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -27,6 +27,7 @@ arrow-schema = { workspace = true, features = ["canonical_extension_types"] } arrow-select = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } +rustc-hash = { workspace = true } simdutf8 = { workspace = true } tracing = { workspace = true } vortex-array = { workspace = true } diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 343c08357af..4a2922dd7e3 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -1,15 +1,22 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::hash_map::Entry; +use std::hash::Hash; use std::sync::Arc; +use arrow_array::Array; use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::ArrowPrimitiveType; use arrow_array::DictionaryArray; use arrow_array::PrimitiveArray; use arrow_array::cast::AsArray; use arrow_array::new_null_array; use arrow_array::types::*; +use arrow_buffer::ArrowNativeType; use arrow_schema::DataType; +use arrow_schema::TimeUnit; +use rustc_hash::FxHashMap; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -56,8 +63,11 @@ pub(super) fn to_arrow_dictionary( }; // Otherwise, we should try and build a dictionary. - // Arrow hides this functionality inside the cast module! let array = execute_arrow_naive(array, Some(values_type), exporter, ctx)?; + if let Some(dict) = primitive_to_dict(&array, codes_type) { + return Ok(dict); + } + // Arrow hides this functionality inside the cast module! arrow_cast::cast( &array, &DataType::Dictionary(Box::new(codes_type.clone()), Box::new(values_type.clone())), @@ -65,6 +75,98 @@ pub(super) fn to_arrow_dictionary( .map_err(VortexError::from) } +/// Dictionary-encodes a primitive Arrow array with codes of `codes_type`. +/// +/// Arrow's cast hashes every value with SipHash. This hashes the values' bits with FxHash instead, +/// which dominates exporting large nested columns that the reader requests as dictionaries. +/// Returns `None` for unsupported types and when the distinct values overflow the codes type, so +/// the caller can fall back to Arrow's cast. +fn primitive_to_dict(array: &ArrowArrayRef, codes_type: &DataType) -> Option { + match codes_type { + DataType::Int8 => primitive_to_dict_with::(array), + DataType::Int16 => primitive_to_dict_with::(array), + DataType::Int32 => primitive_to_dict_with::(array), + DataType::Int64 => primitive_to_dict_with::(array), + DataType::UInt8 => primitive_to_dict_with::(array), + DataType::UInt16 => primitive_to_dict_with::(array), + DataType::UInt32 => primitive_to_dict_with::(array), + DataType::UInt64 => primitive_to_dict_with::(array), + _ => None, + } +} + +fn primitive_to_dict_with( + array: &ArrowArrayRef, +) -> Option { + match array.data_type() { + DataType::Int8 => encode_primitive::(array, |v| v), + DataType::Int16 => encode_primitive::(array, |v| v), + DataType::Int32 => encode_primitive::(array, |v| v), + DataType::Int64 => encode_primitive::(array, |v| v), + DataType::UInt8 => encode_primitive::(array, |v| v), + DataType::UInt16 => encode_primitive::(array, |v| v), + DataType::UInt32 => encode_primitive::(array, |v| v), + DataType::UInt64 => encode_primitive::(array, |v| v), + DataType::Float16 => encode_primitive::(array, |v| v.to_bits()), + DataType::Float32 => encode_primitive::(array, |v| v.to_bits()), + DataType::Float64 => encode_primitive::(array, |v| v.to_bits()), + DataType::Date32 => encode_primitive::(array, |v| v), + DataType::Date64 => encode_primitive::(array, |v| v), + DataType::Timestamp(TimeUnit::Second, _) => { + encode_primitive::(array, |v| v) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + encode_primitive::(array, |v| v) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + encode_primitive::(array, |v| v) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + encode_primitive::(array, |v| v) + } + _ => None, + } +} + +/// Assigns codes in order of first appearance, keyed by `key` of each valid value. Null values +/// get null codes. +fn encode_primitive( + array: &ArrowArrayRef, + key: impl Fn(T::Native) -> B, +) -> Option +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Hash + Eq, +{ + let array = array.as_primitive::(); + let mut lookup: FxHashMap = FxHashMap::default(); + let mut values = Vec::new(); + let mut codes = Vec::with_capacity(array.len()); + for (idx, &value) in array.values().iter().enumerate() { + if array.is_null(idx) { + codes.push(K::Native::default()); + continue; + } + let code = match lookup.entry(key(value)) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + let code = K::Native::from_usize(values.len())?; + values.push(value); + *entry.insert(code) + } + }; + codes.push(code); + } + let codes = PrimitiveArray::::new(codes.into(), array.nulls().cloned()); + let values = + PrimitiveArray::::new(values.into(), None).with_data_type(array.data_type().clone()); + // SAFETY: every valid code indexes a value pushed above. + Some(Arc::new(unsafe { + DictionaryArray::new_unchecked(codes, Arc::new(values)) + })) +} + /// Convert a constant array to a dictionary with a single entry. fn constant_to_dict( array: ConstantArray, @@ -160,7 +262,13 @@ mod tests { use std::sync::LazyLock; use arrow_array::DictionaryArray as ArrowDictArray; + use arrow_array::Float64Array; + use arrow_array::Int16Array; + use arrow_array::Int64Array; use arrow_array::StringArray; + use arrow_array::UInt32Array; + use arrow_array::cast::AsArray; + use arrow_array::types::Int16Type; use arrow_array::types::UInt8Type; use arrow_array::types::UInt32Type; use arrow_schema::DataType; @@ -239,6 +347,22 @@ mod tests { dict_type(DataType::UInt8, DataType::Utf8), Arc::new(vec![Some("a"), None, Some("a"), Some("b"), Some("a")].into_iter().collect::>()) as arrow_array::ArrayRef, )] + #[case::primitive_fallback( + PrimitiveArray::from_option_iter([Some(7i64), None, Some(-3), Some(7), Some(-3)]).into_array(), + dict_type(DataType::UInt32, DataType::Int64), + Arc::new(ArrowDictArray::::new( + UInt32Array::from(vec![Some(0), None, Some(1), Some(0), Some(1)]), + Arc::new(Int64Array::from(vec![7, -3])), + )) as arrow_array::ArrayRef, + )] + #[case::float_fallback_by_bits( + buffer![0.0f64, -0.0, f64::NAN, 0.0, f64::NAN].into_array(), + dict_type(DataType::Int16, DataType::Float64), + Arc::new(ArrowDictArray::::new( + Int16Array::from(vec![0, 1, 2, 0, 2]), + Arc::new(Float64Array::from(vec![0.0, -0.0, f64::NAN])), + )) as arrow_array::ArrayRef, + )] fn to_arrow_dictionary( #[case] input: vortex_array::ArrayRef, #[case] target_type: DataType, @@ -249,6 +373,20 @@ mod tests { Ok(()) } + #[test] + fn primitive_fallback_fills_codes_type() -> VortexResult<()> { + let fits = PrimitiveArray::from_iter(0..=u8::MAX).into_array(); + let actual = execute(fits, &dict_type(DataType::UInt8, DataType::UInt8))?; + assert_eq!(actual.as_dictionary::().values().len(), 256); + Ok(()) + } + + #[test] + fn primitive_fallback_overflowing_codes_type_errors() { + let overflows = PrimitiveArray::from_iter(0..=256u16).into_array(); + assert!(execute(overflows, &dict_type(DataType::UInt8, DataType::UInt16)).is_err()); + } + #[test] fn mask_wrapped_dict_exports() -> VortexResult<()> { // Dictionary behind a lazy `mask` scalar-fn — the shape a scan produces when a row From c5b55fa3ca62d1e27323958c9d320601b33ed584 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 22:34:36 +0000 Subject: [PATCH 03/28] Use a dense table or the values themselves when exporting primitive dictionaries Hash-map growth and hashing still dominated exporting PolarSignals' nested dictionary fields. Keys spanning under 65,536 values (line numbers) are now looked up in a dense table. When a sample of the values is nearly all distinct (addresses), the values become the dictionary with sequential codes, reusing their buffer, since Arrow dictionaries may hold duplicates. Other inputs hash into a map sized from the sample. PolarSignals SF1 on DataFusion (V2): q00-q03 another 22-31% faster. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- vortex-arrow/src/executor/dictionary.rs | 133 +++++++++++++++++++++--- 1 file changed, 121 insertions(+), 12 deletions(-) diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 4a2922dd7e3..d2924a0f1a4 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -14,9 +14,11 @@ use arrow_array::cast::AsArray; use arrow_array::new_null_array; use arrow_array::types::*; use arrow_buffer::ArrowNativeType; +use arrow_buffer::ScalarBuffer; use arrow_schema::DataType; use arrow_schema::TimeUnit; use rustc_hash::FxHashMap; +use rustc_hash::FxHashSet; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -128,8 +130,17 @@ fn primitive_to_dict_with( } } -/// Assigns codes in order of first appearance, keyed by `key` of each valid value. Null values -/// get null codes. +/// Valid values sampled to choose how [`encode_primitive`] assigns codes. +const SAMPLE_LEN: usize = 1024; +/// The largest key range [`encode_primitive`] looks up in a dense table. +const MAX_DENSE_RANGE: usize = 1 << 16; + +/// Dictionary-encodes `array` keyed by `key` of each valid value. Null values get null codes. +/// +/// Keys spanning a small range are looked up in a dense table. When a sample of the values is +/// nearly all distinct, the values themselves become the dictionary with sequential codes, since +/// Arrow dictionaries may hold duplicates and deduplicating would barely shrink them. Otherwise +/// codes are assigned in order of first appearance through a hash map. fn encode_primitive( array: &ArrowArrayRef, key: impl Fn(T::Native) -> B, @@ -137,14 +148,106 @@ fn encode_primitive( where K: ArrowDictionaryKeyType, T: ArrowPrimitiveType, - B: Hash + Eq, + B: Hash + Eq + Copy + Into, { let array = array.as_primitive::(); - let mut lookup: FxHashMap = FxHashMap::default(); + let nulls = array.nulls().filter(|nulls| nulls.null_count() > 0); + let is_valid = |idx: usize| nulls.is_none_or(|nulls| nulls.is_valid(idx)); + let valid_keys = || { + array + .values() + .iter() + .enumerate() + .filter(move |(idx, _)| is_valid(*idx)) + .map(|(_, &value)| key(value)) + }; + + let (min, max) = valid_keys().fold((i128::MAX, i128::MIN), |(min, max), key| { + let key = key.into(); + (min.min(key), max.max(key)) + }); + let range = usize::try_from(max.checked_sub(min)?).ok(); + let (codes, values) = match range { + Some(range) if range < MAX_DENSE_RANGE && range <= array.len().saturating_mul(4) => { + encode_dense::(array, &is_valid, &key, min, range)? + } + _ => { + let mut sample = FxHashSet::default(); + let sampled = valid_keys() + .take(SAMPLE_LEN) + .inspect(|key| { + sample.insert(*key); + }) + .count(); + if sample.len() * 8 >= sampled * 7 && K::Native::from_usize(array.len()).is_some() { + let codes = (0..array.len()).map(K::Native::usize_as).collect(); + (codes, array.values().clone()) + } else { + let distinct = sample.len() * array.len() / sampled.max(1); + encode_hashed::(array, &is_valid, &key, distinct)? + } + } + }; + + let codes = PrimitiveArray::::new(codes, nulls.cloned()); + let values = PrimitiveArray::::new(values, None).with_data_type(array.data_type().clone()); + // SAFETY: every valid code indexes a value of the dictionary. + Some(Arc::new(unsafe { + DictionaryArray::new_unchecked(codes, Arc::new(values)) + })) +} + +/// Assigns codes in order of first appearance through a table indexed by `key - min`. +fn encode_dense( + array: &PrimitiveArray, + is_valid: &impl Fn(usize) -> bool, + key: &impl Fn(T::Native) -> B, + min: i128, + range: usize, +) -> Option<(ScalarBuffer, ScalarBuffer)> +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Into, +{ + const UNSEEN: u32 = u32::MAX; + let mut table = vec![UNSEEN; range + 1]; let mut values = Vec::new(); let mut codes = Vec::with_capacity(array.len()); for (idx, &value) in array.values().iter().enumerate() { - if array.is_null(idx) { + if !is_valid(idx) { + codes.push(K::Native::default()); + continue; + } + let slot = &mut table[usize::try_from(key(value).into() - min).ok()?]; + if *slot == UNSEEN { + K::Native::from_usize(values.len())?; + *slot = u32::try_from(values.len()).ok()?; + values.push(value); + } + codes.push(K::Native::usize_as(*slot as usize)); + } + Some((codes.into(), values.into())) +} + +/// Assigns codes in order of first appearance through a hash map sized for `distinct` keys. +fn encode_hashed( + array: &PrimitiveArray, + is_valid: &impl Fn(usize) -> bool, + key: &impl Fn(T::Native) -> B, + distinct: usize, +) -> Option<(ScalarBuffer, ScalarBuffer)> +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Hash + Eq, +{ + let mut lookup: FxHashMap = + FxHashMap::with_capacity_and_hasher(distinct, Default::default()); + let mut values = Vec::with_capacity(distinct); + let mut codes = Vec::with_capacity(array.len()); + for (idx, &value) in array.values().iter().enumerate() { + if !is_valid(idx) { codes.push(K::Native::default()); continue; } @@ -158,13 +261,7 @@ where }; codes.push(code); } - let codes = PrimitiveArray::::new(codes.into(), array.nulls().cloned()); - let values = - PrimitiveArray::::new(values.into(), None).with_data_type(array.data_type().clone()); - // SAFETY: every valid code indexes a value pushed above. - Some(Arc::new(unsafe { - DictionaryArray::new_unchecked(codes, Arc::new(values)) - })) + Some((codes.into(), values.into())) } /// Convert a constant array to a dictionary with a single entry. @@ -267,6 +364,7 @@ mod tests { use arrow_array::Int64Array; use arrow_array::StringArray; use arrow_array::UInt32Array; + use arrow_array::UInt64Array; use arrow_array::cast::AsArray; use arrow_array::types::Int16Type; use arrow_array::types::UInt8Type; @@ -373,6 +471,17 @@ mod tests { Ok(()) } + #[test] + fn primitive_fallback_distinct_values() -> VortexResult<()> { + let values = (0..2000u64).map(|i| (i % 7 != 3).then(|| i.wrapping_mul(0x9E37_79B9_7F4A_7C15))); + let input = PrimitiveArray::from_option_iter(values.clone()).into_array(); + let actual = execute(input, &dict_type(DataType::UInt32, DataType::UInt64))?; + let expected = UInt64Array::from_iter(values); + let flat = arrow_cast::cast(&actual, &DataType::UInt64)?; + assert_eq!(flat.as_ref(), &expected as &dyn arrow_array::Array); + Ok(()) + } + #[test] fn primitive_fallback_fills_codes_type() -> VortexResult<()> { let fits = PrimitiveArray::from_iter(0..=u8::MAX).into_array(); From ba9d3f15966ebebed5e99b76b6e5f2328c3fed50 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 22:50:38 +0000 Subject: [PATCH 04/28] Export patch-free bit-packed arrays as dictionaries over their bit width A bit-packed array without patches only holds values in 0..2^bit_width, so its unpacked values are already codes into that range. Exporting it to an Arrow dictionary now unpacks straight into the codes type and uses 0..2^bit_width as the values, instead of decoding to the values type and dictionary-encoding the result. Applies up to 16-bit widths when the dictionary is no larger than the array. To unpack into narrower codes, BitPacked's cast kernel now also handles narrowing integer casts when the array has no patches and the target holds every value the bit width allows, skipping the generic bounds check. PolarSignals SF1 on DataFusion (V2): q00-q03 another ~2x faster. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- Cargo.lock | 1 + vortex-arrow/Cargo.toml | 1 + vortex-arrow/src/executor/dictionary.rs | 101 +++++++++++++++++++++++- 3 files changed, 102 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 4a4e93d18c9..cab0d4e29d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11470,6 +11470,7 @@ dependencies = [ "vortex-buffer", "vortex-decimal-byte-parts", "vortex-error", + "vortex-fastlanes", "vortex-fsst", "vortex-mask", "vortex-onpair", diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 3900814ccdc..8904df2d7af 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -33,6 +33,7 @@ tracing = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true, features = ["arrow"] } vortex-error = { workspace = true } +vortex-fastlanes = { workspace = true } vortex-mask = { workspace = true } vortex-runend = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index d2924a0f1a4..8bfeb93aa49 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -26,11 +26,19 @@ use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray as VortexPrimitiveArray; use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; use vortex_array::matcher::Matcher; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::BitWidthsView; use crate::ArrowExporter; use crate::executor::execute_arrow_naive; @@ -42,7 +50,7 @@ impl Matcher for ArrowDictExportable { type Match<'a> = &'a ArrayRef; fn try_match(array: &ArrayRef) -> Option> { - (array.is::() || array.is::()).then_some(array) + (array.is::() || array.is::() || array.is::()).then_some(array) } } @@ -63,6 +71,9 @@ pub(super) fn to_arrow_dictionary( Ok(constant) => return constant_to_dict(constant, codes_type, values_type, exporter, ctx), Err(array) => array, }; + if let Some(dict) = bitpacked_to_dict(&array, codes_type, values_type, exporter, ctx)? { + return Ok(dict); + } // Otherwise, we should try and build a dictionary. let array = execute_arrow_naive(array, Some(values_type), exporter, ctx)?; @@ -77,6 +88,68 @@ pub(super) fn to_arrow_dictionary( .map_err(VortexError::from) } +/// The widest packed values [`bitpacked_to_dict`] exports as codes. +const MAX_BITPACKED_CODE_BITS: u8 = 16; + +/// Exports a patch-free bit-packed array as a dictionary over every value its bit width can hold. +/// +/// Packed values are non-negative and below `2^bit_width`, so the unpacked values are their own +/// codes into the values `0..2^bit_width`. Unpacking straight into the codes type skips decoding +/// to the values type and then dictionary-encoding the result. Returns `None` when the array is not +/// such a bit-packed array, the codes type cannot hold every packed value, or the dictionary would +/// outnumber the array's rows. +fn bitpacked_to_dict( + array: &ArrayRef, + codes_type: &DataType, + values_type: &DataType, + exporter: &ArrowExporter<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some(bitpacked) = array.as_opt::() else { + return Ok(None); + }; + let BitWidthsView::Global(bit_width) = bitpacked.bit_widths() else { + return Ok(None); + }; + let Some(codes_ptype) = codes_ptype(codes_type) else { + return Ok(None); + }; + let dict_len = 1usize << bit_width; + if bit_width > MAX_BITPACKED_CODE_BITS + || usize::from(bit_width) + usize::from(codes_ptype.is_signed_int()) + > codes_ptype.bit_width() + || !values_type.is_integer() + || dict_len > array.len() + || bitpacked.patches().is_some() + { + return Ok(None); + } + + let codes = array + .clone() + .cast(DType::Primitive(codes_ptype, array.dtype().nullability()))?; + let codes = execute_arrow_naive(codes, Some(codes_type), exporter, ctx)?; + let values = match_each_integer_ptype!(array.dtype().as_ptype(), |T| { + VortexPrimitiveArray::from_iter((0..dict_len).map(T::usize_as)).into_array() + }); + let values = execute_arrow_naive(values, Some(values_type), exporter, ctx)?; + make_dict_array(codes_type, codes, values).map(Some) +} + +fn codes_ptype(codes_type: &DataType) -> Option { + Some(match codes_type { + DataType::Int8 => PType::I8, + DataType::Int16 => PType::I16, + DataType::Int32 => PType::I32, + DataType::Int64 => PType::I64, + DataType::UInt8 => PType::U8, + DataType::UInt16 => PType::U16, + DataType::UInt32 => PType::U32, + DataType::UInt64 => PType::U64, + _ => return None, + }) +} + /// Dictionary-encodes a primitive Arrow array with codes of `codes_type`. /// /// Arrow's cast hashes every value with SipHash. This hashes the values' bits with FxHash instead, @@ -382,6 +455,8 @@ mod tests { use vortex_array::scalar_fn::fns::mask::Mask; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_fastlanes::BitPackedArrayExt; + use vortex_fastlanes::BitPackedData; use vortex_session::VortexSession; use crate::ArrowArrayExecutor; @@ -482,6 +557,30 @@ mod tests { Ok(()) } + #[rstest] + #[case::fits(DataType::UInt32, 11)] + #[case::signed_codes(DataType::Int16, 11)] + #[case::codes_narrower_than_bit_width(DataType::UInt16, 20)] + #[case::wide(DataType::UInt32, 20)] + fn bitpacked_exports_as_dict( + #[case] codes_type: DataType, + #[case] bit_width: u8, + ) -> VortexResult<()> { + let max = (1u64 << bit_width) - 1; + let values: Vec<_> = (0..5000u64) + .map(|i| (i % 13 != 4).then(|| i64::try_from(i * 7919 % max).unwrap_or(0))) + .collect(); + let source = PrimitiveArray::from_option_iter(values.clone()).into_array(); + let packed = + BitPackedData::encode(&source, bit_width, &mut SESSION.create_execution_ctx())?; + assert!(packed.patches().is_none()); + + let actual = execute(packed.into_array(), &dict_type(codes_type, DataType::Int64))?; + let flat = arrow_cast::cast(&actual, &DataType::Int64)?; + assert_eq!(flat.as_ref(), &Int64Array::from(values) as &dyn arrow_array::Array); + Ok(()) + } + #[test] fn primitive_fallback_fills_codes_type() -> VortexResult<()> { let fits = PrimitiveArray::from_iter(0..=u8::MAX).into_array(); From 25cb93672031909181029562bdde4d462dfe1317 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 22:53:43 +0000 Subject: [PATCH 05/28] Fix clippy and rustfmt findings Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- vortex-arrow/src/executor/dictionary.rs | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 8bfeb93aa49..3bac2fb3eda 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -270,6 +270,12 @@ where })) } +/// Codes and dictionary values built by [`encode_dense`] and [`encode_hashed`]. +type EncodedDict = ( + ScalarBuffer<::Native>, + ScalarBuffer<::Native>, +); + /// Assigns codes in order of first appearance through a table indexed by `key - min`. fn encode_dense( array: &PrimitiveArray, @@ -277,7 +283,7 @@ fn encode_dense( key: &impl Fn(T::Native) -> B, min: i128, range: usize, -) -> Option<(ScalarBuffer, ScalarBuffer)> +) -> Option> where K: ArrowDictionaryKeyType, T: ArrowPrimitiveType, @@ -309,7 +315,7 @@ fn encode_hashed( is_valid: &impl Fn(usize) -> bool, key: &impl Fn(T::Native) -> B, distinct: usize, -) -> Option<(ScalarBuffer, ScalarBuffer)> +) -> Option> where K: ArrowDictionaryKeyType, T: ArrowPrimitiveType, @@ -548,7 +554,8 @@ mod tests { #[test] fn primitive_fallback_distinct_values() -> VortexResult<()> { - let values = (0..2000u64).map(|i| (i % 7 != 3).then(|| i.wrapping_mul(0x9E37_79B9_7F4A_7C15))); + let values = + (0..2000u64).map(|i| (i % 7 != 3).then(|| i.wrapping_mul(0x9E37_79B9_7F4A_7C15))); let input = PrimitiveArray::from_option_iter(values.clone()).into_array(); let actual = execute(input, &dict_type(DataType::UInt32, DataType::UInt64))?; let expected = UInt64Array::from_iter(values); @@ -577,7 +584,10 @@ mod tests { let actual = execute(packed.into_array(), &dict_type(codes_type, DataType::Int64))?; let flat = arrow_cast::cast(&actual, &DataType::Int64)?; - assert_eq!(flat.as_ref(), &Int64Array::from(values) as &dyn arrow_array::Array); + assert_eq!( + flat.as_ref(), + &Int64Array::from(values) as &dyn arrow_array::Array + ); Ok(()) } From 5ba026b4a46ee3630e4ee81e16bf679d4756798f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 23:22:20 +0000 Subject: [PATCH 06/28] Skip the full min/max scan when sampled dictionary keys span a wide range The primitive dictionary export scanned every key for its range before choosing an encoding, even for high-cardinality columns such as addresses that can never use the dense table. Sample first, and scan the full range only when the sampled keys fit the dense table. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- vortex-arrow/src/executor/dictionary.rs | 63 +++++++++++++++---------- 1 file changed, 38 insertions(+), 25 deletions(-) diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 3bac2fb3eda..d0e7c96352d 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -235,31 +235,44 @@ where .map(|(_, &value)| key(value)) }; - let (min, max) = valid_keys().fold((i128::MAX, i128::MIN), |(min, max), key| { - let key = key.into(); - (min.min(key), max.max(key)) - }); - let range = usize::try_from(max.checked_sub(min)?).ok(); - let (codes, values) = match range { - Some(range) if range < MAX_DENSE_RANGE && range <= array.len().saturating_mul(4) => { - encode_dense::(array, &is_valid, &key, min, range)? - } - _ => { - let mut sample = FxHashSet::default(); - let sampled = valid_keys() - .take(SAMPLE_LEN) - .inspect(|key| { - sample.insert(*key); - }) - .count(); - if sample.len() * 8 >= sampled * 7 && K::Native::from_usize(array.len()).is_some() { - let codes = (0..array.len()).map(K::Native::usize_as).collect(); - (codes, array.values().clone()) - } else { - let distinct = sample.len() * array.len() / sampled.max(1); - encode_hashed::(array, &is_valid, &key, distinct)? - } - } + // A sample decides how to assign codes, and rules out a dense table without scanning every + // key when the sampled keys already span too wide a range. + let mut sample = FxHashSet::default(); + let (mut sample_min, mut sample_max) = (i128::MAX, i128::MIN); + let sampled = valid_keys() + .take(SAMPLE_LEN) + .inspect(|key| { + sample.insert(*key); + let key = (*key).into(); + sample_min = sample_min.min(key); + sample_max = sample_max.max(key); + }) + .count(); + let dense = sample_max + .checked_sub(sample_min) + .is_some_and(|range| range < MAX_DENSE_RANGE as i128) + .then(|| { + let (min, max) = valid_keys().fold((i128::MAX, i128::MIN), |(min, max), key| { + let key = key.into(); + (min.min(key), max.max(key)) + }); + let range = usize::try_from(max - min).ok()?; + (range < MAX_DENSE_RANGE && range <= array.len().saturating_mul(4)) + .then_some((min, range)) + }) + .flatten(); + + let (codes, values) = if let Some((min, range)) = dense { + encode_dense::(array, &is_valid, &key, min, range)? + } else if sampled > 0 + && sample.len() * 8 >= sampled * 7 + && K::Native::from_usize(array.len()).is_some() + { + let codes = (0..array.len()).map(K::Native::usize_as).collect(); + (codes, array.values().clone()) + } else { + let distinct = sample.len() * array.len() / sampled.max(1); + encode_hashed::(array, &is_valid, &key, distinct)? }; let codes = PrimitiveArray::::new(codes, nulls.cloned()); From 337acd0afced0c82e291e964bd9f0d4a0996337a Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 15:33:17 +0000 Subject: [PATCH 07/28] Rank run ends in one pass over the mask when filtering run-end arrays The run-end filter kernel called BitBuffer::count_range once per run. Each call sets up a SIMD popcount, which dominates when runs are a few rows long: TPC-DS store_sales columns average about 6 rows per run, and the kernel was the top CPU cost of TPC-DS q13 for both scans. Short runs now rank their ends in one pass over the mask's 64-bit words, one word popcount per run. Arrays whose runs average at least 256 rows keep one bulk popcount per run. run_end_filter benchmark: 4-row runs 300us -> 73us, 16-row runs 3.2us -> 1.2us, 1024-row runs 1.6us -> 0.8us. TPC-DS SF1 DataFusion geomean: V1 -3.6%, V2 -1.5%. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- encodings/runend/benches/run_end_filter.rs | 14 ++- encodings/runend/src/compute/filter.rs | 138 ++++++++++++++++++--- 2 files changed, 132 insertions(+), 20 deletions(-) diff --git a/encodings/runend/benches/run_end_filter.rs b/encodings/runend/benches/run_end_filter.rs index b012d43ca7d..43d66840334 100644 --- a/encodings/runend/benches/run_end_filter.rs +++ b/encodings/runend/benches/run_end_filter.rs @@ -5,8 +5,8 @@ //! //! This measures the kernel directly rather than going through the lazy //! `ArrayRef::filter` (which only builds a `FilterArray` node and does not run -//! the kernel). The hot work is a per-run popcount of the predicate mask, which -//! now uses `BitBuffer::count_range` (SIMD) instead of a bit-by-bit walk. +//! the kernel). The hot work is counting the predicate mask's set bits in each +//! run, which ranks the run ends in one pass over the mask's words. #![expect(clippy::cast_possible_truncation)] #![expect(clippy::cast_precision_loss)] @@ -81,6 +81,16 @@ const FILTER_ARGS: &[FilterBenchArgs] = &[ run_length: 16, density: 0.9, }, + FilterBenchArgs { + length: 65_536, + run_length: 4, + density: 0.3, + }, + FilterBenchArgs { + length: 65_536, + run_length: 1_024, + density: 0.3, + }, ]; /// Build the run-end boundaries (cumulative run lengths) for `length` rows. diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index d4b4df3b2b0..1c92c6f60fd 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -104,34 +104,46 @@ pub fn filter_run_end_primitive + AsPrim length: u64, mask: &BitBuffer, ) -> VortexResult<(PrimitiveArray, Mask)> { - let mut filtered_run_ends = buffer_mut![R::zero(); run_ends.len()]; + if mask.len() >= run_ends.len().saturating_mul(BULK_COUNT_MIN_RUN_LENGTH) { + let mut position = 0; + let mut counted = 0; + return filter_runs(run_ends, offset, length, |end| { + counted += mask.count_range(position, end); + position = end; + counted + }); + } + let mut ranks = WordRanks::new(mask); + filter_runs(run_ends, offset, length, |end| ranks.rank(end)) +} - let mut run_start = 0u64; +/// Filters run ends given `rank`, the set bits of the mask before each increasing run end. +fn filter_runs>( + run_ends: &[R], + offset: u64, + length: u64, + mut rank: impl FnMut(usize) -> usize, +) -> VortexResult<(PrimitiveArray, Mask)> { + let mut filtered_run_ends = buffer_mut![R::zero(); run_ends.len()]; + let mut run_start_rank = 0; let mut retained_run_count = 0; - let mut filtered_end = R::zero(); let values_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |run_idx| { let run_end = min(run_ends[run_idx].as_() - offset, length); - - // Bulk popcount is SIMD-capable and avoids per-bit reads. The input contract and clamp prove - // `run_start_idx <= run_end_idx <= mask.len()`. - let run_start_idx = run_start - .try_into() - .vortex_expect("run start index must fit in usize"); - let run_end_idx = run_end - .try_into() - .vortex_expect("run end index must fit in usize"); - let selected_in_run = mask.count_range(run_start_idx, run_end_idx); - filtered_end += ::from(selected_in_run) - .vortex_expect("run popcount must fit in run-end native type"); - let retain_run = selected_in_run > 0; + let run_end_rank = rank( + run_end + .try_into() + .vortex_expect("run end index must fit in usize"), + ); + let retain_run = run_end_rank > run_start_rank; + run_start_rank = run_end_rank; // Always write the current end, then advance only for a retained run. This keeps the loop // branchless. - filtered_run_ends[retained_run_count] = filtered_end; + filtered_run_ends[retained_run_count] = ::from(run_end_rank) + .vortex_expect("run popcount must fit in run-end native type"); retained_run_count += retain_run as usize; - run_start = run_end; retain_run }) .into(); @@ -144,15 +156,62 @@ pub fn filter_run_end_primitive + AsPrim )) } +/// Average run length from which one bulk SIMD popcount per run beats [`WordRanks`]. +const BULK_COUNT_MIN_RUN_LENGTH: usize = 256; + +/// Counts the set bits before increasing positions of a mask in one pass over its words. +/// +/// A run's selected rows are the difference of the ranks at its ends, so each short run costs a +/// word popcount rather than a bulk popcount call. +struct WordRanks { + words: Vec, + /// The words whose bits `counted` covers. + word: usize, + /// Set bits in `words[..word]`. + counted: usize, +} + +impl WordRanks { + fn new(mask: &BitBuffer) -> Self { + let chunks = mask.chunks(); + let mut words: Vec = chunks.iter().collect(); + if chunks.remainder_len() > 0 { + words.push(chunks.remainder_bits()); + } + Self { + words, + word: 0, + counted: 0, + } + } + + /// The set bits before `position`. Positions must not decrease between calls. + fn rank(&mut self, position: usize) -> usize { + let word = position / 64; + while self.word < word { + self.counted += self.words[self.word].count_ones() as usize; + self.word += 1; + } + let bits = position % 64; + if bits == 0 { + return self.counted; + } + self.counted + (self.words[word] & ((1u64 << bits) - 1)).count_ones() as usize + } +} + #[cfg(test)] mod tests { + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_mask::Mask; + use super::filter_run_end_primitive; use crate::RunEnd; use crate::RunEndArray; use crate::tests::SESSION; @@ -214,4 +273,47 @@ mod tests { Ok(()) } + + /// Both ways of counting a run's selected rows agree with a bit-by-bit count, including masks + /// that start inside a byte, end inside a word, and run ends past the array's length. + #[rstest] + #[case::short_runs(3)] + #[case::word_runs(64)] + #[case::bulk_runs(300)] + fn filter_run_end_counts_runs(#[case] run_length: u32) -> VortexResult<()> { + let (offset, length, mask_offset) = (5u32, 1_000u32, 3usize); + let run_ends: Vec = (1..) + .map(|run| offset + run * run_length - run % 2) + .take_while(|&end| end < offset + length + run_length) + .collect(); + let bits: Vec = (0..mask_offset + length as usize) + .map(|i| i % 7 == 0 || i % 5 == 1) + .collect(); + let mask = BitBuffer::from(bits.clone()).slice(mask_offset..mask_offset + length as usize); + + let (ends, values_mask) = + filter_run_end_primitive(&run_ends, offset.into(), length.into(), &mask)?; + + let mut expected_ends = Vec::new(); + let mut expected_values = Vec::new(); + let (mut start, mut selected) = (0usize, 0u32); + for &end in &run_ends { + let end = ((end - offset) as usize).min(length as usize); + let in_run = u32::try_from( + bits[mask_offset + start..mask_offset + end] + .iter() + .filter(|&&bit| bit) + .count(), + )?; + selected += in_run; + if in_run > 0 { + expected_ends.push(selected); + } + expected_values.push(in_run > 0); + start = end; + } + assert_eq!(ends.as_slice::(), expected_ends.as_slice()); + assert_eq!(values_mask, Mask::from_iter(expected_values)); + Ok(()) + } } From b3c32176f062a9707908fc393185b9a0d8cff1d7 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 16:30:37 +0000 Subject: [PATCH 08/28] Retain recently read segments across the chunks of V2 scans A V2 scan reads each chunk through a session of its own, and the file's segment source shares a read only while it is in flight. A stored segment that spans a chunk boundary was therefore read again by every chunk that started after its neighbour had finished, and by every partition and every concurrent scan of the file. TPC-DS q28 on SF1 read 188 MB where V1 reads 114 MB, every segment V1 reads up to two and a half times as often. The file's ScanCache now retains up to 64 MiB of recently read host segments in least-recently-used order, and every scan of the file reads through it. ScanCache::with_retained_bytes sets the budget; zero disables it. q28 now reads 125 MB. TPC-DS SF1 on DataFusion, V2 vs V1 geomean: 1.052 -> 1.002. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- vortex-layout/src/scan/v2/cache.rs | 25 ++++- vortex-layout/src/scan/v2/mod.rs | 7 +- vortex-layout/src/scan/v2/retain.rs | 150 ++++++++++++++++++++++++++++ vortex-layout/src/scan/v2/tests.rs | 52 +++++++++- 4 files changed, 226 insertions(+), 8 deletions(-) create mode 100644 vortex-layout/src/scan/v2/retain.rs diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index 21efeb3ac3f..ca2d93a5bdd 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -13,9 +13,12 @@ use vortex_array::expr::ExactBoundExpr; use vortex_error::VortexResult; use super::io::SegmentRanges; +use super::retain::RetainedSegments; +use super::retain::RetainingSegmentSource; use crate::plan::physical::PlanRef; use crate::plan::physical::pipeline::ValueCache; use crate::scan::execution::FilterPlans; +use crate::segments::SegmentSource; /// Reusable state for repeated V2 scans of one file, the V2 counterpart of a cached /// [`LayoutReader`](crate::LayoutReader) tree. @@ -26,7 +29,8 @@ use crate::scan::execution::FilterPlans; /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps /// its optimized plan, each distinct filter keeps its conjunct and zone-proof plans along with the /// zones those proofs have already excluded, and dictionary values stay keyed by the cached plan -/// that produced them. +/// that produced them. Recently read segments are retained up to a byte budget, so the chunks of +/// the file's scans that share a stored segment read it once. /// /// A cache describes one file's layout, segment locations, and segment source. Use a new cache /// when any of them changes. Cached plans and values live as long as the cache, one set per @@ -47,9 +51,20 @@ struct Inner { /// Serialized bytes the whole-file projection reads, by projection. projection_bytes: Mutex>, values: ValueCache, + /// Recently read segments, shared by the chunks of every scan of the file. + retained: Arc, } impl ScanCache { + /// A cache whose scans retain up to `bytes` of recently read segments, so chunks that share + /// a stored segment read it once. The default retains 64 MiB; zero retains nothing. + pub fn with_retained_bytes(bytes: u64) -> Self { + Self(Arc::new(Inner { + retained: Arc::new(RetainedSegments::new(bytes)), + ..Inner::default() + })) + } + /// The lowered root plan an earlier scan kept. pub(super) fn root(&self) -> Option { self.0.root.lock().clone() @@ -131,6 +146,14 @@ impl ScanCache { Ok(bytes) } + /// Serves `segments` through the segments the file's scans retain. + pub(super) fn retaining(&self, segments: Arc) -> Arc { + Arc::new(RetainingSegmentSource { + inner: segments, + retained: Arc::clone(&self.0.retained), + }) + } + pub(super) fn values(&self) -> ValueCache { self.0.values.clone() } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 3cfc947abf4..f4df511c591 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -17,11 +17,12 @@ //! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high. //! //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and proved -//! zones, segment index, and dictionary values with every other scan of that file, like a cached -//! layout reader does for V1 scans. +//! zones, segment index, dictionary values, and recently read segments with every other scan of +//! that file, like a cached layout reader does for V1 scans. mod cache; mod io; +mod retain; #[cfg(test)] mod tests; @@ -336,7 +337,7 @@ impl ScanRun { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)), }; - let io = Arc::new(SegmentScanIo::new(file.segments, ranges)); + let io = Arc::new(SegmentScanIo::new(cache.retaining(file.segments), ranges)); Ok(Self { session: builder.session.clone(), scan: Arc::new(PreparedScan { diff --git a/vortex-layout/src/scan/v2/retain.rs b/vortex-layout/src/scan/v2/retain.rs new file mode 100644 index 00000000000..5f50c4b80df --- /dev/null +++ b/vortex-layout/src/scan/v2/retain.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Keeps recently read segments so the chunks of a scan that share a segment read it once. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use futures::FutureExt; +use futures::TryFutureExt; +use parking_lot::Mutex; +use rustc_hash::FxHashMap; +use vortex_array::buffer::BufferHandle; + +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// The host bytes of recently read segments, evicting the least recently used past a budget. +/// +/// A V2 scan cuts its rows into chunks that each read through a session of their own, and a +/// stored segment that spans a chunk boundary is read by every chunk it overlaps. The file's segment +/// source shares a read only while it is in flight, so a chunk that starts after its neighbour +/// finished would read the segment again. Neighbouring chunks run close together, so a small +/// budget is enough to read each such segment once. +pub(super) struct RetainedSegments { + budget: u64, + state: Mutex, +} + +/// The bytes a file's scans retain by default. +const DEFAULT_RETAINED_BYTES: u64 = 64 << 20; + +impl Default for RetainedSegments { + fn default() -> Self { + Self::new(DEFAULT_RETAINED_BYTES) + } +} + +#[derive(Default)] +struct State { + bytes: u64, + tick: u64, + /// Each retained segment and the tick of its last use. + segments: FxHashMap, + /// Retained segments by the tick of their last use, oldest first. + by_use: BTreeMap, +} + +impl RetainedSegments { + pub(super) fn new(budget: u64) -> Self { + Self { + budget, + state: Mutex::default(), + } + } + + fn get(&self, id: SegmentId) -> Option { + let mut state = self.state.lock(); + state.tick += 1; + let tick = state.tick; + let (buffer, used) = state.segments.get_mut(&id)?; + let buffer = buffer.clone(); + let previous = std::mem::replace(used, tick); + state.by_use.remove(&previous); + state.by_use.insert(tick, id); + Some(buffer) + } + + fn insert(&self, id: SegmentId, buffer: &BufferHandle) { + let len = buffer.len() as u64; + if !buffer.is_on_host() || len > self.budget / 4 { + return; + } + let mut state = self.state.lock(); + if state.segments.contains_key(&id) { + return; + } + state.tick += 1; + let tick = state.tick; + state.segments.insert(id, (buffer.clone(), tick)); + state.by_use.insert(tick, id); + state.bytes += len; + while state.bytes > self.budget { + let Some((_, oldest)) = state.by_use.pop_first() else { + break; + }; + if let Some((evicted, _)) = state.segments.remove(&oldest) { + state.bytes -= evicted.len() as u64; + } + } + } +} + +/// Serves segments from [`RetainedSegments`] before reading them from `inner`, and retains +/// what it reads. +pub(super) struct RetainingSegmentSource { + pub(super) inner: Arc, + pub(super) retained: Arc, +} + +impl SegmentSource for RetainingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + if let Some(buffer) = self.retained.get(id) { + return futures::future::ready(Ok(buffer)).boxed(); + } + let retained = Arc::clone(&self.retained); + self.inner + .request(id) + .inspect_ok(move |buffer| retained.insert(id, buffer)) + .boxed() + } +} + +#[cfg(test)] +mod tests { + use vortex_array::buffer::BufferHandle; + use vortex_buffer::ByteBuffer; + + use super::RetainedSegments; + use crate::segments::SegmentId; + + fn buffer(len: usize) -> BufferHandle { + BufferHandle::new_host(ByteBuffer::from(vec![0u8; len])) + } + + #[test] + fn evicts_least_recently_used_past_budget() { + let retained = RetainedSegments::new(400); + retained.insert(SegmentId::from(0), &buffer(100)); + retained.insert(SegmentId::from(1), &buffer(100)); + retained.insert(SegmentId::from(2), &buffer(100)); + // Using segment 0 makes segment 1 the least recently used. + assert!(retained.get(SegmentId::from(0)).is_some()); + retained.insert(SegmentId::from(3), &buffer(100)); + retained.insert(SegmentId::from(4), &buffer(100)); + + assert!(retained.get(SegmentId::from(1)).is_none()); + for id in [0, 2, 3, 4] { + assert!(retained.get(SegmentId::from(id)).is_some(), "segment {id}"); + } + } + + #[test] + fn skips_segments_over_a_quarter_of_the_budget() { + let retained = RetainedSegments::new(400); + retained.insert(SegmentId::from(0), &buffer(101)); + assert!(retained.get(SegmentId::from(0)).is_none()); + } +} diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index b69f5390704..db0de52055b 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -520,7 +520,10 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { reads: AtomicUsize::new(0), }); let segments = Arc::clone(&counted) as Arc; - let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + let mut file = scan_file(&segments, &layout)?; + // Without retained segments, so a chunk reads every codes chunk it overlaps. + file.cache = Some(ScanCache::with_retained_bytes(0)); + let scan = ScanRun::try_new(ScanBuilder::new(session.clone(), file))?; let chunks = row_chunks(scan.rows.clone(), scan.scan.chunk_rows); assert_eq!( chunks.len(), @@ -553,6 +556,44 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { Ok(()) } +/// Work chunks that each overlap half of a stored segment read it once while the file's cache +/// retains segments, and once per chunk when it retains nothing. +#[rstest] +#[case::retained(None, 1)] +#[case::not_retained(Some(0), 2)] +#[tokio::test] +async fn chunks_sharing_a_segment_read_it_once( + #[case] retained_bytes: Option, + #[case] reads_per_segment: usize, +) -> VortexResult<()> { + let session = new_session().with_tokio(); + let (segments, layout) = write(&session, chunked_flat().as_ref()).await?; + let counted = Arc::new(CountingSegments { + inner: segments, + reads: AtomicUsize::new(0), + }); + let segments = Arc::clone(&counted) as Arc; + let mut file = scan_file(&segments, &layout)?; + file.cache = + Some(retained_bytes.map_or_else(ScanCache::default, ScanCache::with_retained_bytes)); + let projection = get_item("ints", root()).bind(layout.dtype())?; + let mut scan = ScanRun::try_new(ScanBuilder::new(session, file).with_projection(projection))?; + let chunk_rows = CHUNK_ROWS / 2; + Arc::get_mut(&mut scan.scan) + .ok_or_else(|| vortex_err!("the prepared scan is not shared yet"))? + .chunk_rows = chunk_rows; + + // One chunk at a time, so no read is still in flight when the next chunk starts. + for rows in row_chunks(scan.rows.clone(), chunk_rows) { + run_rows(&scan, rows).await?; + } + assert_eq!( + counted.reads.load(Ordering::Relaxed), + CHUNKS * reads_per_segment + ); + Ok(()) +} + /// V2 rejects list layouts: their element ranges are only known once their offsets are read. #[tokio::test] async fn list_layouts_return_an_error() -> VortexResult<()> { @@ -764,7 +805,8 @@ async fn scan_cache_reuses_values_between_scans( }); let segments = Arc::clone(&counted) as Arc; let mut file = scan_file(&segments, &layout)?; - file.cache = cached.then(ScanCache::default); + // Without retained segments, so only the values come from the cache. + file.cache = cached.then(|| ScanCache::with_retained_bytes(0)); let mut reads = Vec::new(); for _ in 0..2 { @@ -818,7 +860,8 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let mut file = scan_file(&segments, &layout)?; - file.cache = Some(ScanCache::default()); + // Without retained segments, so only the values come from the cache. + file.cache = Some(ScanCache::with_retained_bytes(0)); let projection = gt(root(), lit(15i32)).bind(layout.dtype())?; let mut reads = Vec::new(); @@ -901,7 +944,8 @@ async fn shared_cache_proves_zones_once() -> VortexResult<()> { }); let segments = Arc::clone(&counted) as Arc; let mut file = scan_file(&segments, &layout)?; - let cache = ScanCache::default(); + // Without retained segments, so only the proved zones come from the cache. + let cache = ScanCache::with_retained_bytes(0); file.cache = Some(cache.clone()); let filter = gt(get_item("ints", root()), lit(2_999i32)).bind(layout.dtype())?; let projection = get_item("ints", root()).bind(layout.dtype())?; From 78afd60c91404b2d7ce06bff99f593864bda428e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 14:05:58 +0000 Subject: [PATCH 09/28] Default datafusion-bench to the V2 scan for TPC-DS and ClickBench TPC-H already defaults to the planner-driven V2 scan. Extend the default to TPC-DS and ClickBench so the SQL benchmarks measure V2 on them too. `--scan-v2 false` or `VORTEX_SCAN_V2=0` still selects V1. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- benchmarks/datafusion-bench/src/main.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 1a83d374fcb..3d3983f43ee 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -120,7 +120,7 @@ struct Args { #[arg(long, value_delimiter = ',', value_parser = value_parser!(Format))] formats: Vec, - /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H). + /// Use the planner-driven V2 Vortex scan (defaults to true for TPC-H, TPC-DS and ClickBench). #[arg(long, action = clap::ArgAction::Set)] scan_v2: Option, @@ -137,8 +137,10 @@ async fn main() -> anyhow::Result<()> { let args = Args::parse(); let opts = Opts::from(args.options); let scan_v2 = args.scan_v2.unwrap_or_else(|| { - matches!(args.benchmark, BenchmarkArg::TpcH) - && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") + matches!( + args.benchmark, + BenchmarkArg::TpcH | BenchmarkArg::TpcDS | BenchmarkArg::ClickBench + ) && !std::env::var("VORTEX_SCAN_V2").is_ok_and(|v| v == "0") }); set_join_set_tracer(get_static_tracer())?; From 052c2346c9bc19bcc5cc9dfabaeed3af6601b081 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 26 Oct 2026 21:34:05 +0000 Subject: [PATCH 10/28] Reduce V2 scan predicate planning and sparse OR filtering work Avoid expanding whole structs for single-field predicates and evaluate infallible cross-column OR predicates before applying sparse selections. Increase only unfiltered scan and output batch limits. Signed-off-by: Joe Isaacs --- .../src/plan/physical/operators/filter.rs | 7 +++- .../src/plan/physical/plans/filter.rs | 18 ++++++--- vortex-layout/src/plan/physical/plans/pack.rs | 36 ++++++++++++------ vortex-layout/src/scan/execution/chunk.rs | 10 +++-- vortex-layout/src/scan/execution/filter.rs | 20 +++++++++- .../src/scan/execution/projection.rs | 14 +++++-- vortex-layout/src/scan/v2/mod.rs | 2 +- vortex-layout/src/scan/v2/tests.rs | 37 ++++++++++++++++--- 8 files changed, 110 insertions(+), 34 deletions(-) diff --git a/vortex-layout/src/plan/physical/operators/filter.rs b/vortex-layout/src/plan/physical/operators/filter.rs index ecd0943f97c..55c55309289 100644 --- a/vortex-layout/src/plan/physical/operators/filter.rs +++ b/vortex-layout/src/plan/physical/operators/filter.rs @@ -26,14 +26,17 @@ pub(crate) struct MaskStage { /// The batches are a predicate's result, which, mostly selected, is cheaper to execute /// whole than to filter lazily. predicate: bool, + /// The predicate is safe to evaluate outside the selected rows. + whole_predicate: bool, } impl MaskStage { - pub(crate) fn new(mask: Mask, predicate: bool) -> Self { + pub(crate) fn new(mask: Mask, predicate: bool, whole_predicate: bool) -> Self { Self { mask, cursor: 0, predicate, + whole_predicate, } } } @@ -53,7 +56,7 @@ impl Operator for MaskStage { if mask.all_true() { return Ok(Step::Final(batch)); } - if self.predicate && mask.density() >= EXPR_EVAL_THRESHOLD { + if self.predicate && (self.whole_predicate || mask.density() >= EXPR_EVAL_THRESHOLD) { return Ok(Step::Final( batch .execute::(cx.exec())? diff --git a/vortex-layout/src/plan/physical/plans/filter.rs b/vortex-layout/src/plan/physical/plans/filter.rs index 857a69a376e..37b56f14bd3 100644 --- a/vortex-layout/src/plan/physical/plans/filter.rs +++ b/vortex-layout/src/plan/physical/plans/filter.rs @@ -33,7 +33,10 @@ pub struct Filter; /// Operator-specific data for a [`Filter`] plan. #[derive(Clone, Debug)] -pub struct FilterData; +pub struct FilterData { + /// Evaluate a proven-infallible predicate before applying its selection. + whole_predicate: bool, +} /// A plan that keeps only the selected rows of its child. pub type FilterPlan = Plan; @@ -41,12 +44,16 @@ pub type FilterPlan = Plan; impl FilterPlan { /// Creates a filter over `child`. pub fn new(child: PlanRef) -> Self { + Self::for_predicate(child, false) + } + + pub(crate) fn for_predicate(child: PlanRef, whole_predicate: bool) -> Self { PlanParts { vtable: Filter, dtype: child.dtype().clone(), row_count: child.row_count(), children: vec![child].into(), - data: FilterData, + data: FilterData { whole_predicate }, } .into_typed() } @@ -66,8 +73,8 @@ impl PlanVTable for Filter { *ID } - fn metadata(_plan: &Plan) -> Option { - Some(EmptyMetadata) + fn metadata(plan: &Plan) -> Option { + (!plan.data().whole_predicate).then_some(EmptyMetadata) } fn with_children( @@ -110,12 +117,13 @@ impl PlanVTable for Filter { } let len = usize::try_from(rows.end - rows.start)?; let predicate = plan.dtype().is_boolean(); + let whole_predicate = plan.data().whole_predicate; let chain = compiler.compile(&child, rows, &Mask::new_true(len))?; Ok(chain.map(|chain| { if mask.all_true() { chain } else { - chain.with(MaskStage::new(mask.clone(), predicate)) + chain.with(MaskStage::new(mask.clone(), predicate, whole_predicate)) } })) } diff --git a/vortex-layout/src/plan/physical/plans/pack.rs b/vortex-layout/src/plan/physical/plans/pack.rs index d3886de19ad..fbe4bcee76f 100644 --- a/vortex-layout/src/plan/physical/plans/pack.rs +++ b/vortex-layout/src/plan/physical/plans/pack.rs @@ -344,8 +344,23 @@ impl PlanParentReduceRule for ExpressionPackRule { .get(&ExactBoundExpr(expression.clone())) .vortex_expect("Bound expression missing free-field annotations") .clone(); - let expanded_root = expanded_struct_root(child.dtype(), fields)?; - let expanded = expand_struct_root(expression.clone(), &expanded_root, fields)?; + // Most predicates reference one field. Avoid expanding the struct and building a + // partition tree when rebinding the expression to that field is sufficient. + if referenced_fields.len() == 1 { + let name = referenced_fields + .iter() + .next() + .vortex_expect("one referenced field"); + let index = fields.find(name).ok_or_else(|| { + vortex_err!("Struct expression references unknown field '{name}'") + })?; + let field = field_plan(child, index)?; + let lowered = step_into_struct_field(expression.clone(), name, field.dtype().clone())?; + if lowered.is_root_bound_to(field.dtype()) { + return Ok(Some(EvalPlan::try_new(lowered, field)?.into_plan())); + } + } + let expanded = expand_struct_root(expression.clone(), fields)?; let partitioned = partition_bound(expanded.clone(), make_bound_free_field_annotator(fields))?; @@ -533,14 +548,13 @@ fn is_identity_expression(expression: &BoundExpression, input_dtype: &DType) -> fn expand_struct_root( expression: BoundExpression, - expanded_root: &BoundExpression, fields: &StructFields, ) -> VortexResult { Ok(expression .transform_down(|node| { if node.is_root() { return Ok(Transformed { - value: expanded_root.clone(), + value: expanded_struct_root(node.dtype(), fields)?, changed: true, order: TraversalOrder::Skip, }); @@ -558,27 +572,25 @@ fn expand_struct_root( } if let Some(field_name) = scalar_fn.as_opt::() { - let index = fields.find(field_name).ok_or_else(|| { + fields.find(field_name).ok_or_else(|| { vortex_err!("Field {field_name} not found while expanding struct root") })?; return Ok(Transformed { - value: expanded_root.children()[index].clone(), - changed: true, + value: node, + changed: false, order: TraversalOrder::Skip, }); } if let Some(selection) = scalar_fn.as_opt::