Dont use slot index

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-array/src/arrays/masked/execute.rs

@@ -25,6 +25,7 @@ use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::arrays::struct_::StructArrayExt;
 use crate::arrays::variant::VariantArrayExt;
+use crate::builtins::ArrayBuiltins;
 use crate::executor::ExecutionCtx;
 use crate::validity::Validity;
 
@@ -178,10 +179,9 @@ fn mask_validity_variant(
             core_storage
         }
         Validity::Array(_) => {
-            // Core storage has an array-backed validity stored as its first child.
-            // Combine with the mask and replace that child via with_children.
-            let combined = Validity::and(core_validity, validity)?;
-            core_storage.with_slot(0, combined.to_array(len))?
+            // Core storage already has nulls, but its physical validity layout depends on the
+            // actual encoding. Use the mask operation instead of rewriting a presumed slot.
+            core_storage.mask(validity.to_array(len))?
         }
     };
     let masked_shredded = if let Some(shredded) = array.shredded() {

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

@@ -142,27 +142,13 @@ mod tests {
         expected_core: &[Option<i32>],
         expected_shredded: &[Option<i32>],
     ) -> VortexResult<()> {
-        assert_eq!(array.len(), expected_core.len());
+        assert_variant_core_rows(array, expected_core)?;
         assert_eq!(array.len(), expected_shredded.len());
 
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
-        for (idx, expected) in expected_core.iter().enumerate() {
-            let scalar = array.core_storage().execute_scalar(idx, &mut ctx)?;
-            let variant = scalar.as_variant();
-            match expected {
-                Some(expected) => {
-                    let value = variant
-                        .value()
-                        .ok_or_else(|| vortex_err!("expected non-null variant row"))?;
-                    assert_eq!(value.as_primitive().typed_value::<i32>(), Some(*expected));
-                }
-                None => assert!(variant.is_null()),
-            }
-        }
-
         let shredded = array
             .shredded()
             .ok_or_else(|| vortex_err!("expected shredded child"))?;
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
         let shredded = shredded.clone().execute::<PrimitiveArray>(&mut ctx)?;
         let expected_shredded_array = if let Some(values) = expected_shredded
             .iter()
@@ -178,6 +164,30 @@ mod tests {
         Ok(())
     }
 
+    fn assert_variant_core_rows(
+        array: &VariantArray,
+        expected_core: &[Option<i32>],
+    ) -> VortexResult<()> {
+        assert_eq!(array.len(), expected_core.len());
+
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        for (idx, expected) in expected_core.iter().enumerate() {
+            let scalar = array.core_storage().execute_scalar(idx, &mut ctx)?;
+            let variant = scalar.as_variant();
+            match expected {
+                Some(expected) => {
+                    let value = variant
+                        .value()
+                        .ok_or_else(|| vortex_err!("expected non-null variant row"))?;
+                    assert_eq!(value.as_primitive().typed_value::<i32>(), Some(*expected));
+                }
+                None => assert!(variant.is_null()),
+            }
+        }
+
+        Ok(())
+    }
+
     #[test]
     fn try_new_exposes_core_storage_without_shredded() -> VortexResult<()> {
         let core_storage = core_storage(2);
@@ -331,6 +341,31 @@ mod tests {
         )
     }
 
+    #[test]
+    fn mask_preserves_chunked_core_storage_validity() -> VortexResult<()> {
+        let dtype = DType::Variant(Nullability::Nullable);
+        let core_chunks = [Some(1i32), None, Some(3), Some(4)]
+            .into_iter()
+            .map(|value| {
+                let scalar = match value {
+                    Some(value) => {
+                        Scalar::variant(Scalar::primitive(value, Nullability::NonNullable))
+                            .cast(&dtype)?
+                    }
+                    None => Scalar::null(dtype.clone()),
+                };
+                Ok(ConstantArray::new(scalar, 1).into_array())
+            })
+            .collect::<VortexResult<Vec<_>>>()?;
+        let core_storage = ChunkedArray::try_new(core_chunks, dtype)?.into_array();
+        let variant = VariantArray::try_new(core_storage, None)?;
+        let mask = BoolArray::from_iter([true, true, false, true]).into_array();
+
+        let masked = execute_variant(variant.into_array().mask(mask)?)?;
+
+        assert_variant_core_rows(&masked, &[Some(1), None, None, Some(4)])
+    }
+
     #[test]
     fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> {
         let core_storage = row_storage([1, 2, 3])?;