fix[vortex-array]: struct validity + runend validity (#3645)

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

Author: joseph-isaacs

encodings/runend/benches/run_end_filter.rs

@@ -2,6 +2,7 @@
 #![allow(clippy::cast_possible_truncation)]
 
 use divan::Bencher;
+use vortex_array::validity::Validity;
 use vortex_array::{Array, IntoArray};
 use vortex_buffer::Buffer;
 use vortex_mask::Mask;
@@ -42,7 +43,9 @@ fn take_indices(bencher: Bencher, (n, run_step, filter_density): (usize, usize,
 
     bencher
         .with_inputs(|| (&array, indices))
-        .bench_refs(|(array, indices)| take_indices_unchecked(array, indices).unwrap());
+        .bench_refs(|(array, indices)| {
+            take_indices_unchecked(array, indices, &Validity::NonNullable).unwrap()
+        });
 }
 
 #[divan::bench(args = BENCH_ARGS)]

encodings/runend/src/compute/filter.rs

@@ -27,7 +27,7 @@ impl FilterKernel for RunEndVTable {
 
                 if runs_ratio < FILTER_TAKE_THRESHOLD || mask_values.true_count() < 25 {
                     // This strategy is directly proportional to the number of indices.
-                    take_indices_unchecked(array, mask_values.indices())
+                    take_indices_unchecked(array, mask_values.indices(), &Validity::NonNullable)
                 } else {
                     // This strategy ends up being close to fixed cost based on the number of runs,
                     // rather than the number of indices.

encodings/runend/src/compute/take.rs

@@ -1,7 +1,10 @@
 use num_traits::{AsPrimitive, NumCast};
+use vortex_array::arrays::PrimitiveArray;
 use vortex_array::compute::{TakeKernel, TakeKernelAdapter, take};
 use vortex_array::search_sorted::{SearchResult, SearchSorted, SearchSortedSide};
-use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
+use vortex_array::validity::Validity;
+use vortex_array::vtable::ValidityHelper;
+use vortex_array::{Array, ArrayRef, ToCanonical, register_kernel};
 use vortex_buffer::Buffer;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::{VortexResult, vortex_bail};
@@ -28,7 +31,7 @@ impl TakeKernel for RunEndVTable {
                 .collect::<VortexResult<Vec<_>>>()?
         });
 
-        take_indices_unchecked(array, &checked_indices)
+        take_indices_unchecked(array, &checked_indices, primitive_indices.validity())
     }
 }
 
@@ -38,13 +41,15 @@ register_kernel!(TakeKernelAdapter(RunEndVTable).lift());
 pub fn take_indices_unchecked<T: AsPrimitive<usize>>(
     array: &RunEndArray,
     indices: &[T],
+    validity: &Validity,
 ) -> VortexResult<ArrayRef> {
     let ends = array.ends().to_primitive()?;
     let ends_len = ends.len();
 
+    // TODO(joe): use the validity mask to skip search sorted.
     let physical_indices = match_each_integer_ptype!(ends.ptype(), |I| {
         let end_slices = ends.as_slice::<I>();
-        indices
+        let buffer = indices
             .iter()
             .map(|idx| idx.as_() + array.offset())
             .map(|idx| {
@@ -57,18 +62,21 @@ pub fn take_indices_unchecked<T: AsPrimitive<usize>>(
                 }
             })
             .map(|result| result.to_ends_index(ends_len) as u64)
-            .collect::<Buffer<u64>>()
-            .into_array()
+            .collect::<Buffer<u64>>();
+
+        PrimitiveArray::new(buffer, validity.clone())
     });
 
-    take(array.values(), &physical_indices)
+    take(array.values(), physical_indices.as_ref())
 }
 
 #[cfg(test)]
 mod test {
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::compute::take;
     use vortex_array::{Array, IntoArray, ToCanonical};
+    use vortex_dtype::{DType, Nullability, PType};
+    use vortex_scalar::{Scalar, ScalarValue};
 
     use crate::RunEndArray;
 
@@ -126,4 +134,25 @@ mod test {
         assert_eq!(taken.scalar_at(1).unwrap(), 2.into());
         assert_eq!(taken.scalar_at(2).unwrap(), 5.into());
     }
+
+    #[test]
+    fn ree_take_nullable() {
+        let taken = take(
+            ree_array().as_ref(),
+            PrimitiveArray::from_option_iter([Some(1), None]).as_ref(),
+        )
+        .unwrap();
+
+        assert_eq!(
+            taken.scalar_at(0).unwrap(),
+            Scalar::new(
+                DType::Primitive(PType::I32, Nullability::Nullable),
+                ScalarValue::from(1i32)
+            )
+        );
+        assert_eq!(
+            taken.scalar_at(1).unwrap(),
+            Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable))
+        );
+    }
 }

vortex-array/src/arrays/struct_/compute/mod.rs

@@ -3,24 +3,43 @@ mod filter;
 mod mask;
 
 use itertools::Itertools;
+use vortex_dtype::Nullability::NonNullable;
 use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
 
 use crate::arrays::StructVTable;
 use crate::arrays::struct_::StructArray;
 use crate::compute::{
     IsConstantKernel, IsConstantKernelAdapter, IsConstantOpts, MinMaxKernel, MinMaxKernelAdapter,
-    MinMaxResult, TakeKernel, TakeKernelAdapter, is_constant_opts, take,
+    MinMaxResult, TakeKernel, TakeKernelAdapter, fill_null, is_constant_opts, take,
 };
+use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 use crate::{Array, ArrayRef, IntoArray, register_kernel};
 
 impl TakeKernel for StructVTable {
     fn take(&self, array: &StructArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
+        // If the struct array is empty then the indices must be all null, otherwise it will access
+        // an out of bounds element
+        if array.is_empty() {
+            return StructArray::try_new_with_dtype(
+                array.fields().to_vec(),
+                array.struct_fields().clone(),
+                indices.len(),
+                Validity::AllInvalid,
+            )
+            .map(StructArray::into_array);
+        }
+        // The validity is applied to the struct validity,
+        let inner_indices = &fill_null(
+            indices,
+            &Scalar::default_value(indices.dtype().with_nullability(NonNullable)),
+        )?;
         StructArray::try_new_with_dtype(
             array
                 .fields()
                 .iter()
-                .map(|field| take(field, indices))
+                .map(|field| take(field, inner_indices))
                 .try_collect()?,
             array.struct_fields().clone(),
             indices.len(),
@@ -71,13 +90,15 @@ register_kernel!(IsConstantKernelAdapter(StructVTable).lift());
 mod tests {
     use std::sync::Arc;
 
+    use Nullability::{NonNullable, Nullable};
     use vortex_buffer::buffer;
     use vortex_dtype::{DType, FieldNames, Nullability, PType, StructFields};
     use vortex_mask::Mask;
+    use vortex_scalar::Scalar;
 
     use crate::arrays::{BoolArray, BooleanBuffer, PrimitiveArray, StructArray, VarBinArray};
     use crate::compute::conformance::mask::test_mask;
-    use crate::compute::{cast, filter};
+    use crate::compute::{cast, filter, take};
     use crate::validity::Validity;
     use crate::{Array, IntoArray as _};
 
@@ -92,6 +113,52 @@ mod tests {
         assert_eq!(filtered.len(), 5);
     }
 
+    #[test]
+    fn take_empty_struct() {
+        let struct_arr =
+            StructArray::try_new(vec![].into(), vec![], 10, Validity::NonNullable).unwrap();
+        let indices = PrimitiveArray::from_option_iter([Some(1), None]);
+        let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap();
+        assert_eq!(taken.len(), 2);
+
+        assert_eq!(
+            taken.scalar_at(0).unwrap(),
+            Scalar::struct_(
+                DType::Struct(Arc::new(StructFields::new([].into(), vec![])), Nullable),
+                vec![]
+            )
+        );
+        assert_eq!(
+            taken.scalar_at(1).unwrap(),
+            Scalar::null(DType::Struct(
+                Arc::new(StructFields::new([].into(), vec![])),
+                Nullable
+            ))
+        );
+    }
+
+    #[test]
+    fn take_field_struct() {
+        let struct_arr =
+            StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).to_array())])
+                .unwrap();
+        let indices = PrimitiveArray::from_option_iter([Some(1), None]);
+        let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap();
+        assert_eq!(taken.len(), 2);
+
+        assert_eq!(
+            taken.scalar_at(0).unwrap(),
+            Scalar::struct_(
+                struct_arr.dtype().union_nullability(Nullable),
+                vec![Scalar::primitive(1, NonNullable)],
+            )
+        );
+        assert_eq!(
+            taken.scalar_at(1).unwrap(),
+            Scalar::null(struct_arr.dtype().union_nullability(Nullable),)
+        );
+    }
+
     #[test]
     fn filter_empty_struct_with_empty_filter() {
         let struct_arr =
@@ -114,7 +181,7 @@ mod tests {
         let xs = buffer![0i64, 1, 2, 3, 4].into_array();
         let ys = VarBinArray::from_iter(
             [Some("a"), Some("b"), None, Some("d"), None],
-            DType::Utf8(Nullability::Nullable),
+            DType::Utf8(Nullable),
         )
         .into_array();
         let zs =
@@ -148,17 +215,13 @@ mod tests {
         let array = StructArray::try_new(vec![].into(), vec![], 5, Validity::NonNullable)
             .unwrap()
             .into_array();
-        let non_nullable_dtype = DType::Struct(
-            Arc::from(StructFields::new([].into(), vec![])),
-            Nullability::NonNullable,
-        );
+        let non_nullable_dtype =
+            DType::Struct(Arc::from(StructFields::new([].into(), vec![])), NonNullable);
         let casted = cast(&array, &non_nullable_dtype).unwrap();
         assert_eq!(casted.dtype(), &non_nullable_dtype);
 
-        let nullable_dtype = DType::Struct(
-            Arc::from(StructFields::new([].into(), vec![])),
-            Nullability::Nullable,
-        );
+        let nullable_dtype =
+            DType::Struct(Arc::from(StructFields::new([].into(), vec![])), Nullable);
         let casted = cast(&array, &nullable_dtype).unwrap();
         assert_eq!(casted.dtype(), &nullable_dtype);
     }
@@ -177,7 +240,7 @@ mod tests {
         )
         .unwrap();
 
-        let tu8 = DType::Primitive(PType::U8, Nullability::NonNullable);
+        let tu8 = DType::Primitive(PType::U8, NonNullable);
 
         let result = cast(
             array.as_ref(),
@@ -186,7 +249,7 @@ mod tests {
                     FieldNames::from(["ys".into(), "xs".into(), "zs".into()]),
                     vec![tu8.clone(), tu8.clone(), tu8],
                 )),
-                Nullability::NonNullable,
+                NonNullable,
             ),
         );
         assert!(
@@ -201,10 +264,7 @@ mod tests {
     #[test]
     fn test_cast_complex_struct() {
         let xs = PrimitiveArray::from_option_iter([Some(0i64), Some(1), Some(2), Some(3), Some(4)]);
-        let ys = VarBinArray::from_vec(
-            vec!["a", "b", "c", "d", "e"],
-            DType::Utf8(Nullability::Nullable),
-        );
+        let ys = VarBinArray::from_vec(vec!["a", "b", "c", "d", "e"], DType::Utf8(Nullable));
         let zs = BoolArray::new(
             BooleanBuffer::from_iter([true, true, false, false, true]),
             Validity::AllValid,
@@ -241,17 +301,17 @@ mod tests {
                         Arc::from(StructFields::new(
                             ["left".into(), "right".into()].into(),
                             vec![
-                                DType::Primitive(PType::I64, Nullability::NonNullable),
-                                DType::Primitive(PType::I64, Nullability::Nullable),
+                                DType::Primitive(PType::I64, NonNullable),
+                                DType::Primitive(PType::I64, Nullable),
                             ],
                         )),
-                        Nullability::Nullable,
+                        Nullable,
                     ),
-                    DType::Utf8(Nullability::Nullable),
-                    DType::Bool(Nullability::Nullable),
+                    DType::Utf8(Nullable),
+                    DType::Bool(Nullable),
                 ],
             )),
-            Nullability::Nullable,
+            Nullable,
         );
         let casted = cast(&fully_nullable_array, &non_null_xs_right).unwrap();
         assert_eq!(casted.dtype(), &non_null_xs_right);
@@ -264,17 +324,17 @@ mod tests {
                         Arc::from(StructFields::new(
                             ["left".into(), "right".into()].into(),
                             vec![
-                                DType::Primitive(PType::I64, Nullability::Nullable),
-                                DType::Primitive(PType::I64, Nullability::Nullable),
+                                DType::Primitive(PType::I64, Nullable),
+                                DType::Primitive(PType::I64, Nullable),
                             ],
                         )),
-                        Nullability::NonNullable,
+                        NonNullable,
                     ),
-                    DType::Utf8(Nullability::Nullable),
-                    DType::Bool(Nullability::Nullable),
+                    DType::Utf8(Nullable),
+                    DType::Bool(Nullable),
                 ],
             )),
-            Nullability::Nullable,
+            Nullable,
         );
         let casted = cast(&fully_nullable_array, &non_null_xs).unwrap();
         assert_eq!(casted.dtype(), &non_null_xs);