Fix: masked dtype mismatch (#6008)

Fixes: #5989

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-array/src/arrays/masked/vtable/mod.rs

@@ -130,12 +130,23 @@ impl VTable for MaskedVTable {
     }
 
     fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<Canonical> {
-        if let Some(canonical) = execute_fast_path(array, ctx)? {
-            return Ok(canonical);
+        let validity_mask = array.validity_mask();
+
+        // Fast path: all masked means result is all nulls.
+        if validity_mask.all_false() {
+            return ConstantArray::new(Scalar::null(array.dtype().as_nullable()), array.len())
+                .into_array()
+                .execute::<Canonical>(ctx);
         }
 
+        // NB: We intentionally do NOT have a fast path for `validity_mask.all_true()`.
+        // `MaskedArray`'s dtype is always `Nullable`, but the child has `NonNullable` `DType` (by
+        // invariant). Simply returning the child's canonical would cause a dtype mismatch.
+        // While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for
+        // `AllTrue` masks (no data copying), so there's no benefit.
+
         let child = array.child().clone().execute::<Canonical>(ctx)?;
-        let canonical = mask_validity_canonical(child, &array.validity_mask())?;
+        let canonical = mask_validity_canonical(child, &validity_mask)?;
 
         vortex_ensure!(
             canonical.as_ref().dtype() == array.dtype(),
@@ -176,42 +187,18 @@ impl VTable for MaskedVTable {
     }
 }
 
-/// Check for fast-path execution conditions.
-pub(super) fn execute_fast_path(
-    array: &MaskedArray,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<Option<Canonical>> {
-    let validity_mask = array.validity_mask();
-
-    // All valid - no masking needed
-    if validity_mask.all_true() {
-        return Ok(Some(array.child.clone().execute(ctx)?));
-    }
-
-    // All masked - result is all nulls
-    if validity_mask.all_false() {
-        return Ok(Some(
-            ConstantArray::new(Scalar::null(array.dtype().as_nullable()), array.len())
-                .into_array()
-                .execute::<Canonical>(ctx)?,
-        ));
-    }
-
-    // Child is already all nulls - masking has no effect
-    if array.child.all_invalid() {
-        return Ok(Some(array.child.clone().execute(ctx)?));
-    }
-
-    Ok(None)
-}
-
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
     use vortex_buffer::ByteBufferMut;
+    use vortex_dtype::Nullability;
+    use vortex_error::VortexError;
 
     use crate::ArrayContext;
+    use crate::Canonical;
     use crate::IntoArray;
+    use crate::LEGACY_SESSION;
+    use crate::VortexSessionExecute;
     use crate::arrays::MaskedArray;
     use crate::arrays::MaskedVTable;
     use crate::arrays::PrimitiveArray;
@@ -265,4 +252,33 @@ mod tests {
             decoded.display_values().to_string()
         );
     }
+
+    /// Regression test for issue #5989: execute_fast_path returns child with wrong dtype.
+    ///
+    /// When MaskedArray's validity mask is all true, returning the child's canonical form
+    /// directly would cause a dtype mismatch because the child has NonNullable dtype while
+    /// MaskedArray always has Nullable dtype.
+    #[test]
+    fn test_execute_with_all_valid_preserves_nullable_dtype() -> Result<(), VortexError> {
+        // Create a MaskedArray with AllValid validity.
+
+        // Child has NonNullable dtype, but MaskedArray's dtype is Nullable.
+        let child = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
+        assert_eq!(child.dtype().nullability(), Nullability::NonNullable);
+
+        let array = MaskedArray::try_new(child, Validity::AllValid)?;
+        assert_eq!(array.dtype().nullability(), Nullability::Nullable);
+
+        // Execute the array. This should produce a Canonical with Nullable dtype.
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let result: Canonical = array.into_array().execute(&mut ctx)?;
+
+        assert_eq!(
+            result.as_ref().dtype().nullability(),
+            Nullability::Nullable,
+            "MaskedArray execute should produce Nullable dtype"
+        );
+
+        Ok(())
+    }
 }