Skip to content

Commit

Permalink
Handle filtering empty struct arrays (#827)
Browse files Browse the repository at this point in the history
  • Loading branch information
robert3005 authored Sep 16, 2024
1 parent 4c80613 commit 4a02dba
Showing 1 changed file with 21 additions and 0 deletions.
21 changes: 21 additions & 0 deletions vortex-array/src/array/struct_/compute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use vortex_scalar::Scalar;
use crate::array::struct_::StructArray;
use crate::compute::unary::{scalar_at, scalar_at_unchecked, ScalarAtFn};
use crate::compute::{filter, slice, take, ArrayCompute, FilterFn, SliceFn, TakeFn};
use crate::stats::ArrayStatistics;
use crate::variants::StructArrayTrait;
use crate::{Array, ArrayDType, IntoArray};

Expand Down Expand Up @@ -85,6 +86,7 @@ impl FilterFn for StructArray {
let length = fields
.first()
.map(|a| a.len())
.or_else(|| predicate.statistics().compute_true_count())
.ok_or_else(|| vortex_err!("Struct arrays should have at least one field"))?;

Self::try_new(
Expand All @@ -96,3 +98,22 @@ impl FilterFn for StructArray {
.map(|a| a.into_array())
}
}

#[cfg(test)]
mod tests {
use crate::array::{BoolArray, StructArray};
use crate::compute::filter;
use crate::validity::Validity;
use crate::IntoArray;

#[test]
fn filter_empty_struct() {
let struct_arr =
StructArray::try_new(vec![].into(), vec![], 10, Validity::NonNullable).unwrap();
let mask = vec![
false, true, false, true, false, true, false, true, false, true,
];
let filtered = filter(struct_arr.as_ref(), &BoolArray::from(mask).into_array()).unwrap();
assert_eq!(filtered.len(), 5);
}
}

0 comments on commit 4a02dba

Please sign in to comment.