fix: ConstantArray#take handles nullable indices

Author: robert3005

fuzz/src/sort.rs

@@ -16,13 +16,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
             let mut opt_values = bool_array
                 .boolean_buffer()
                 .iter()
-                .zip(
-                    bool_array
-                        .validity_mask()
-                        .vortex_expect("Failed to get logical validity")
-                        .to_boolean_buffer()
-                        .iter(),
-                )
+                .zip(bool_array.validity_mask()?.to_boolean_buffer().iter())
                 .map(|(b, v)| v.then_some(b))
                 .collect::<Vec<_>>();
             opt_values.sort();

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

@@ -4,6 +4,7 @@ mod cast;
 mod compare;
 mod invert;
 mod search_sorted;
+mod take;
 
 use num_traits::{CheckedMul, ToPrimitive};
 use vortex_dtype::{NativePType, PType, match_each_native_ptype};
@@ -58,15 +59,15 @@ impl ComputeVTable for ConstantEncoding {
         Some(self)
     }
 
-    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
+    fn sum_fn(&self) -> Option<&dyn SumFn<&dyn Array>> {
         Some(self)
     }
 
-    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
+    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
         Some(self)
     }
 
-    fn sum_fn(&self) -> Option<&dyn SumFn<&dyn Array>> {
+    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
         Some(self)
     }
 }
@@ -77,12 +78,6 @@ impl ScalarAtFn<&ConstantArray> for ConstantEncoding {
     }
 }
 
-impl TakeFn<&ConstantArray> for ConstantEncoding {
-    fn take(&self, array: &ConstantArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
-        Ok(ConstantArray::new(array.scalar().clone(), indices.len()).into_array())
-    }
-}
-
 impl SliceFn<&ConstantArray> for ConstantEncoding {
     fn slice(&self, array: &ConstantArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
         Ok(ConstantArray::new(array.scalar().clone(), stop - start).into_array())

vortex-array/src/arrays/constant/compute/take.rs

@@ -0,0 +1,66 @@
+use vortex_error::VortexResult;
+use vortex_mask::AllOr;
+use vortex_scalar::Scalar;
+
+use crate::arrays::{ConstantArray, ConstantEncoding};
+use crate::builders::{ArrayBuilderExt, builder_with_capacity};
+use crate::compute::TakeFn;
+use crate::{Array, ArrayRef};
+
+impl TakeFn<&ConstantArray> for ConstantEncoding {
+    fn take(&self, array: &ConstantArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
+        match indices.validity_mask()?.boolean_buffer() {
+            AllOr::All => {
+                Ok(ConstantArray::new(array.scalar().clone(), indices.len()).into_array())
+            }
+            AllOr::None => Ok(ConstantArray::new(
+                Scalar::null(array.dtype().clone()),
+                indices.len(),
+            )
+            .into_array()),
+            AllOr::Some(v) => {
+                let mut result_builder =
+                    builder_with_capacity(&array.dtype().as_nullable(), indices.len());
+                for valid in v.iter() {
+                    if valid {
+                        result_builder.append_scalar_value(array.scalar().value().clone())?;
+                    } else {
+                        result_builder.append_null();
+                    }
+                }
+                Ok(result_builder.finish())
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+    use vortex_mask::AllOr;
+
+    use crate::arrays::{ConstantArray, PrimitiveArray};
+    use crate::compute::take;
+    use crate::validity::Validity;
+    use crate::{Array, ToCanonical};
+
+    #[test]
+    fn take_nullable_indices() {
+        let array = ConstantArray::new(42, 10).to_array();
+        let taken = take(
+            &array,
+            &PrimitiveArray::new(
+                buffer![0, 5, 7],
+                Validity::from_iter(vec![false, true, false]),
+            )
+            .into_array(),
+        )
+        .unwrap();
+        let valid_indices: &[usize] = &[1usize];
+        assert_eq!(taken.to_primitive().unwrap().as_slice::<i32>(), &[0, 42, 0]);
+        assert_eq!(
+            taken.validity_mask().unwrap().indices(),
+            AllOr::Some(valid_indices)
+        );
+    }
+}

vortex-array/src/validity.rs

@@ -147,7 +147,7 @@ impl Validity {
             Self::AllInvalid => Ok(Self::AllInvalid),
             Self::Array(is_valid) => {
                 let maybe_is_valid = take(is_valid, indices)?;
-                // Null indices invalidite that position.
+                // Null indices invalidate that position.
                 let is_valid = fill_null(&maybe_is_valid, Scalar::from(false))?;
                 Ok(Self::Array(is_valid))
             }