Implement take for Mask and FromIterator<Mask> for Mask (#2689)

Author: robert3005

encodings/bytebool/src/array.rs

@@ -117,7 +117,7 @@ impl ArrayValidityImpl for ByteBoolArray {
     }
 
     fn _validity_mask(&self) -> VortexResult<Mask> {
-        self.validity.to_logical(self.len())
+        self.validity.to_mask(self.len())
     }
 }
 

encodings/datetime-parts/src/canonical.rs

@@ -131,7 +131,7 @@ mod test {
 
         assert_eq!(
             date_times.validity_mask().unwrap(),
-            validity.to_logical(date_times.len()).unwrap()
+            validity.to_mask(date_times.len()).unwrap()
         );
 
         let primitive_values = decode_to_temporal(&date_times)

encodings/fastlanes/src/bitpacking/compress.rs

@@ -16,7 +16,7 @@ use vortex_dtype::{
     match_each_unsigned_integer_ptype,
 };
 use vortex_error::{VortexExpect, VortexResult, vortex_bail};
-use vortex_mask::AllOr;
+use vortex_mask::{AllOr, Mask};
 use vortex_scalar::Scalar;
 
 use crate::BitPackedArray;
@@ -274,52 +274,52 @@ fn apply_patches<T: NativePType>(dst: &mut UninitRange<T>, patches: &Patches) ->
 
     let indices = indices.to_primitive()?;
     let values = values.to_primitive()?;
-    let validity = values.validity();
+    let validity = values.validity_mask()?;
     let values = values.as_slice::<T>();
     match_each_unsigned_integer_ptype!(indices.ptype(), |$P| {
-        insert_values_and_validity_at_indices::<T, $P>(
+        insert_values_and_validity_at_indices::<T, _>(
             dst,
-            indices,
+            indices.as_slice::<$P>(),
             values,
-            validity.clone(),
+            validity,
             indices_offset,
         )
-    })
+    });
+    Ok(())
 }
 
 fn insert_values_and_validity_at_indices<
     T: NativePType,
     IndexT: NativePType + AsPrimitive<usize>,
 >(
     dst: &mut UninitRange<T>,
-    indices: PrimitiveArray,
+    indices: &[IndexT],
     values: &[T],
-    validity: Validity,
+    values_validity: Mask,
     indices_offset: usize,
-) -> VortexResult<()> {
-    match validity {
-        Validity::NonNullable => {
-            for (compressed_index, decompressed_index) in
-                indices.as_slice::<IndexT>().iter().enumerate()
-            {
+) {
+    match values_validity {
+        Mask::AllTrue(_) => {
+            for (compressed_index, decompressed_index) in indices.iter().enumerate() {
                 dst[decompressed_index.as_() - indices_offset] =
                     MaybeUninit::new(values[compressed_index]);
             }
         }
-        _ => {
-            let validity = validity.to_logical(indices.len())?;
-            for (compressed_index, decompressed_index) in
-                indices.as_slice::<IndexT>().iter().enumerate()
-            {
+        Mask::AllFalse(_) => {
+            for (compressed_index, decompressed_index) in indices.iter().enumerate() {
                 let out_index = decompressed_index.as_() - indices_offset;
-                dst[decompressed_index.as_() - indices_offset] =
-                    MaybeUninit::new(values[compressed_index]);
-                dst.set_bit(out_index, validity.value(out_index));
+                dst[out_index] = MaybeUninit::new(values[compressed_index]);
+                dst.set_bit(out_index, false);
+            }
+        }
+        Mask::Values(vb) => {
+            for (compressed_index, decompressed_index) in indices.iter().enumerate() {
+                let out_index = decompressed_index.as_() - indices_offset;
+                dst[out_index] = MaybeUninit::new(values[compressed_index]);
+                dst.set_bit(out_index, vb.value(out_index));
             }
         }
     }
-
-    Ok(())
 }
 
 fn unpack_values_into<T: NativePType, UnsignedT: NativePType + BitPacking, F, G>(

encodings/fastlanes/src/bitpacking/mod.rs

@@ -279,7 +279,7 @@ impl ArrayValidityImpl for BitPackedArray {
     }
 
     fn _validity_mask(&self) -> VortexResult<Mask> {
-        self.validity.to_logical(self.len())
+        self.validity.to_mask(self.len())
     }
 }
 

encodings/fastlanes/src/delta/mod.rs

@@ -250,7 +250,7 @@ impl ArrayValidityImpl for DeltaArray {
     }
 
     fn _validity_mask(&self) -> VortexResult<Mask> {
-        self.validity.to_logical(self.len)
+        self.validity.to_mask(self.len)
     }
 }
 

vortex-array/src/arrays/bool/array.rs

@@ -152,7 +152,7 @@ impl ArrayValidityImpl for BoolArray {
 
     #[inline]
     fn _validity_mask(&self) -> VortexResult<Mask> {
-        self.validity.to_logical(self.len())
+        self.validity.to_mask(self.len())
     }
 }
 

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

@@ -1,12 +1,12 @@
 use arrow_buffer::BooleanBuffer;
 use itertools::Itertools;
 use num_traits::AsPrimitive;
-use vortex_dtype::{NativePType, match_each_integer_ptype};
+use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
 
-use crate::arrays::{BoolArray, BoolEncoding, ConstantArray, PrimitiveArray};
+use crate::arrays::{BoolArray, BoolEncoding, ConstantArray};
 use crate::builders::ArrayBuilder;
 use crate::compute::{TakeFn, fill_null};
 use crate::variants::PrimitiveArrayTrait;
@@ -27,7 +27,7 @@ impl TakeFn<&BoolArray> for BoolEncoding {
         };
         let indices_nulls_zeroed = indices_nulls_zeroed.to_primitive()?;
         let buffer = match_each_integer_ptype!(indices_nulls_zeroed.ptype(), |$I| {
-            take_valid_indices::<$I>(array, &indices_nulls_zeroed)
+            take_valid_indices(array.boolean_buffer(), indices_nulls_zeroed.as_slice::<$I>())
         });
 
         Ok(BoolArray::new(buffer, array.validity().take(indices)?).into_array())
@@ -43,30 +43,30 @@ impl TakeFn<&BoolArray> for BoolEncoding {
     }
 }
 
-fn take_valid_indices<I: AsPrimitive<usize> + NativePType>(
-    array: &BoolArray,
-    indices: &PrimitiveArray,
+fn take_valid_indices<I: AsPrimitive<usize>>(
+    bools: &BooleanBuffer,
+    indices: &[I],
 ) -> BooleanBuffer {
     // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth
     // the overhead to convert to a Vec<bool>.
-    if array.len() <= 4096 {
-        let bools = array.boolean_buffer().into_iter().collect_vec();
-        take_byte_bool(bools, indices.as_slice::<I>())
+    if bools.len() <= 4096 {
+        let bools = bools.into_iter().collect_vec();
+        take_byte_bool(bools, indices)
     } else {
-        take_bool(array.boolean_buffer(), indices.as_slice::<I>())
+        take_bool(bools, indices)
     }
 }
 
 fn take_byte_bool<I: AsPrimitive<usize>>(bools: Vec<bool>, indices: &[I]) -> BooleanBuffer {
     BooleanBuffer::collect_bool(indices.len(), |idx| {
-        bools[unsafe { (*indices.get_unchecked(idx)).as_() }]
+        bools[unsafe { indices.get_unchecked(idx).as_() }]
     })
 }
 
 fn take_bool<I: AsPrimitive<usize>>(bools: &BooleanBuffer, indices: &[I]) -> BooleanBuffer {
     BooleanBuffer::collect_bool(indices.len(), |idx| {
         // We can always take from the indices unchecked since collect_bool just iterates len.
-        bools.value(unsafe { (*indices.get_unchecked(idx)).as_() })
+        bools.value(unsafe { indices.get_unchecked(idx).as_() })
     })
 }
 

vortex-array/src/arrays/chunked/mod.rs

@@ -17,7 +17,6 @@ use crate::iter::{ArrayIterator, ArrayIteratorAdapter};
 use crate::nbytes::NBytes;
 use crate::stats::{ArrayStats, StatsSetRef};
 use crate::stream::{ArrayStream, ArrayStreamAdapter};
-use crate::validity::Validity;
 use crate::vtable::VTableRef;
 use crate::{Array, ArrayImpl, ArrayRef, ArrayStatisticsImpl, EmptyMetadata, Encoding, IntoArray};
 
@@ -240,14 +239,10 @@ impl ArrayValidityImpl for ChunkedArray {
     }
 
     fn _validity_mask(&self) -> VortexResult<Mask> {
-        // TODO(ngates): implement FromIterator<LogicalValidity> for LogicalValidity.
-        // TODO(ngates): or use a boolean array builder?
-        let validity: Validity = self
-            .chunks()
+        self.chunks()
             .iter()
             .map(|a| a.validity_mask())
-            .try_collect()?;
-        validity.to_logical(self.len())
+            .try_collect()
     }
 }
 

vortex-array/src/arrays/list/mod.rs

@@ -173,7 +173,7 @@ impl ArrayValidityImpl for ListArray {
     }
 
     fn _validity_mask(&self) -> VortexResult<Mask> {
-        self.validity.to_logical(self.len())
+        self.validity.to_mask(self.len())
     }
 }
 

vortex-array/src/arrays/primitive/compute/cast.rs

@@ -22,9 +22,7 @@ impl CastFn<&PrimitiveArray> for PrimitiveEncoding {
         } else if new_nullability == Nullability::Nullable {
             // from non-nullable to nullable
             array.validity().clone().into_nullable()
-        } else if new_nullability == Nullability::NonNullable
-            && array.validity().to_logical(array.len())?.all_true()
-        {
+        } else if new_nullability == Nullability::NonNullable && array.validity().all_valid()? {
             // from nullable but all valid, to non-nullable
             Validity::NonNullable
         } else {