bug[vortex-array]: take compare nullable bugs (#3628)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

encodings/dict/src/compute/mod.rs

@@ -7,7 +7,7 @@ mod like;
 mod min_max;
 
 use vortex_array::compute::{
-    FilterKernel, FilterKernelAdapter, TakeKernel, TakeKernelAdapter, filter, take,
+    FilterKernel, FilterKernelAdapter, TakeKernel, TakeKernelAdapter, cast, filter, take,
 };
 use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
 use vortex_error::VortexResult;
@@ -17,8 +17,13 @@ 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)?;
-        DictArray::try_new(codes, array.values().clone()).map(|a| a.into_array())
+        let values_dtype = array
+            .values()
+            .dtype()
+            .union_nullability(codes.dtype().nullability());
+        DictArray::try_new(codes, cast(array.values(), &values_dtype)?).map(|a| a.into_array())
     }
 }
 
@@ -38,8 +43,9 @@ mod test {
     use vortex_array::accessor::ArrayAccessor;
     use vortex_array::arrays::{ConstantArray, PrimitiveArray, VarBinArray, VarBinViewArray};
     use vortex_array::compute::conformance::mask::test_mask;
-    use vortex_array::compute::{Operator, compare};
+    use vortex_array::compute::{Operator, compare, take};
     use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
+    use vortex_dtype::PType::I32;
     use vortex_dtype::{DType, Nullability};
     use vortex_scalar::Scalar;
 
@@ -184,4 +190,19 @@ mod test {
         .unwrap();
         test_mask(array.as_ref());
     }
+
+    #[test]
+    fn test_take_dict() {
+        let array = dict_encode(PrimitiveArray::from_iter([1, 2]).as_ref()).unwrap();
+
+        assert_eq!(
+            take(
+                array.as_ref(),
+                PrimitiveArray::from_option_iter([Option::<i32>::None]).as_ref()
+            )
+            .unwrap()
+            .dtype(),
+            &DType::Primitive(I32, Nullability::Nullable)
+        );
+    }
 }

encodings/fsst/src/compute/mod.rs

@@ -13,7 +13,10 @@ impl TakeKernel for FSSTVTable {
     // Take on an FSSTArray is a simple take on the codes array.
     fn take(&self, array: &FSSTArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
         Ok(FSSTArray::try_new(
-            array.dtype().clone(),
+            array
+                .dtype()
+                .clone()
+                .union_nullability(indices.dtype().nullability()),
             array.symbols().clone(),
             array.symbol_lengths().clone(),
             take(array.codes().as_ref(), indices)?
@@ -32,3 +35,33 @@ impl TakeKernel for FSSTVTable {
 }
 
 register_kernel!(TakeKernelAdapter(FSSTVTable).lift());
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::arrays::{PrimitiveArray, VarBinArray};
+    use vortex_array::compute::take;
+    use vortex_dtype::{DType, Nullability};
+
+    use crate::{fsst_compress, fsst_train_compressor};
+
+    #[test]
+    fn test_take_null() {
+        let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable));
+        let compr = fsst_train_compressor(arr.as_ref()).unwrap();
+        let fsst = fsst_compress(arr.as_ref(), &compr).unwrap();
+
+        let idx1: PrimitiveArray = (0..1).collect();
+
+        assert_eq!(
+            take(fsst.as_ref(), idx1.as_ref()).unwrap().dtype(),
+            &DType::Utf8(Nullability::NonNullable)
+        );
+
+        let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]);
+
+        assert_eq!(
+            take(fsst.as_ref(), idx2.as_ref()).unwrap().dtype(),
+            &DType::Utf8(Nullability::Nullable)
+        );
+    }
+}

vortex-array/src/arrays/varbin/compute/compare.rs

@@ -52,7 +52,13 @@ impl CompareKernel for VarBinVTable {
                 };
 
                 return Ok(Some(
-                    BoolArray::new(buffer, lhs.validity().clone()).into_array(),
+                    BoolArray::new(
+                        buffer,
+                        lhs.validity()
+                            .clone()
+                            .union_nullability(rhs.dtype().nullability()),
+                    )
+                    .into_array(),
                 ));
             }
 
@@ -177,3 +183,27 @@ mod test {
         );
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_dtype::{DType, Nullability};
+    use vortex_scalar::Scalar;
+
+    use crate::Array;
+    use crate::arrays::{ConstantArray, VarBinArray};
+    use crate::compute::{Operator, compare};
+
+    #[test]
+    fn test_null_compare() {
+        let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable));
+
+        let const_ = ConstantArray::new(Scalar::utf8("", Nullability::Nullable), 1);
+
+        assert_eq!(
+            compare(arr.as_ref(), const_.as_ref(), Operator::Eq)
+                .unwrap()
+                .dtype(),
+            &DType::Bool(Nullability::Nullable)
+        );
+    }
+}

vortex-array/src/compute/mod.rs

@@ -113,10 +113,14 @@ impl ComputeFn {
 
         if output.dtype() != &expected_dtype {
             vortex_bail!(
-                "Internal error: compute function {} returned a result of type {} but expected {}",
+                "Internal error: compute function {} returned a result of type {} but expected {}\n{}",
                 self.id,
                 output.dtype(),
                 &expected_dtype,
+                args.inputs
+                    .iter()
+                    .filter_map(|input| input.array())
+                    .format_with(",", |array, f| f(&array.tree_display()))
             );
         }
         if output.len() != expected_len {