From a94ac8b738ed1b75ab6e2bd9bc986f71050758a0 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 19:48:30 +0000 Subject: [PATCH 1/6] 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 --- .../src/scan/execution/projection.rs | 75 +++++++++++++------ vortex-layout/src/scan/v2/cache.rs | 19 +++++ vortex-layout/src/scan/v2/mod.rs | 71 +++++++++++++++--- 3 files changed, 134 insertions(+), 31 deletions(-) diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 00c0a0614f6..71bf4c5c653 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; 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 { plans: Option, 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 plans = self @@ -185,7 +217,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 3daa0975a0e..5a02816a00d 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -9,6 +9,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; @@ -36,6 +37,8 @@ struct Inner { /// them, so a projection must reuse its plan for its values to be found again. projections: Mutex>, ranges: Mutex>, + /// Serialized bytes the whole-file projection reads, by projection. + projection_bytes: Mutex>, values: ValueCache, } @@ -80,6 +83,22 @@ impl ScanCache { Arc::clone(self.0.ranges.lock().get_or_insert(ranges)) } + /// 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 c6b359fa328..1b552f37838 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, segment index, and dictionary @@ -45,6 +46,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; +use vortex_mask::Mask; use vortex_scan::planning::planner::WorkScope; use vortex_scan::selection::Selection; use vortex_scan::strict_sorted_buffer::StrictSortedBuffer; @@ -55,6 +57,7 @@ pub use self::cache::ScanCache; 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::FilterPlans; @@ -76,13 +79,39 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize { 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 task 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. +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); + by_bytes.min(by_threads).min(max_blocks) as usize * 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 +} + +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. @@ -307,18 +336,40 @@ impl PreparedScan { None => { let plan = optimize(plan_row_idx_expression(builder.projection.clone(), root)?)?; - 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 { MAX_ROWS_PER_CHUNK }; - let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows); + let splits = row_chunks(rows, chunk_rows); let ranges = match cache.ranges() { Some(ranges) => ranges, None => cache.insert_ranges(segment_ranges(&file.locations)), From 87a48d8e4c9605523fb95f40bb616d5352d13b13 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 19:48:30 +0000 Subject: [PATCH 2/6] 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 a471337bea4..9d8aad0727a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11473,6 +11473,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 8605b37a3f0..a9662d8d7f2 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 8f16ec28528..8205abb4a04 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, @@ -159,7 +261,13 @@ mod tests { use std::sync::Arc; 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; @@ -236,6 +344,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, @@ -246,6 +370,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 e71f1373ed688274962680f77ac4c40fa3e0ea38 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 22:34:36 +0000 Subject: [PATCH 3/6] 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 8205abb4a04..1041bf17af2 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. @@ -266,6 +363,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; @@ -370,6 +468,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 bcbe8e3b02dfd9bc36b609f73ed9692237e46a2d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 22:50:38 +0000 Subject: [PATCH 4/6] 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 + .../fastlanes/src/bitpacking/compute/cast.rs | 52 ++++++++- vortex-arrow/Cargo.toml | 1 + vortex-arrow/src/executor/dictionary.rs | 104 +++++++++++++++++- 4 files changed, 153 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9d8aad0727a..35498023cb1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11479,6 +11479,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-fastlanes", "vortex-fsst", "vortex-mask", "vortex-onpair", diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index b41a3dbb14f..22521884c45 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -32,6 +32,14 @@ fn is_widening_int_cast(src: PType, tgt: PType) -> bool { && (src.is_unsigned_int() || tgt.is_signed_int()) } +/// Returns `true` if every value a patch-free array packed at `bit_width` can hold fits `tgt`. +/// Packed values are non-negative and below `2^bit_width`. +fn packed_values_fit(bit_width: u8, src: PType, tgt: PType) -> bool { + src.is_int() + && tgt.is_int() + && usize::from(bit_width) + usize::from(tgt.is_signed_int()) <= tgt.bit_width() +} + fn build_with_validity( array: ArrayView<'_, BitPacked>, bit_width: u8, @@ -89,15 +97,17 @@ impl CastKernel for BitPacked { return build_with_validity(array, bit_width, dtype, new_validity).map(Some); } - // Widening integer cast: unpack each FastLanes chunk into a cache-resident scratch buffer - // and cast-copy straight into the wide output, avoiding a full-length intermediate buffer - // and the generic cast kernel's bounds-check scan (unnecessary when widening). + // Widening integer cast, or a narrowing one whose target holds every packed value: unpack + // each FastLanes chunk into a cache-resident scratch buffer and cast-copy straight into the + // output, avoiding a full-length intermediate buffer and the generic cast kernel's + // bounds-check scan (unnecessary when every value fits). let DType::Primitive(tgt, tgt_nullability) = dtype else { return Ok(None); }; let (tgt, tgt_nullability) = (*tgt, *tgt_nullability); let src = array.dtype().as_ptype(); - if !is_widening_int_cast(src, tgt) { + let fits = || array.patches().is_none() && packed_values_fit(bit_width, src, tgt); + if !is_widening_int_cast(src, tgt) && !fits() { return Ok(None); } @@ -144,6 +154,7 @@ mod tests { use vortex_session::VortexSession; use crate::BitPackedArray; + use crate::BitPackedArrayExt; use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { @@ -266,6 +277,39 @@ mod tests { Ok(()) } + /// Narrowing casts of patch-free arrays whose bit width fits the target unpack directly and + /// match a plain primitive cast, including sliced and nullable inputs. + #[rstest] + #[case(PType::I64, PType::U8)] + #[case(PType::I64, PType::I8)] + #[case(PType::U64, PType::U16)] + #[case(PType::I32, PType::I16)] + fn test_cast_bitpacked_narrowing_fits( + #[case] src: PType, + #[case] tgt: PType, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source = match_each_integer_ptype!(src, |S| { + PrimitiveArray::from_option_iter((0..2051usize).map(|i| { + (i % 11 != 5).then(|| { + ::from_usize(i % 128) + .expect("test values fit every integer ptype") + }) + })) + }) + .into_array(); + let packed = bp(&source, 7); + assert!(packed.patches().is_none()); + let target = DType::Primitive(tgt, Nullability::Nullable); + + let casted = packed.clone().into_array().cast(target.clone())?; + assert_arrays_eq!(casted, source.clone().cast(target.clone())?, &mut ctx); + + let sliced = packed.into_array().slice(500..1700)?.cast(target.clone())?; + assert_arrays_eq!(sliced, source.slice(500..1700)?.cast(target)?, &mut ctx); + Ok(()) + } + #[rstest] #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))] #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))] diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index a9662d8d7f2..948e03a9865 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 1041bf17af2..fd1cf75019d 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 crate::ArrowArrayExecutor; use crate::executor::dictionary::ConstantArray; @@ -479,6 +554,33 @@ 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 array_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 e922b779e78d73d1a20bebd0e57caa2f58e6f086 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 22:53:43 +0000 Subject: [PATCH 5/6] Fix clippy and rustfmt findings Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LYor8ydYMRcjUSh622VTEC Signed-off-by: Claude --- .../fastlanes/src/bitpacking/compute/cast.rs | 2 +- vortex-arrow/src/executor/dictionary.rs | 18 ++++++++++++++---- vortex-layout/src/scan/v2/mod.rs | 3 ++- 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 22521884c45..b57b4bfe04b 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -303,7 +303,7 @@ mod tests { let target = DType::Primitive(tgt, Nullability::Nullable); let casted = packed.clone().into_array().cast(target.clone())?; - assert_arrays_eq!(casted, source.clone().cast(target.clone())?, &mut ctx); + assert_arrays_eq!(casted, source.cast(target.clone())?, &mut ctx); let sliced = packed.into_array().slice(500..1700)?.cast(target.clone())?; assert_arrays_eq!(sliced, source.slice(500..1700)?.cast(target)?, &mut ctx); diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index fd1cf75019d..968a11b73e7 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, @@ -545,7 +551,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(()) } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 1b552f37838..cbb65c5959f 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -98,7 +98,8 @@ fn filtered_rows_per_chunk( 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); - by_bytes.min(by_threads).min(max_blocks) as usize * ROW_BLOCK_SIZE + 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. From a47ec03eba357d86d166c40a363a7fbbd7a688f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 23:22:20 +0000 Subject: [PATCH 6/6] 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());