From 9059cbf0abe07f64ae69183bebe5e1ec78bde738 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 20 Jun 2022 02:10:16 -0400 Subject: [PATCH] Add `DictionaryArray::key` function (#1912) --- arrow/src/array/array_dictionary.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/arrow/src/array/array_dictionary.rs b/arrow/src/array/array_dictionary.rs index 0fbd5a34eb6..b9da60ccfe8 100644 --- a/arrow/src/array/array_dictionary.rs +++ b/arrow/src/array/array_dictionary.rs @@ -169,6 +169,17 @@ impl<'a, K: ArrowPrimitiveType> DictionaryArray { .iter() .map(|key| key.map(|k| k.to_usize().expect("Dictionary index not usize"))) } + + /// Return the value of `keys` (the dictionary key) at index `i`, + /// cast to `usize`, `None` if the value at `i` is `NULL`. + pub fn key(&self, i: usize) -> Option { + self.keys.is_valid(i).then(|| { + self.keys + .value(i) + .to_usize() + .expect("Dictionary index not usize") + }) + } } /// Constructs a `DictionaryArray` from an array data reference. @@ -534,6 +545,17 @@ mod tests { assert!(iter.next().is_none()); } + #[test] + fn test_dictionary_key() { + let keys = Int8Array::from(vec![Some(2), None, Some(1)]); + let values = StringArray::from(vec!["foo", "bar", "baz", "blarg"]); + + let array = DictionaryArray::try_new(&keys, &values).unwrap(); + assert_eq!(array.key(0), Some(2)); + assert_eq!(array.key(1), None); + assert_eq!(array.key(2), Some(1)); + } + #[test] fn test_try_new() { let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]