diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index 4bb5f29719fb..fcce0ecbba70 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,74 @@ 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), + } +} + +/// 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, +) -> 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 +3458,60 @@ 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); + } + + #[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(