Fix: simplify dict array compute functions and report correct dictarray dtype now that codes and values don't have same nullability (#4579)

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

encodings/dict/src/array.rs

@@ -44,6 +44,7 @@ pub struct DictArray {
     codes: ArrayRef,
     values: ArrayRef,
     stats_set: ArrayStats,
+    dtype: DType,
 }
 
 #[derive(Clone, Debug)]
@@ -57,10 +58,14 @@ impl DictArray {
     /// by the safe [`DictArray::try_new`] constructor are valid, for example when
     /// you are filtering or slicing an existing valid `DictArray`.
     pub unsafe fn new_unchecked(codes: ArrayRef, values: ArrayRef) -> Self {
+        let dtype = values
+            .dtype()
+            .union_nullability(codes.dtype().nullability());
         Self {
             codes,
             values,
             stats_set: Default::default(),
+            dtype,
         }
     }
 
@@ -88,11 +93,7 @@ impl DictArray {
             vortex_bail!(MismatchedTypes: "unsigned int", codes.dtype());
         }
 
-        Ok(Self {
-            codes,
-            values,
-            stats_set: Default::default(),
-        })
+        Ok(unsafe { Self::new_unchecked(codes, values) })
     }
 
     #[inline]
@@ -112,7 +113,7 @@ impl ArrayVTable<DictVTable> for DictVTable {
     }
 
     fn dtype(array: &DictArray) -> &DType {
-        array.values.dtype()
+        &array.dtype
     }
 
     fn stats(array: &DictArray) -> StatsSetRef<'_> {

encodings/dict/src/compute/cast.rs

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_array::compute::{CastKernel, CastKernelAdapter, cast};
-use vortex_array::{ArrayRef, IntoArray, register_kernel};
+use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
@@ -13,7 +13,8 @@ impl CastKernel for DictVTable {
         // Cast the dictionary values to the target type
         let casted_values = cast(array.values(), dtype)?;
 
-        let casted_codes = if dtype.nullability() != array.codes().dtype().nullability() {
+        // If the codes are nullable but we are casting to non nullable dtype we have to remove nullability from codes as well
+        let casted_codes = if array.codes().dtype().is_nullable() && !dtype.is_nullable() {
             cast(
                 array.codes(),
                 &array.codes().dtype().with_nullability(dtype.nullability()),
@@ -23,11 +24,9 @@ impl CastKernel for DictVTable {
         };
 
         // SAFETY: casting does not alter invariants of the codes
-        unsafe {
-            Ok(Some(
-                DictArray::new_unchecked(casted_codes, casted_values).into_array(),
-            ))
-        }
+        Ok(Some(
+            unsafe { DictArray::new_unchecked(casted_codes, casted_values) }.into_array(),
+        ))
     }
 }
 
@@ -132,7 +131,7 @@ mod tests {
         let nullable_dict = nullable.as_::<DictVTable>();
         assert_eq!(
             nullable_dict.codes().dtype().nullability(),
-            Nullability::Nullable
+            Nullability::NonNullable
         );
         assert_eq!(
             nullable_dict.values().dtype().nullability(),

encodings/dict/src/compute/mod.rs

@@ -10,7 +10,7 @@ mod like;
 mod min_max;
 
 use vortex_array::compute::{
-    FilterKernel, FilterKernelAdapter, TakeKernel, TakeKernelAdapter, cast, filter, take,
+    FilterKernel, FilterKernelAdapter, TakeKernel, TakeKernelAdapter, filter, take,
 };
 use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
 use vortex_error::VortexResult;
@@ -20,16 +20,9 @@ use crate::{DictArray, DictVTable};
 
 impl TakeKernel for DictVTable {
     fn take(&self, array: &DictArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
-        // TODO(joe): can we remove the cast and allow dict arrays to have nullable codes and values
         let codes = take(array.codes(), indices)?;
-        let values_dtype = array
-            .values()
-            .dtype()
-            .union_nullability(codes.dtype().nullability());
         // SAFETY: selecting codes doesn't change the invariants of DictArray
-        unsafe {
-            Ok(DictArray::new_unchecked(codes, cast(array.values(), &values_dtype)?).into_array())
-        }
+        Ok(unsafe { DictArray::new_unchecked(codes, array.values().clone()) }.into_array())
     }
 }
 

encodings/dict/src/ops.rs

@@ -19,30 +19,33 @@ impl OperationsVTable<DictVTable> for DictVTable {
             return if let Some(code) = code {
                 ConstantArray::new(array.values().scalar_at(*code), sliced_code.len()).into_array()
             } else {
-                let dtype = array.values().dtype().with_nullability(
-                    array.values().dtype().nullability() | array.codes().dtype().nullability(),
-                );
-                ConstantArray::new(Scalar::null(dtype), sliced_code.len()).to_array()
+                ConstantArray::new(Scalar::null(array.dtype().clone()), sliced_code.len())
+                    .to_array()
             };
         }
         // SAFETY: slicing the codes preserves invariants
         unsafe { DictArray::new_unchecked(sliced_code, array.values().clone()).into_array() }
     }
 
     fn scalar_at(array: &DictArray, index: usize) -> Scalar {
-        let dict_index: usize = array
-            .codes()
-            .scalar_at(index)
-            .as_ref()
-            .try_into()
-            .vortex_expect("code overflowed usize");
-        array.values().scalar_at(dict_index)
+        let Some(dict_index) = array.codes().scalar_at(index).as_primitive().as_::<usize>() else {
+            return Scalar::null(array.dtype().clone());
+        };
+
+        array
+            .values()
+            .scalar_at(dict_index)
+            .cast(array.dtype())
+            .vortex_expect("Array dtype will only differ by nullability")
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use vortex_array::IntoArray;
     use vortex_array::arrays::PrimitiveArray;
+    use vortex_buffer::buffer;
+    use vortex_dtype::Nullability;
     use vortex_scalar::Scalar;
 
     use crate::DictArray;
@@ -65,4 +68,19 @@ mod tests {
             dict.slice(1..2).as_constant()
         );
     }
+
+    #[test]
+    fn test_scalar_at_null_code() {
+        let dict = DictArray::try_new(
+            PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).to_array(),
+            buffer![1i32].into_array(),
+        )
+        .unwrap();
+
+        assert_eq!(dict.scalar_at(0), Scalar::null(dict.dtype().clone()));
+        assert_eq!(
+            dict.scalar_at(1),
+            Scalar::primitive(1, Nullability::Nullable)
+        );
+    }
 }