Fill out a few small pieces in Variant

Author: AdamGS

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

@@ -5,7 +5,6 @@
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 
 use crate::Canonical;
 use crate::ExecutionCtx;
@@ -27,6 +26,7 @@ use crate::arrays::Struct;
 use crate::arrays::StructArray;
 use crate::arrays::VarBinView;
 use crate::arrays::VarBinViewArray;
+use crate::arrays::VariantArray;
 use crate::arrays::dict::TakeExecute;
 use crate::arrays::dict::TakeReduce;
 
@@ -51,8 +51,12 @@ pub fn take_canonical(
         }
         Canonical::Struct(a) => Canonical::Struct(take_struct(&a, codes)),
         Canonical::Extension(a) => Canonical::Extension(take_extension(&a, codes, ctx)),
-        Canonical::Variant(_) => {
-            vortex_bail!("Variant arrays don't support Take")
+        Canonical::Variant(a) => {
+            let taken_child = a
+                .child()
+                .take(codes.clone().into_array())
+                .vortex_expect("VariantArray child could not be taken");
+            Canonical::Variant(VariantArray::new(taken_child))
         }
     })
 }

vortex-array/src/arrays/filter/execute/mod.rs

@@ -9,7 +9,6 @@ use std::sync::Arc;
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_mask::MaskValues;
 
@@ -21,6 +20,7 @@ use crate::arrays::ConstantArray;
 use crate::arrays::ExtensionArray;
 use crate::arrays::FilterArray;
 use crate::arrays::NullArray;
+use crate::arrays::VariantArray;
 use crate::scalar::Scalar;
 use crate::validity::Validity;
 
@@ -95,8 +95,12 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Ca
                 .vortex_expect("ExtensionArray storage type somehow could not be filtered");
             Canonical::Extension(ExtensionArray::new(a.ext_dtype().clone(), filtered_storage))
         }
-        Canonical::Variant(_) => {
-            vortex_panic!("Variant arrays don't support filtering")
+        Canonical::Variant(a) => {
+            let filtered_child = a
+                .child()
+                .filter(values_to_mask(mask))
+                .vortex_expect("VariantArray child could not be filtered");
+            Canonical::Variant(VariantArray::new(filtered_child))
         }
     }
 }

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

@@ -6,20 +6,22 @@
 use std::ops::BitAnd;
 
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 use vortex_mask::Mask;
 
 use crate::Canonical;
 use crate::IntoArray;
+use crate::ArrayVisitor;
 use crate::arrays::BoolArray;
 use crate::arrays::DecimalArray;
 use crate::arrays::ExtensionArray;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListViewArray;
+use crate::arrays::MaskedArray;
 use crate::arrays::NullArray;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::StructArray;
 use crate::arrays::VarBinViewArray;
+use crate::arrays::VariantArray;
 use crate::dtype::Nullability;
 use crate::executor::ExecutionCtx;
 use crate::match_each_decimal_value_type;
@@ -54,8 +56,8 @@ pub fn mask_validity_canonical(
         Canonical::Extension(a) => {
             Canonical::Extension(mask_validity_extension(a, validity_mask, ctx)?)
         }
-        Canonical::Variant(_) => {
-            vortex_bail!("Variant arrays don't masking validity")
+        Canonical::Variant(a) => {
+            Canonical::Variant(mask_validity_variant(a, validity_mask, ctx)?)
         }
     })
 }
@@ -200,3 +202,35 @@ fn mask_validity_extension(
         masked_storage,
     ))
 }
+
+fn mask_validity_variant(
+    array: VariantArray,
+    mask: &Mask,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<VariantArray> {
+    let child = array.child().clone();
+    let len = child.len();
+    let child_validity = child.validity()?;
+
+    match child_validity {
+        Validity::NonNullable | Validity::AllValid => {
+            // Child has no nulls — wrap in MaskedArray to apply the mask.
+            let new_validity = Validity::from_mask(mask.clone(), Nullability::Nullable);
+            let masked_child = MaskedArray::try_new(child, new_validity)?;
+            Ok(VariantArray::new(masked_child.into_array()))
+        }
+        Validity::AllInvalid => {
+            // Already all-null, ANDing with any mask is still all-null.
+            Ok(array)
+        }
+        Validity::Array(_) => {
+            // Child has an array-backed validity stored as its first child.
+            // Combine with the mask and replace that child via with_children.
+            let combined = combine_validity(&child_validity, mask, len, ctx)?;
+            let mut children = child.children();
+            children[0] = combined.to_array(len);
+            let new_child = child.with_children(children)?;
+            Ok(VariantArray::new(new_child))
+        }
+    }
+}