From 9a086f3aad65dfb0f94f0405100760a6b980bf49 Mon Sep 17 00:00:00 2001 From: Chinmay Shrivastava <67478015+JonSnow1807@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:51:29 -0600 Subject: [PATCH 1/2] Optimize byte-array dictionary gathering in arrow-row --- arrow-row/src/lib.rs | 159 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 157 insertions(+), 2 deletions(-) diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index 4bb5f29719fb..aa83c0da0fa5 100644 --- a/arrow-row/src/lib.rs +++ b/arrow-row/src/lib.rs @@ -160,9 +160,12 @@ use std::slice::Windows; use std::sync::Arc; use arrow_array::cast::*; -use arrow_array::types::{ArrowDictionaryKeyType, ByteArrayType, ByteViewType}; +use arrow_array::types::{ + ArrowDictionaryKeyType, BinaryType, ByteArrayType, ByteViewType, LargeBinaryType, + LargeUtf8Type, Utf8Type, +}; use arrow_array::*; -use arrow_buffer::{ArrowNativeType, Buffer, OffsetBuffer, ScalarBuffer}; +use arrow_buffer::{ArrowNativeType, Buffer, NullBufferBuilder, OffsetBuffer, ScalarBuffer}; use arrow_data::transform::MutableArrayData; use arrow_schema::*; use variable::{decode_binary_view, decode_string_view}; @@ -2293,6 +2296,72 @@ pub fn encode_dictionary_values( fn gather_dictionary_values( array: &DictionaryArray, +) -> Result { + let values = array.values(); + match values.data_type() { + DataType::Utf8 => gather_dictionary_byte_values::( + array, + values.as_any().downcast_ref::().unwrap(), + ), + DataType::LargeUtf8 => gather_dictionary_byte_values::( + array, + values.as_any().downcast_ref::().unwrap(), + ), + DataType::Binary => gather_dictionary_byte_values::( + array, + values.as_any().downcast_ref::().unwrap(), + ), + DataType::LargeBinary => gather_dictionary_byte_values::( + array, + values.as_any().downcast_ref::().unwrap(), + ), + _ => gather_dictionary_values_generic(array), + } +} + +fn gather_dictionary_byte_values( + array: &DictionaryArray, + values: &GenericByteArray, +) -> Result { + let mut offsets = Vec::with_capacity(array.len() + 1); + offsets.push(T::Offset::from_usize(0).unwrap()); + let mut bytes = Vec::new(); + let mut nulls = NullBufferBuilder::new(array.len()); + + for key in array.keys() { + match key { + Some(key) if values.is_valid(key.as_usize()) => { + let index = key.as_usize(); + let start = values.value_offsets()[index].as_usize(); + let end = values.value_offsets()[index + 1].as_usize(); + let value = &values.value_data()[start..end]; + let end = bytes + .len() + .checked_add(value.len()) + .ok_or(ArrowError::OffsetOverflowError(usize::MAX))?; + let offset = + T::Offset::from_usize(end).ok_or(ArrowError::OffsetOverflowError(end))?; + bytes.extend_from_slice(value); + offsets.push(offset); + nulls.append(true); + } + _ => { + offsets.push(*offsets.last().unwrap()); + nulls.append(false); + } + } + } + + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + Ok(Arc::new(GenericByteArray::::try_new( + offsets, + Buffer::from_vec(bytes), + nulls.finish(), + )?)) +} + +fn gather_dictionary_values_generic( + array: &DictionaryArray, ) -> Result { let data = array.values().to_data(); let mut mutable = MutableArrayData::new(vec![&data], false, array.len()); @@ -3387,6 +3456,92 @@ mod tests { dictionary_eq(&decoded[0], &expected); } + #[test] + fn test_gather_dictionary_byte_values() { + let keys = Int8Array::from(vec![Some(0), Some(3), None, Some(1), Some(3), Some(2)]); + let values = Arc::new(StringArray::from(vec![ + Some("unused"), + None, + Some("é"), + Some("selected"), + ])); + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + let sliced = dictionary.slice(1, 4); + let gathered = gather_dictionary_values(&sliced).unwrap(); + let expected = StringArray::from(vec![Some("selected"), None, None, Some("selected")]); + assert_eq!(gathered.as_ref(), &expected as &dyn Array); + + let keys = UInt64Array::from(vec![Some(2), None, Some(0), Some(1), Some(2)]); + let values = Arc::new(LargeStringArray::from(vec![Some(""), Some("wide"), None])); + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + let sliced = dictionary.slice(1, 4); + let gathered = gather_dictionary_values(&sliced).unwrap(); + let expected = LargeStringArray::from(vec![None, Some(""), Some("wide"), None]); + assert_eq!(gathered.as_ref(), &expected as &dyn Array); + + let keys = UInt16Array::from(vec![Some(0), Some(2), None, Some(1), Some(2)]); + let values = Arc::new(BinaryArray::from(vec![ + Some(&b"unused"[..]), + Some(&b"\xff\0"[..]), + None, + ])); + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + let sliced = dictionary.slice(1, 4); + let gathered = gather_dictionary_values(&sliced).unwrap(); + let expected = BinaryArray::from(vec![None, None, Some(&b"\xff\0"[..]), None]); + assert_eq!(gathered.as_ref(), &expected as &dyn Array); + + let keys = Int64Array::from(vec![Some(0), Some(2), None, Some(1), Some(2)]); + let values = Arc::new(LargeBinaryArray::from(vec![ + Some(&b"unused"[..]), + Some(&b"\0\xff"[..]), + None, + ])); + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + let sliced = dictionary.slice(1, 4); + let gathered = gather_dictionary_values(&sliced).unwrap(); + let expected = LargeBinaryArray::from(vec![None, None, Some(&b"\0\xff"[..]), None]); + assert_eq!(gathered.as_ref(), &expected as &dyn Array); + } + + #[test] + fn test_gather_dictionary_generic_values() { + let keys = UInt8Array::from(vec![Some(2), None, Some(0), Some(1), Some(2)]); + let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + let sliced = dictionary.slice(1, 4); + let gathered = gather_dictionary_values(&sliced).unwrap(); + let expected = Int32Array::from(vec![Some(1), Some(1), None, Some(3)]); + assert_eq!(gathered.as_ref(), &expected as &dyn Array); + } + + #[test] + fn test_sparse_binary_dictionary_with_null_values() { + let values = BinaryArray::from_iter((0..DICTIONARY_GATHER_MIN_VALUES).map(|index| { + if index == 0 { + None + } else { + Some(&b"value"[..]) + } + })); + let keys = UInt64Array::from(vec![Some(0), Some(32_767), None]); + let dictionary = DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let input = Arc::new(dictionary) as ArrayRef; + let expected = + Arc::new(BinaryArray::from(vec![None, Some(&b"value"[..]), None])) as ArrayRef; + let converter = RowConverter::new(vec![SortField::new(input.data_type().clone())]).unwrap(); + let actual_rows = converter.convert_columns(&[input]).unwrap(); + let expected_converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap(); + let expected_rows = expected_converter + .convert_columns(&[Arc::clone(&expected)]) + .unwrap(); + for (actual, expected) in actual_rows.iter().zip(&expected_rows) { + assert_eq!(actual, expected); + } + let decoded = converter.convert_rows(&actual_rows).unwrap(); + assert_eq!(decoded[0].as_ref(), expected.as_ref()); + } + #[test] fn test_sparse_dictionary_gather_overflow_falls_back() { let inner = RunArray::::try_new( From 76ee5adf76ecf58c70a09eb6fe5b73282813bf54 Mon Sep 17 00:00:00 2001 From: Chinmay Shrivastava <67478015+JonSnow1807@users.noreply.github.com> Date: Sat, 7 Nov 2026 12:23:47 -0600 Subject: [PATCH 2/2] Simplify dictionary gather test and document byte helper --- arrow-row/src/lib.rs | 34 ++-------------------------------- 1 file changed, 2 insertions(+), 32 deletions(-) diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index aa83c0da0fa5..fcce0ecbba70 100644 --- a/arrow-row/src/lib.rs +++ b/arrow-row/src/lib.rs @@ -2319,6 +2319,8 @@ fn gather_dictionary_values( } } +/// Gathers byte values in dictionary-key order, preserving null keys and values. +/// Returns an error if the accumulated byte offsets overflow. fn gather_dictionary_byte_values( array: &DictionaryArray, values: &GenericByteArray, @@ -3470,38 +3472,6 @@ mod tests { let gathered = gather_dictionary_values(&sliced).unwrap(); let expected = StringArray::from(vec![Some("selected"), None, None, Some("selected")]); assert_eq!(gathered.as_ref(), &expected as &dyn Array); - - let keys = UInt64Array::from(vec![Some(2), None, Some(0), Some(1), Some(2)]); - let values = Arc::new(LargeStringArray::from(vec![Some(""), Some("wide"), None])); - let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); - let sliced = dictionary.slice(1, 4); - let gathered = gather_dictionary_values(&sliced).unwrap(); - let expected = LargeStringArray::from(vec![None, Some(""), Some("wide"), None]); - assert_eq!(gathered.as_ref(), &expected as &dyn Array); - - let keys = UInt16Array::from(vec![Some(0), Some(2), None, Some(1), Some(2)]); - let values = Arc::new(BinaryArray::from(vec![ - Some(&b"unused"[..]), - Some(&b"\xff\0"[..]), - None, - ])); - let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); - let sliced = dictionary.slice(1, 4); - let gathered = gather_dictionary_values(&sliced).unwrap(); - let expected = BinaryArray::from(vec![None, None, Some(&b"\xff\0"[..]), None]); - assert_eq!(gathered.as_ref(), &expected as &dyn Array); - - let keys = Int64Array::from(vec![Some(0), Some(2), None, Some(1), Some(2)]); - let values = Arc::new(LargeBinaryArray::from(vec![ - Some(&b"unused"[..]), - Some(&b"\0\xff"[..]), - None, - ])); - let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); - let sliced = dictionary.slice(1, 4); - let gathered = gather_dictionary_values(&sliced).unwrap(); - let expected = LargeBinaryArray::from(vec![None, None, Some(&b"\0\xff"[..]), None]); - assert_eq!(gathered.as_ref(), &expected as &dyn Array); } #[test]