fix: Validity::to_logical asserts mask length equals argument length (#2397)

I added an assertion to validity to_logical and fixed a bug which it
found in VarBinViewArray's `take_into`.

Author: danking

encodings/sparse/src/compute/take.rs

@@ -8,7 +8,12 @@ use crate::{SparseArray, SparseEncoding};
 impl TakeFn<SparseArray> for SparseEncoding {
     fn take(&self, array: &SparseArray, take_indices: &Array) -> VortexResult<Array> {
         let Some(new_patches) = array.patches().take(take_indices)? else {
-            return Ok(ConstantArray::new(array.fill_scalar(), take_indices.len()).into_array());
+            let result_nullability =
+                array.dtype().nullability() | take_indices.dtype().nullability();
+            let result_fill_scalar = array
+                .fill_scalar()
+                .cast(&array.dtype().with_nullability(result_nullability))?;
+            return Ok(ConstantArray::new(result_fill_scalar, take_indices.len()).into_array());
         };
 
         SparseArray::try_new_from_patches(new_patches, take_indices.len(), array.fill_scalar())

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

@@ -78,7 +78,7 @@ impl TakeFn<VarBinViewArray> for VarBinViewEncoding {
         // min-max valid range.
         // TODO(joe): impl validity_mask take
         let validity = array.validity().take(indices)?;
-        let mask = validity.to_logical(array.len())?;
+        let mask = validity.to_logical(indices.len())?;
         let indices = indices.clone().into_primitive()?;
 
         match_each_integer_ptype!(indices.ptype(), |$I| {

vortex-array/src/compute/take.rs

@@ -113,12 +113,15 @@ pub fn take(array: impl AsRef<Array>, indices: impl AsRef<Array>) -> VortexResul
     #[cfg(debug_assertions)]
     {
         // If either the indices or the array are nullable, the result should be nullable.
-        let expected_nullability =
-            (indices.dtype().is_nullable() || array.dtype().is_nullable()).into();
+        let expected_nullability = indices.dtype().nullability() | array.dtype().nullability();
         assert_eq!(
             taken.dtype(),
             &array.dtype().with_nullability(expected_nullability),
-            "Take result should be nullable if either the indices or the array are nullable"
+            "Take result ({}) should be nullable if either the indices ({}) or the array ({}) are nullable. ({})",
+            taken.dtype(),
+            indices.dtype().nullability().verbose_display(),
+            array.dtype().nullability().verbose_display(),
+            array.encoding(),
         );
     }
 
@@ -136,12 +139,15 @@ pub fn take_into(
     #[cfg(debug_assertions)]
     {
         // If either the indices or the array are nullable, the result should be nullable.
-        let expected_nullability =
-            (indices.dtype().is_nullable() || array.dtype().is_nullable()).into();
+        let expected_nullability = indices.dtype().nullability() | array.dtype().nullability();
         assert_eq!(
             builder.dtype(),
             &array.dtype().with_nullability(expected_nullability),
-            "Take_into result should be nullable if either the indices or the array are nullable"
+            "Take_into result ({}) should be nullable if, and only if, either the indices ({}) or the array ({}) are nullable. ({})",
+            builder.dtype(),
+            indices.dtype().nullability().verbose_display(),
+            array.dtype().nullability().verbose_display(),
+            array.encoding(),
         );
     }
 

vortex-array/src/validity.rs

@@ -300,7 +300,16 @@ impl Validity {
         Ok(match self {
             Self::NonNullable | Self::AllValid => Mask::AllTrue(length),
             Self::AllInvalid => Mask::AllFalse(length),
-            Self::Array(a) => Mask::try_from(a.clone())?,
+            Self::Array(is_valid) => {
+                assert_eq!(
+                    is_valid.len(),
+                    length,
+                    "Validity::Array length must equal to_logical's argument: {}, {}.",
+                    is_valid.len(),
+                    length,
+                );
+                Mask::try_from(is_valid.clone())?
+            }
         })
     }
 

vortex-dtype/src/nullability.rs

@@ -11,6 +11,18 @@ pub enum Nullability {
     Nullable,
 }
 
+impl Nullability {
+    /// A self-describing displayed form.
+    ///
+    /// The usual Display renders [Nullability::NonNullable] as the empty string.
+    pub fn verbose_display(&self) -> impl Display {
+        match self {
+            Nullability::NonNullable => "NonNullable",
+            Nullability::Nullable => "Nullable",
+        }
+    }
+}
+
 impl BitOr for Nullability {
     type Output = Nullability;