From 82679754004baa36217f68cd2b410be3f62e9d8c Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 30 Mar 2026 15:05:12 +0100 Subject: [PATCH 1/2] Fill out a few small pieces in Variant Signed-off-by: Adam Gutglick --- vortex-array/src/arrays/dict/execute.rs | 10 +++-- vortex-array/src/arrays/filter/execute/mod.rs | 10 +++-- vortex-array/src/arrays/masked/execute.rs | 40 +++++++++++++++++-- 3 files changed, 51 insertions(+), 9 deletions(-) diff --git a/vortex-array/src/arrays/dict/execute.rs b/vortex-array/src/arrays/dict/execute.rs index 9fd0a773327..d7ccced91b0 100644 --- a/vortex-array/src/arrays/dict/execute.rs +++ b/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)) } }) } diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 0dbffad14b5..e44117da110 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/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) -> 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)) } } } diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index a1b49fedef5..55cb260d454 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/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 { + 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)) + } + } +} From 483364ee0d88831797fc095e6a98dd2f2b161ed0 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 30 Mar 2026 15:10:03 +0100 Subject: [PATCH 2/2] lint Signed-off-by: Adam Gutglick --- vortex-array/src/arrays/masked/execute.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index 55cb260d454..f6025299f28 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -8,9 +8,9 @@ use std::ops::BitAnd; use vortex_error::VortexResult; use vortex_mask::Mask; +use crate::ArrayVisitor; use crate::Canonical; use crate::IntoArray; -use crate::ArrayVisitor; use crate::arrays::BoolArray; use crate::arrays::DecimalArray; use crate::arrays::ExtensionArray; @@ -56,9 +56,7 @@ pub fn mask_validity_canonical( Canonical::Extension(a) => { Canonical::Extension(mask_validity_extension(a, validity_mask, ctx)?) } - Canonical::Variant(a) => { - Canonical::Variant(mask_validity_variant(a, validity_mask, ctx)?) - } + Canonical::Variant(a) => Canonical::Variant(mask_validity_variant(a, validity_mask, ctx)?), }) }