diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs index 9a8147a36a6..8f4de0f4bd4 100644 --- a/src/array/dictionary/mod.rs +++ b/src/array/dictionary/mod.rs @@ -14,6 +14,7 @@ pub use iterator::*; pub use mutable::*; use super::{new_empty_array, primitive::PrimitiveArray, Array}; +use crate::scalar::NullScalar; /// Trait denoting [`NativeType`]s that can be used as keys of a dictionary. pub trait DictionaryKey: NativeType + num_traits::NumCast + num_traits::FromPrimitive { @@ -140,8 +141,12 @@ impl DictionaryArray { /// Returns the value of the [`DictionaryArray`] at position `i`. #[inline] pub fn value(&self, index: usize) -> Box { - let index = self.keys.value(index).to_usize().unwrap(); - new_scalar(self.values.as_ref(), index) + if self.keys.is_null(index) { + Box::new(NullScalar::new()) + } else { + let index = self.keys.value(index).to_usize().unwrap(); + new_scalar(self.values.as_ref(), index) + } } } diff --git a/src/array/growable/dictionary.rs b/src/array/growable/dictionary.rs index 4be93a53a2b..7fe1040cea4 100644 --- a/src/array/growable/dictionary.rs +++ b/src/array/growable/dictionary.rs @@ -100,7 +100,8 @@ impl<'a, T: DictionaryKey> Growable<'a> for GrowableDictionary<'a, T> { self.key_values.extend( values .iter() - .map(|x| T::from_usize(offset + x.to_usize().unwrap()).unwrap()), + // `.unwrap_or(0)` because this operation does not check for null values, which may contain any key. + .map(|x| T::from_usize(offset + x.to_usize().unwrap_or(0)).unwrap()), ); } diff --git a/tests/it/array/growable/dictionary.rs b/tests/it/array/growable/dictionary.rs index 61bacafd9be..8c27fc50070 100644 --- a/tests/it/array/growable/dictionary.rs +++ b/tests/it/array/growable/dictionary.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use arrow2::array::growable::{Growable, GrowableDictionary}; use arrow2::array::*; +use arrow2::datatypes::DataType; use arrow2::error::Result; #[test] @@ -29,6 +30,25 @@ fn test_single() -> Result<()> { Ok(()) } +#[test] +fn test_negative_keys() { + let vals = vec![Some("a"), Some("b"), Some("c")]; + let keys = vec![0, 1, 2, -1]; + + let keys = PrimitiveArray::from_data( + DataType::Int32, + keys.into(), + Some(vec![true, true, true, false].into()), + ); + + let arr = DictionaryArray::from_data(keys, Arc::new(Utf8Array::::from(vals))); + // check that we don't panic with negative keys to usize conversion + let mut growable = GrowableDictionary::new(&[&arr], false, 0); + growable.extend(0, 0, 4); + let out: DictionaryArray = growable.into(); + assert_eq!(out, arr); +} + #[test] fn test_multi() -> Result<()> { let mut original_data1 = vec![Some("a"), Some("b"), None, Some("a")];