Dictionary-encode primitive Arrow exports without arrow_cast

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 }

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,15 +65,224 @@ 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())),
     )
     .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<ArrowArrayRef> {
+    match codes_type {
+        DataType::Int8 => primitive_to_dict_with::<Int8Type>(array),
+        DataType::Int16 => primitive_to_dict_with::<Int16Type>(array),
+        DataType::Int32 => primitive_to_dict_with::<Int32Type>(array),
+        DataType::Int64 => primitive_to_dict_with::<Int64Type>(array),
+        DataType::UInt8 => primitive_to_dict_with::<UInt8Type>(array),
+        DataType::UInt16 => primitive_to_dict_with::<UInt16Type>(array),
+        DataType::UInt32 => primitive_to_dict_with::<UInt32Type>(array),
+        DataType::UInt64 => primitive_to_dict_with::<UInt64Type>(array),
+        _ => None,
+    }
+}
+
+fn primitive_to_dict_with<K: ArrowDictionaryKeyType>(
+    array: &ArrowArrayRef,
+) -> Option<ArrowArrayRef> {
+    match array.data_type() {
+        DataType::Int8 => encode_primitive::<K, Int8Type, _>(array, |v| v),
+        DataType::Int16 => encode_primitive::<K, Int16Type, _>(array, |v| v),
+        DataType::Int32 => encode_primitive::<K, Int32Type, _>(array, |v| v),
+        DataType::Int64 => encode_primitive::<K, Int64Type, _>(array, |v| v),
+        DataType::UInt8 => encode_primitive::<K, UInt8Type, _>(array, |v| v),
+        DataType::UInt16 => encode_primitive::<K, UInt16Type, _>(array, |v| v),
+        DataType::UInt32 => encode_primitive::<K, UInt32Type, _>(array, |v| v),
+        DataType::UInt64 => encode_primitive::<K, UInt64Type, _>(array, |v| v),
+        DataType::Float16 => encode_primitive::<K, Float16Type, _>(array, |v| v.to_bits()),
+        DataType::Float32 => encode_primitive::<K, Float32Type, _>(array, |v| v.to_bits()),
+        DataType::Float64 => encode_primitive::<K, Float64Type, _>(array, |v| v.to_bits()),
+        DataType::Date32 => encode_primitive::<K, Date32Type, _>(array, |v| v),
+        DataType::Date64 => encode_primitive::<K, Date64Type, _>(array, |v| v),
+        DataType::Timestamp(TimeUnit::Second, _) => {
+            encode_primitive::<K, TimestampSecondType, _>(array, |v| v)
+        }
+        DataType::Timestamp(TimeUnit::Millisecond, _) => {
+            encode_primitive::<K, TimestampMillisecondType, _>(array, |v| v)
+        }
+        DataType::Timestamp(TimeUnit::Microsecond, _) => {
+            encode_primitive::<K, TimestampMicrosecondType, _>(array, |v| v)
+        }
+        DataType::Timestamp(TimeUnit::Nanosecond, _) => {
+            encode_primitive::<K, TimestampNanosecondType, _>(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<K, T, B>(
+    array: &ArrowArrayRef,
+    key: impl Fn(T::Native) -> B,
+) -> Option<ArrowArrayRef>
+where
+    K: ArrowDictionaryKeyType,
+    T: ArrowPrimitiveType,
+    B: Hash + Eq + Copy + Into<i128>,
+{
+    let array = array.as_primitive::<T>();
+    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::<K, T, B>(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::<K, T, B>(array, &is_valid, &key, distinct)?
+    };
+
+    let codes = PrimitiveArray::<K>::new(codes, nulls.cloned());
+    let values = PrimitiveArray::<T>::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<K, T> = (
+    ScalarBuffer<<K as ArrowPrimitiveType>::Native>,
+    ScalarBuffer<<T as ArrowPrimitiveType>::Native>,
+);
+
+/// Assigns codes in order of first appearance through a table indexed by `key - min`.
+fn encode_dense<K, T, B>(
+    array: &PrimitiveArray<T>,
+    is_valid: &impl Fn(usize) -> bool,
+    key: &impl Fn(T::Native) -> B,
+    min: i128,
+    range: usize,
+) -> Option<EncodedDict<K, T>>
+where
+    K: ArrowDictionaryKeyType,
+    T: ArrowPrimitiveType,
+    B: Into<i128>,
+{
+    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<K, T, B>(
+    array: &PrimitiveArray<T>,
+    is_valid: &impl Fn(usize) -> bool,
+    key: &impl Fn(T::Native) -> B,
+    distinct: usize,
+) -> Option<EncodedDict<K, T>>
+where
+    K: ArrowDictionaryKeyType,
+    T: ArrowPrimitiveType,
+    B: Hash + Eq,
+{
+    let mut lookup: FxHashMap<B, K::Native> =
+        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::<ArrowDictArray<UInt8Type>>()) 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::<UInt32Type>::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::<Int16Type>::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::<UInt8Type>().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