From a6867e9c36313e2c3fd1bc05b32d26001f8493fb Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 26 Oct 2026 14:15:10 +0000 Subject: [PATCH] Dictionary-encode primitive Arrow exports without arrow_cast Exporting a primitive array as an Arrow dictionary went through arrow_cast, which hashes each value with SipHash and builds an ArrowError per value. On PolarSignals this export was ~90% of scan CPU. Encode primitive, float (by bits), date and timestamp arrays natively, choosing in order: a dense table when the keys span fewer than 65,536 values; the values themselves with sequential codes when a sample is nearly all distinct (Arrow dictionaries may hold duplicates); otherwise an FxHash map sized from the sample. The full min/max scan only runs when the sample's range could fit the dense table. Nulls stay in the codes. Other types, and codes that would overflow, fall back to arrow_cast. Split out of vortex-data/vortex#10761. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SQocaYaDvJHfmorKXnyvJC --- Cargo.lock | 1 + vortex-arrow/Cargo.toml | 1 + vortex-arrow/src/executor/dictionary.rs | 269 +++++++++++++++++++++++- 3 files changed, 270 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 7aad4bb823c..e76cf88f5c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11450,6 +11450,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 343c08357af..945b24e8ec6 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -1,15 +1,24 @@ // 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_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; @@ -56,8 +65,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 +77,212 @@ 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, + } +} + +/// 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, +) -> Option +where + K: ArrowDictionaryKeyType, + T: ArrowPrimitiveType, + B: Hash + Eq + Copy + Into, +{ + let array = array.as_primitive::(); + 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)) + }; + + // 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()); + 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)) + })) +} + +/// 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, + is_valid: &impl Fn(usize) -> bool, + key: &impl Fn(T::Native) -> B, + min: i128, + range: usize, +) -> Option> +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 !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> +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; + } + 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); + } + Some((codes.into(), values.into())) +} + /// Convert a constant array to a dictionary with a single entry. fn constant_to_dict( array: ConstantArray, @@ -160,7 +378,14 @@ 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::UInt64Array; + 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 +464,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 +490,32 @@ 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(); + 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