fix: take_into mask length must equal _indices_ length (#2396)

I added a test that catches this behavior. In particular, the indices
must differ in length from the array and also it must be either all
valid or all invalid.

I additionally changed `Validity::to_logical` to assert the length
matches the array length. Since mixed validity is more common than the
other two, I hope this new assertion will surface these kinds of issues
in more tests.

An example of a failure on `develop` with these new tests:
```
---- array::primitive::compute::take::test::test_take_into stdout ----
thread 'array::primitive::compute::take::test::test_take_into' panicked at vortex-array/src/builders/primitive.rs:98:9:
assertion `left == right` failed: null count must equal value count
  left: 5
 right: 3
stack backtrace:
   0: rust_begin_unwind
             at /rustc/409998c4e8cae45344fd434b358b697cc93870d0/library/std/src/panicking.rs:676:5
   1: core::panicking::panic_fmt
             at /rustc/409998c4e8cae45344fd434b358b697cc93870d0/library/core/src/panicking.rs:75:14
   2: core::panicking::assert_failed_inner
   3: core::panicking::assert_failed
             at /rustc/409998c4e8cae45344fd434b358b697cc93870d0/library/core/src/panicking.rs:364:5
   4: vortex_array::builders::primitive::PrimitiveBuilder<T>::finish_into_primitive
             at ./src/builders/primitive.rs:98:9
   5: <vortex_array::builders::primitive::PrimitiveBuilder<T> as vortex_array::builders::ArrayBuilder>::finish
             at ./src/builders/primitive.rs:176:9
   6: vortex_array::array::primitive::compute::take::test::test_take_into
             at ./src/array/primitive/compute/take.rs:150:22
   7: vortex_array::array::primitive::compute::take::test::test_take_into::{{closure}}
             at ./src/array/primitive/compute/take.rs:142:24
```

Author: danking

vortex-array/src/array/primitive/compute/take.rs

@@ -50,7 +50,7 @@ impl TakeFn<PrimitiveArray> for PrimitiveEncoding {
         let indices = indices.clone().into_primitive()?;
         // TODO(joe): impl take over mask and use `Array::validity_mask`, instead of `validity()`.
         let validity = array.validity().take(indices.as_ref())?;
-        let mask = validity.to_logical(array.len())?;
+        let mask = validity.to_logical(indices.len())?;
 
         match_each_native_ptype!(array.ptype(), |$T| {
             match_each_integer_ptype!(indices.ptype(), |$I| {
@@ -66,6 +66,8 @@ fn take_into_impl<T: NativePType, I: NativePType + AsPrimitive<usize>>(
     mask: Mask,
     builder: &mut dyn ArrayBuilder,
 ) -> VortexResult<()> {
+    assert_eq!(indices.len(), mask.len());
+
     let array = array.as_slice::<T>();
     let indices = indices.as_slice::<I>();
     let builder = builder
@@ -103,11 +105,13 @@ unsafe fn take_primitive_unchecked<T: NativePType, I: NativePType + AsPrimitive<
 #[cfg(test)]
 mod test {
     use vortex_buffer::buffer;
+    use vortex_dtype::Nullability;
     use vortex_scalar::Scalar;
 
     use crate::array::primitive::compute::take::take_primitive;
     use crate::array::{BoolArray, PrimitiveArray};
-    use crate::compute::{scalar_at, take};
+    use crate::builders::{ArrayBuilder as _, PrimitiveBuilder};
+    use crate::compute::{scalar_at, take, take_into};
     use crate::validity::Validity;
     use crate::IntoArray as _;
 
@@ -135,4 +139,50 @@ mod test {
         // the third index is null
         assert_eq!(scalar_at(&actual, 2).unwrap(), Scalar::null_typed::<i32>());
     }
+
+    #[test]
+    fn test_take_into() {
+        let values = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::NonNullable);
+        let all_valid_indices = PrimitiveArray::new(
+            buffer![0, 3, 4],
+            Validity::Array(BoolArray::from_iter([true, true, true]).into_array()),
+        );
+        let mut builder = PrimitiveBuilder::<i32>::new(Nullability::Nullable);
+        take_into(&values, all_valid_indices, &mut builder).unwrap();
+        let actual = builder.finish().unwrap();
+        assert_eq!(scalar_at(&actual, 0).unwrap(), Scalar::from(Some(1)));
+        assert_eq!(scalar_at(&actual, 1).unwrap(), Scalar::from(Some(4)));
+        assert_eq!(scalar_at(&actual, 2).unwrap(), Scalar::from(Some(5)));
+
+        let mixed_valid_indices = PrimitiveArray::new(
+            buffer![0, 3, 4],
+            Validity::Array(BoolArray::from_iter([true, true, false]).into_array()),
+        );
+        let mut builder = PrimitiveBuilder::<i32>::new(Nullability::Nullable);
+        take_into(&values, mixed_valid_indices, &mut builder).unwrap();
+        let actual = builder.finish().unwrap();
+        assert_eq!(scalar_at(&actual, 0).unwrap(), Scalar::from(Some(1)));
+        assert_eq!(scalar_at(&actual, 1).unwrap(), Scalar::from(Some(4)));
+        // the third index is null
+        assert_eq!(scalar_at(&actual, 2).unwrap(), Scalar::null_typed::<i32>());
+
+        let all_invalid_indices = PrimitiveArray::new(
+            buffer![0, 3, 4],
+            Validity::Array(BoolArray::from_iter([false, false, false]).into_array()),
+        );
+        let mut builder = PrimitiveBuilder::<i32>::new(Nullability::Nullable);
+        take_into(&values, all_invalid_indices, &mut builder).unwrap();
+        let actual = builder.finish().unwrap();
+        assert_eq!(scalar_at(&actual, 0).unwrap(), Scalar::null_typed::<i32>());
+        assert_eq!(scalar_at(&actual, 1).unwrap(), Scalar::null_typed::<i32>());
+        assert_eq!(scalar_at(&actual, 2).unwrap(), Scalar::null_typed::<i32>());
+
+        let non_null_indices = PrimitiveArray::new(buffer![0, 3, 4], Validity::NonNullable);
+        let mut builder = PrimitiveBuilder::<i32>::new(Nullability::NonNullable);
+        take_into(&values, non_null_indices, &mut builder).unwrap();
+        let actual = builder.finish().unwrap();
+        assert_eq!(scalar_at(&actual, 0).unwrap(), Scalar::from(1));
+        assert_eq!(scalar_at(&actual, 1).unwrap(), Scalar::from(4));
+        assert_eq!(scalar_at(&actual, 2).unwrap(), Scalar::from(5));
+    }
 }

vortex-array/src/compute/take.rs

@@ -224,12 +224,14 @@ fn take_into_impl(
     indices: &Array,
     builder: &mut dyn ArrayBuilder,
 ) -> VortexResult<()> {
-    if array.dtype() != builder.dtype() {
+    let result_nullability = array.dtype().nullability() | indices.dtype().nullability();
+    let result_dtype = array.dtype().with_nullability(result_nullability);
+    if &result_dtype != builder.dtype() {
         vortex_bail!(
-            "TakeIntoFn {} had a builder with a different dtype {} to the array dtype {}",
+            "TakeIntoFn {} had a builder with a different dtype {} to the resulting array dtype {}",
             array.encoding(),
-            array.dtype(),
-            builder.dtype()
+            builder.dtype(),
+            result_dtype,
         );
     }
     if let Some(take_fn) = array.vtable().take_fn() {

vortex-dtype/src/nullability.rs

@@ -1,4 +1,5 @@
 use std::fmt::{Display, Formatter};
+use std::ops::BitOr;
 
 /// Whether an instance of a DType can be `null or not
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
@@ -10,6 +11,17 @@ pub enum Nullability {
     Nullable,
 }
 
+impl BitOr for Nullability {
+    type Output = Nullability;
+
+    fn bitor(self, rhs: Self) -> Self::Output {
+        match (self, rhs) {
+            (Self::NonNullable, Self::NonNullable) => Self::NonNullable,
+            _ => Self::Nullable,
+        }
+    }
+}
+
 impl From<bool> for Nullability {
     fn from(value: bool) -> Self {
         if value {