fix: ConstantArray#take handles nullable indices (#2631)

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,70 @@
+use vortex_error::VortexResult;
+use vortex_mask::{AllOr, Mask};
+use vortex_scalar::Scalar;
+
+use crate::arrays::{ConstantArray, ConstantEncoding};
+use crate::builders::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 arr = ConstantArray::new(array.scalar().clone(), indices.len()).into_array();
+
+                if array.scalar().is_null() {
+                    return Ok(arr);
+                }
+
+                let mut result_builder =
+                    builder_with_capacity(&array.dtype().as_nullable(), indices.len());
+                result_builder.extend_from_array(&arr)?;
+                result_builder.set_validity(Mask::from_buffer(v.clone()));
+                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>(),
+            &[42, 42, 42]
+        );
+        assert_eq!(
+            taken.validity_mask().unwrap().indices(),
+            AllOr::Some(valid_indices)
+        );
+    }
+}

vortex-array/src/builders/bool.rs

@@ -3,6 +3,7 @@ use std::any::Any;
 use arrow_buffer::BooleanBufferBuilder;
 use vortex_dtype::{DType, Nullability};
 use vortex_error::{VortexResult, vortex_bail};
+use vortex_mask::Mask;
 
 use crate::arrays::BoolArray;
 use crate::builders::ArrayBuilder;
@@ -85,6 +86,11 @@ impl ArrayBuilder for BoolBuilder {
         Ok(())
     }
 
+    fn set_validity(&mut self, validity: Mask) {
+        self.nulls = LazyNullBufferBuilder::new(validity.len());
+        self.nulls.append_validity_mask(validity);
+    }
+
     fn finish(&mut self) -> ArrayRef {
         assert_eq!(
             self.nulls.len(),

vortex-array/src/builders/extension.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 
 use vortex_dtype::{DType, ExtDType};
 use vortex_error::{VortexResult, vortex_bail};
+use vortex_mask::Mask;
 use vortex_scalar::ExtScalar;
 
 use crate::arrays::ExtensionArray;
@@ -82,6 +83,10 @@ impl ArrayBuilder for ExtensionBuilder {
         array.storage().append_to_builder(self.storage.as_mut())
     }
 
+    fn set_validity(&mut self, validity: Mask) {
+        self.storage.set_validity(validity);
+    }
+
     fn finish(&mut self) -> ArrayRef {
         let storage = self.storage.finish();
         ExtensionArray::new(self.ext_dtype(), storage).into_array()

vortex-array/src/builders/lazy_validity_builder.rs

@@ -52,7 +52,6 @@ impl LazyNullBufferBuilder {
             .append_n(n, false);
     }
 
-    #[allow(dead_code)]
     #[inline]
     pub fn append_null(&mut self) {
         self.materialize_if_needed();
@@ -62,7 +61,6 @@ impl LazyNullBufferBuilder {
             .append(false);
     }
 
-    #[allow(dead_code)]
     #[inline]
     pub fn append(&mut self, not_null: bool) {
         if not_null {

vortex-array/src/builders/list.rs

@@ -4,6 +4,7 @@ use std::sync::Arc;
 use vortex_dtype::Nullability::NonNullable;
 use vortex_dtype::{DType, NativePType, Nullability};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
+use vortex_mask::Mask;
 use vortex_scalar::{BinaryNumericOperator, ListScalar};
 
 use crate::arrays::{ConstantArray, ListArray, OffsetPType};
@@ -150,6 +151,11 @@ impl<O: OffsetPType> ArrayBuilder for ListBuilder<O> {
         Ok(())
     }
 
+    fn set_validity(&mut self, validity: Mask) {
+        self.nulls = LazyNullBufferBuilder::new(validity.len());
+        self.nulls.append_validity_mask(validity);
+    }
+
     fn finish(&mut self) -> ArrayRef {
         assert_eq!(
             self.index_builder.len(),

vortex-array/src/builders/mod.rs

@@ -17,6 +17,7 @@ pub use primitive::*;
 pub use varbinview::*;
 use vortex_dtype::{DType, match_each_native_ptype};
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_mask::Mask;
 use vortex_scalar::{
     BinaryScalar, BoolScalar, ExtScalar, ListScalar, PrimitiveScalar, Scalar, ScalarValue,
     StructScalar, Utf8Scalar,
@@ -57,6 +58,9 @@ pub trait ArrayBuilder: Send {
     /// Extends the array with the provided array, canonicalizing if necessary.
     fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()>;
 
+    /// Override builders validity with the one provided
+    fn set_validity(&mut self, validity: Mask);
+
     /// Constructs an Array from the builder components.
     ///
     /// # Panics

vortex-array/src/builders/null.rs

@@ -2,6 +2,7 @@ use std::any::Any;
 
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
+use vortex_mask::Mask;
 
 use crate::arrays::NullArray;
 use crate::builders::ArrayBuilder;
@@ -54,6 +55,10 @@ impl ArrayBuilder for NullBuilder {
         Ok(())
     }
 
+    fn set_validity(&mut self, validity: Mask) {
+        self.length = validity.len();
+    }
+
     fn finish(&mut self) -> ArrayRef {
         NullArray::new(self.length).into_array()
     }

vortex-array/src/builders/primitive.rs

@@ -178,6 +178,11 @@ impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
         Ok(())
     }
 
+    fn set_validity(&mut self, validity: Mask) {
+        self.nulls = LazyNullBufferBuilder::new(validity.len());
+        self.nulls.append_validity_mask(validity);
+    }
+
     fn finish(&mut self) -> ArrayRef {
         self.finish_into_primitive().into_array()
     }