Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions vortex-array/src/arrays/dict/execute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;

use crate::Canonical;
use crate::ExecutionCtx;
Expand All @@ -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;

Expand All @@ -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))
}
})
}
Expand Down
10 changes: 7 additions & 3 deletions vortex-array/src/arrays/filter/execute/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;

Expand Down Expand Up @@ -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))
}
}
}
40 changes: 36 additions & 4 deletions vortex-array/src/arrays/masked/execute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,22 @@
use std::ops::BitAnd;

use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_mask::Mask;

use crate::ArrayVisitor;
use crate::Canonical;
use crate::IntoArray;
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;
Expand Down Expand Up @@ -54,9 +56,7 @@ 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)?),
})
}

Expand Down Expand Up @@ -200,3 +200,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))
}
}
}
Loading