This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 9059cbf0a Add `DictionaryArray::key` function (#1912)
9059cbf0a is described below
commit 9059cbf0abe07f64ae69183bebe5e1ec78bde738
Author: Andrew Lamb <[email protected]>
AuthorDate: Mon Jun 20 02:10:16 2022 -0400
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 0fbd5a34e..b9da60ccf 100644
--- a/arrow/src/array/array_dictionary.rs
+++ b/arrow/src/array/array_dictionary.rs
@@ -169,6 +169,17 @@ impl<'a, K: ArrowPrimitiveType> DictionaryArray<K> {
.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<usize> {
+ 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")]
