From f62415321bc74c739e380199281d55f846438f0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B6rn=20Horstmann?= Date: Sat, 12 Jun 2021 14:46:27 +0200 Subject: [PATCH] Remove DictionaryArray::keys_array method and replace usages by the keys method (#419) --- arrow/src/array/array_dictionary.rs | 21 +++------------------ arrow/src/array/builder.rs | 2 +- arrow/src/array/ord.rs | 4 ++-- arrow/src/compute/kernels/cast.rs | 6 ++++-- arrow/src/compute/kernels/sort.rs | 4 ++-- arrow/src/compute/kernels/take.rs | 2 +- arrow/src/util/display.rs | 2 +- 7 files changed, 14 insertions(+), 27 deletions(-) diff --git a/arrow/src/array/array_dictionary.rs b/arrow/src/array/array_dictionary.rs index 9b036da7ae5d..168a4395b767 100644 --- a/arrow/src/array/array_dictionary.rs +++ b/arrow/src/array/array_dictionary.rs @@ -70,26 +70,11 @@ pub struct DictionaryArray { } impl<'a, K: ArrowPrimitiveType> DictionaryArray { - /// Return an iterator to the keys of this dictionary. + /// Return an array view of the keys of this dictionary as a PrimitiveArray. pub fn keys(&self) -> &PrimitiveArray { &self.keys } - /// Returns an array view of the keys of this dictionary - pub fn keys_array(&self) -> PrimitiveArray { - let data = self.data_ref(); - let keys_data = ArrayData::new( - K::DATA_TYPE, - data.len(), - Some(data.null_count()), - data.null_buffer().cloned(), - data.offset(), - data.buffers().to_vec(), - vec![], - ); - PrimitiveArray::::from(keys_data) - } - /// Returns the lookup key by doing reverse dictionary lookup pub fn lookup_key(&self, value: &str) -> Option { let rd_buf: &StringArray = @@ -379,7 +364,7 @@ mod tests { let test = vec!["a", "b", "c", "a"]; let array: DictionaryArray = test.into_iter().collect(); - let keys = array.keys_array(); + let keys = array.keys(); assert_eq!(&DataType::Int8, keys.data_type()); assert_eq!(0, keys.null_count()); assert_eq!(&[0, 1, 2, 0], keys.values()); @@ -390,7 +375,7 @@ mod tests { let test = vec![Some("a"), None, Some("b"), None, None, Some("a")]; let array: DictionaryArray = test.into_iter().collect(); - let keys = array.keys_array(); + let keys = array.keys(); assert_eq!(&DataType::Int32, keys.data_type()); assert_eq!(3, keys.null_count()); diff --git a/arrow/src/array/builder.rs b/arrow/src/array/builder.rs index 71fd9a90e7a0..3bbeaf794577 100644 --- a/arrow/src/array/builder.rs +++ b/arrow/src/array/builder.rs @@ -3282,7 +3282,7 @@ mod tests { assert_eq!(array.is_null(1), true); assert_eq!(array.is_valid(1), false); - let keys = array.keys_array(); + let keys = array.keys(); assert_eq!(keys.value(0), 1); assert_eq!(keys.is_null(1), true); diff --git a/arrow/src/array/ord.rs b/arrow/src/array/ord.rs index c2c94f5ffc7c..d7b18290be50 100644 --- a/arrow/src/array/ord.rs +++ b/arrow/src/array/ord.rs @@ -90,8 +90,8 @@ where { let left = left.as_any().downcast_ref::>().unwrap(); let right = right.as_any().downcast_ref::>().unwrap(); - let left_keys = left.keys_array(); - let right_keys = right.keys_array(); + let left_keys = left.keys(); + let right_keys = right.keys(); let left_values = StringArray::from(left.values().data().clone()); let right_values = StringArray::from(right.values().data().clone()); diff --git a/arrow/src/compute/kernels/cast.rs b/arrow/src/compute/kernels/cast.rs index 463da7c6fdfe..150f1f64386f 100644 --- a/arrow/src/compute/kernels/cast.rs +++ b/arrow/src/compute/kernels/cast.rs @@ -1370,7 +1370,8 @@ fn dictionary_cast( ) })?; - let keys_array: ArrayRef = Arc::new(dict_array.keys_array()); + let keys_array: ArrayRef = + Arc::new(PrimitiveArray::::from(dict_array.keys().data().clone())); let values_array = dict_array.values(); let cast_keys = cast_with_options(&keys_array, to_index_type, &cast_options)?; let cast_values = @@ -1450,7 +1451,8 @@ where cast_with_options(&dict_array.values(), to_type, cast_options)?; // Note take requires first casting the indices to u32 - let keys_array: ArrayRef = Arc::new(dict_array.keys_array()); + let keys_array: ArrayRef = + Arc::new(PrimitiveArray::::from(dict_array.keys().data().clone())); let indicies = cast_with_options(&keys_array, &DataType::UInt32, cast_options)?; let u32_indicies = indicies diff --git a/arrow/src/compute/kernels/sort.rs b/arrow/src/compute/kernels/sort.rs index feef00cda24b..dff5695750d1 100644 --- a/arrow/src/compute/kernels/sort.rs +++ b/arrow/src/compute/kernels/sort.rs @@ -600,7 +600,7 @@ fn sort_string_dictionary( ) -> Result { let values: &DictionaryArray = as_dictionary_array::(values); - let keys: &PrimitiveArray = &values.keys_array(); + let keys: &PrimitiveArray = values.keys(); let dict = values.values(); let dict: &StringArray = as_string_array(&dict); @@ -1071,7 +1071,7 @@ mod tests { .as_any() .downcast_ref::() .expect("Unable to get dictionary values"); - let sorted_keys = sorted.keys_array(); + let sorted_keys = sorted.keys(); assert_eq!(sorted_dict, dict); diff --git a/arrow/src/compute/kernels/take.rs b/arrow/src/compute/kernels/take.rs index d325ce44a185..66bfd613fd80 100644 --- a/arrow/src/compute/kernels/take.rs +++ b/arrow/src/compute/kernels/take.rs @@ -755,7 +755,7 @@ where I: ArrowNumericType, I::Native: ToPrimitive, { - let new_keys = take_primitive::(&values.keys_array(), indices)?; + let new_keys = take_primitive::(values.keys(), indices)?; let new_keys_data = new_keys.data_ref(); let data = ArrayData::new( diff --git a/arrow/src/util/display.rs b/arrow/src/util/display.rs index 61f549af7dd2..13d9f1959f6e 100644 --- a/arrow/src/util/display.rs +++ b/arrow/src/util/display.rs @@ -297,7 +297,7 @@ fn dict_array_value_to_string( ) -> Result { let dict_array = colum.as_any().downcast_ref::>().unwrap(); - let keys_array = dict_array.keys_array(); + let keys_array = dict_array.keys(); if keys_array.is_null(row) { return Ok(String::from(""));