Skip to content
Merged
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
17 changes: 9 additions & 8 deletions vortex-array/src/arrays/bool/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
use std::fmt::Display;
use std::fmt::Formatter;

use arrow_array::BooleanArray;
use smallvec::smallvec;
use vortex_buffer::BitBuffer;
use vortex_buffer::BitBufferMeta;
Expand Down Expand Up @@ -338,14 +337,16 @@ impl FromIterator<bool> for BoolArray {

impl FromIterator<Option<bool>> for BoolArray {
fn from_iter<I: IntoIterator<Item = Option<bool>>>(iter: I) -> Self {
let (buffer, nulls) = BooleanArray::from_iter(iter).into_parts();
let iter = iter.into_iter();
let capacity = iter.size_hint().0;
let mut bits = BitBufferMut::with_capacity(capacity);
let mut validity = BitBufferMut::with_capacity(capacity);
for value in iter {
bits.append(value.unwrap_or_default());
validity.append(value.is_some());
}

BoolArray::new(
BitBuffer::from(buffer),
nulls
.map(|n| Validity::from(BitBuffer::from(n.into_inner())))
.unwrap_or(Validity::AllValid),
)
BoolArray::new(bits.freeze(), Validity::from(validity.freeze()))
}
}

Expand Down
10 changes: 2 additions & 8 deletions vortex-array/src/arrays/filter/execute/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -83,19 +83,13 @@ pub(super) fn execute_filter_fast_paths(
}

/// Filter a canonical array by a mask, returning a new canonical array.
pub(super) fn execute_filter(
canonical: Canonical,
mask: &Arc<MaskValues>,
ctx: &mut ExecutionCtx,
) -> Canonical {
pub(super) fn execute_filter(canonical: Canonical, mask: &Arc<MaskValues>) -> Canonical {
match canonical {
Canonical::Null(_) => Canonical::Null(NullArray::new(mask.true_count())),
Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)),
Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)),
Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)),
Canonical::VarBinView(a) => {
Canonical::VarBinView(varbinview::filter_varbinview(&a, mask, ctx))
}
Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)),
Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)),
Canonical::FixedSizeList(a) => {
Canonical::FixedSizeList(fixed_size_list::filter_fixed_size_list(&a, mask))
Expand Down
68 changes: 23 additions & 45 deletions vortex-array/src/arrays/filter/execute/varbinview.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,54 +3,32 @@

use std::sync::Arc;

use arrow_array::BooleanArray;
use vortex_error::VortexExpect;
use vortex_mask::Mask;
use vortex_buffer::Buffer;
use vortex_mask::MaskValues;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::arrays::VarBinView;
use crate::arrays::VarBinViewArray;
use crate::arrow::ArrowSessionExt;
use crate::arrow::FromArrowArray;

pub fn filter_varbinview(
array: &VarBinViewArray,
mask: &Arc<MaskValues>,
ctx: &mut ExecutionCtx,
) -> VarBinViewArray {
// Delegate to the Arrow implementation of filter over `VarBinView`.
arrow_filter_fn(
&array.clone().into_array(),
&Mask::Values(Arc::clone(mask)),
ctx,
)
.vortex_expect("VarBinViewArray is Arrow-compatible and supports arrow_filter_fn")
.as_::<VarBinView>()
.into_owned()
}

fn arrow_filter_fn(
array: &ArrayRef,
mask: &Mask,
ctx: &mut ExecutionCtx,
) -> vortex_error::VortexResult<ArrayRef> {
let values = match &mask {
Mask::Values(values) => values,
Mask::AllTrue(_) | Mask::AllFalse(_) => unreachable!("check in filter invoke"),
};

let array_ref = ctx
.session()
.arrow()
.clone()
.execute_arrow(array.clone(), None, ctx)?;
let mask_array = BooleanArray::new(values.bit_buffer().clone().into(), None);
let filtered = arrow_select::filter::filter(array_ref.as_ref(), &mask_array)?;

ArrayRef::from_arrow(filtered.as_ref(), array.dtype().is_nullable())
use crate::arrays::filter::execute::buffer;
use crate::arrays::filter::execute::filter_validity;
use crate::arrays::varbinview::BinaryView;
use crate::arrays::varbinview::VarBinViewArrayExt;
use crate::buffer::BufferHandle;

pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc<MaskValues>) -> VarBinViewArray {
let filtered_validity = filter_validity(array.varbinview_validity(), mask);

let views = Buffer::<BinaryView>::from_byte_buffer(array.views_handle().as_host().clone());
let filtered_views = buffer::filter_buffer(views, mask.as_ref());

// SAFETY: the filtered views are a subset of the original views and reference the same data
// buffers, and the validity is filtered by the same mask so lengths stay aligned.
unsafe {
VarBinViewArray::new_handle_unchecked(
BufferHandle::new_host(filtered_views.into_byte_buffer()),
Arc::clone(array.data_buffers()),
array.dtype().clone(),
filtered_validity,
)
}
}

#[cfg(test)]
Expand Down
2 changes: 1 addition & 1 deletion vortex-array/src/arrays/filter/vtable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ impl VTable for Filter {
// TODO(joe): fix the ownership of AnyCanonical
let child = Canonical::from(array.child().as_::<AnyCanonical>());
Ok(ExecutionResult::done(
execute_filter(child, &mask_values, ctx).into_array(),
execute_filter(child, &mask_values).into_array(),
))
}

Expand Down
Loading
Loading