1use std::ops::AddAssign;
21use std::sync::Arc;
22
23use arrow_array::builder::BooleanBufferBuilder;
24use arrow_array::cast::AsArray;
25use arrow_array::types::{
26 ArrowDictionaryKeyType, ArrowPrimitiveType, ByteArrayType, ByteViewType, RunEndIndexType,
27};
28use arrow_array::*;
29use arrow_buffer::{
30 ArrowNativeType, BooleanBuffer, NullBuffer, OffsetBuffer, RunEndBuffer, ScalarBuffer, bit_util,
31};
32use arrow_buffer::{Buffer, MutableBuffer};
33use arrow_data::bit_iterator::{BitIndexIterator, BitSliceIterator};
34use arrow_data::transform::MutableArrayData;
35use arrow_schema::*;
36
37const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8;
44
45#[derive(Debug)]
57pub struct SlicesIterator<'a>(BitSliceIterator<'a>);
58
59impl<'a> SlicesIterator<'a> {
60 pub fn new(filter: &'a BooleanArray) -> Self {
62 filter.values().into()
63 }
64}
65
66impl<'a> From<&'a BooleanBuffer> for SlicesIterator<'a> {
67 fn from(filter: &'a BooleanBuffer) -> Self {
68 Self(filter.set_slices())
69 }
70}
71
72impl Iterator for SlicesIterator<'_> {
73 type Item = (usize, usize);
74
75 fn next(&mut self) -> Option<Self::Item> {
76 self.0.next()
77 }
78}
79
80pub(crate) struct IndexIterator<'a> {
85 remaining: usize,
86 iter: BitIndexIterator<'a>,
87}
88
89impl<'a> IndexIterator<'a> {
90 pub(crate) fn new(filter: &'a BooleanArray, remaining: usize) -> Self {
91 assert_eq!(filter.null_count(), 0);
92 let iter = filter.values().set_indices();
93 Self { remaining, iter }
94 }
95
96 pub fn collect(mut self) -> Vec<usize> {
100 let len = self.remaining;
101 let mut result = Vec::with_capacity(len);
102 let ptr: *mut usize = result.as_mut_ptr();
103 for i in 0..len {
104 let next = self.iter.next();
107 debug_assert!(next.is_some(), "IndexIterator exhausted early");
108 unsafe {
109 *ptr.add(i) = next.unwrap_unchecked();
110 }
111 }
112 unsafe {
114 result.set_len(len);
115 }
116 result
117 }
118}
119
120impl Iterator for IndexIterator<'_> {
121 type Item = usize;
122
123 fn next(&mut self) -> Option<Self::Item> {
124 if self.remaining != 0 {
125 let next = self.iter.next().expect("IndexIterator exhausted early");
128 self.remaining -= 1;
129 return Some(next);
131 }
132 None
133 }
134
135 fn size_hint(&self) -> (usize, Option<usize>) {
136 (self.remaining, Some(self.remaining))
137 }
138}
139
140pub fn prep_null_mask_filter(filter: &BooleanArray) -> BooleanArray {
168 let nulls = filter.nulls().unwrap();
169 let mask = filter.values() & nulls.inner();
170 BooleanArray::new(mask, None)
171}
172
173pub fn filter(values: &dyn Array, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
202 let mut filter_builder = FilterBuilder::new(predicate);
203
204 if FilterBuilder::is_optimize_beneficial(values.data_type()) {
205 filter_builder = filter_builder.optimize();
208 }
209
210 let predicate = filter_builder.build();
211
212 filter_array(values, &predicate)
213}
214
215pub fn filter_record_batch(
226 record_batch: &RecordBatch,
227 predicate: &BooleanArray,
228) -> Result<RecordBatch, ArrowError> {
229 let mut filter_builder = FilterBuilder::new(predicate);
230 let num_cols = record_batch.num_columns();
231 if num_cols > 1
232 || (num_cols > 0
233 && FilterBuilder::is_optimize_beneficial(
234 record_batch.schema_ref().field(0).data_type(),
235 ))
236 {
237 filter_builder = filter_builder.optimize();
240 }
241 let filter = filter_builder.build();
242
243 filter.filter_record_batch(record_batch)
244}
245
246#[derive(Debug)]
248pub struct FilterBuilder {
249 filter: BooleanArray,
250 count: usize,
251 strategy: IterationStrategy,
252}
253
254impl FilterBuilder {
255 pub fn new(filter: &BooleanArray) -> Self {
257 Self::new_with_count(filter, filter.true_count())
258 }
259
260 pub(crate) fn new_with_count(filter: &BooleanArray, count: usize) -> Self {
261 let filter = match filter.null_count() {
262 0 => filter.clone(),
263 _ => prep_null_mask_filter(filter),
264 };
265
266 let strategy = IterationStrategy::default_strategy(filter.len(), count);
267
268 Self {
269 filter,
270 count,
271 strategy,
272 }
273 }
274
275 pub fn optimize(mut self) -> Self {
286 match self.strategy {
287 IterationStrategy::SlicesIterator => {
288 let slices = SlicesIterator::new(&self.filter).collect();
289 self.strategy = IterationStrategy::Slices(slices)
290 }
291 IterationStrategy::IndexIterator => {
292 let indices = IndexIterator::new(&self.filter, self.count).collect();
293 self.strategy = IterationStrategy::Indices(indices)
294 }
295 _ => {}
296 }
297 self
298 }
299
300 pub fn is_optimize_beneficial(data_type: &DataType) -> bool {
305 match data_type {
306 DataType::Struct(fields) => {
307 fields.len() > 1
308 || fields.len() == 1
309 && FilterBuilder::is_optimize_beneficial(fields[0].data_type())
310 }
311 DataType::Union(fields, UnionMode::Sparse) => !fields.is_empty(),
312 _ => false,
313 }
314 }
315
316 pub fn build(self) -> FilterPredicate {
318 FilterPredicate {
319 filter: self.filter,
320 count: self.count,
321 strategy: self.strategy,
322 }
323 }
324}
325
326#[derive(Debug)]
328enum IterationStrategy {
329 SlicesIterator,
331 IndexIterator,
333 Indices(Vec<usize>),
335 Slices(Vec<(usize, usize)>),
337 All,
339 None,
341}
342
343impl IterationStrategy {
344 fn default_strategy(filter_length: usize, filter_count: usize) -> Self {
347 if filter_length == 0 || filter_count == 0 {
348 return IterationStrategy::None;
349 }
350
351 if filter_count == filter_length {
352 return IterationStrategy::All;
353 }
354
355 let selectivity_frac = filter_count as f64 / filter_length as f64;
360 if selectivity_frac > FILTER_SLICES_SELECTIVITY_THRESHOLD {
361 return IterationStrategy::SlicesIterator;
362 }
363 IterationStrategy::IndexIterator
364 }
365}
366
367pub(crate) enum FilterSelection<'a> {
369 None,
370 All { len: usize },
371 Slices(FilterSlices<'a>),
372 Indices(FilterIndices<'a>),
373}
374
375pub(crate) type FilterSlices<'a> =
376 FilterIterator<std::iter::Copied<std::slice::Iter<'a, (usize, usize)>>, SlicesIterator<'a>>;
377
378pub(crate) type FilterIndices<'a> =
379 FilterIterator<std::iter::Copied<std::slice::Iter<'a, usize>>, IndexIterator<'a>>;
380
381pub(crate) enum FilterIterator<M, I> {
387 Materialized(M),
388 Lazy(I),
389}
390
391impl<M, I> FilterIterator<M, I>
392where
393 M: Iterator,
394 I: Iterator<Item = M::Item>,
395{
396 pub(crate) fn for_each<F>(self, f: F)
397 where
398 F: FnMut(M::Item),
399 {
400 match self {
401 Self::Materialized(iter) => iter.for_each(f),
402 Self::Lazy(iter) => iter.for_each(f),
403 }
404 }
405
406 pub(crate) fn try_for_each<F, E>(self, mut f: F) -> Result<(), E>
407 where
408 F: FnMut(M::Item) -> Result<(), E>,
409 {
410 match self {
411 Self::Materialized(iter) => {
412 for item in iter {
413 f(item)?;
414 }
415 }
416 Self::Lazy(iter) => {
417 for item in iter {
418 f(item)?;
419 }
420 }
421 }
422
423 Ok(())
424 }
425}
426
427#[derive(Debug)]
429pub struct FilterPredicate {
430 filter: BooleanArray,
431 count: usize,
432 strategy: IterationStrategy,
433}
434
435impl FilterPredicate {
436 pub fn filter(&self, values: &dyn Array) -> Result<ArrayRef, ArrowError> {
438 filter_array(values, self)
439 }
440
441 pub fn filter_record_batch(
446 &self,
447 record_batch: &RecordBatch,
448 ) -> Result<RecordBatch, ArrowError> {
449 let filtered_arrays = record_batch
450 .columns()
451 .iter()
452 .map(|a| filter_array(a, self))
453 .collect::<Result<Vec<_>, _>>()?;
454
455 unsafe {
458 Ok(RecordBatch::new_unchecked(
459 record_batch.schema(),
460 filtered_arrays,
461 self.count,
462 ))
463 }
464 }
465
466 pub fn count(&self) -> usize {
468 self.count
469 }
470
471 pub(crate) fn selection(&self) -> FilterSelection<'_> {
472 match &self.strategy {
473 IterationStrategy::None => FilterSelection::None,
474 IterationStrategy::All => FilterSelection::All { len: self.count },
475 IterationStrategy::Slices(slices) => {
476 FilterSelection::Slices(FilterIterator::Materialized(slices.iter().copied()))
477 }
478 IterationStrategy::SlicesIterator => {
479 FilterSelection::Slices(FilterIterator::Lazy(SlicesIterator::new(&self.filter)))
480 }
481 IterationStrategy::Indices(indices) => {
482 FilterSelection::Indices(FilterIterator::Materialized(indices.iter().copied()))
483 }
484 IterationStrategy::IndexIterator => FilterSelection::Indices(FilterIterator::Lazy(
485 IndexIterator::new(&self.filter, self.count),
486 )),
487 }
488 }
489
490 pub fn filter_nulls(&self, nulls: Option<&NullBuffer>) -> Option<NullBuffer> {
501 let nulls = nulls?;
502 if nulls.null_count() == 0 {
503 return None;
504 }
505
506 assert!(nulls.len() >= self.filter.len());
507 match self.strategy {
508 IterationStrategy::None => return None,
509 IterationStrategy::All => {
510 let nulls = nulls.slice(0, self.count);
511 return (nulls.null_count() != 0).then_some(nulls);
512 }
513 _ => {}
514 }
515
516 let nulls = filter_bits(nulls.inner(), self);
517 let null_count = self.count - nulls.count_set_bits_offset(0, self.count);
520
521 if null_count == 0 {
522 return None;
523 }
524
525 let buffer = BooleanBuffer::new(nulls, 0, self.count);
526 debug_assert_eq!(null_count, buffer.len() - buffer.count_set_bits());
527 Some(unsafe { NullBuffer::new_unchecked(buffer, null_count) })
530 }
531}
532
533fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result<ArrayRef, ArrowError> {
534 if predicate.filter.len() > values.len() {
535 return Err(ArrowError::InvalidArgumentError(format!(
536 "Filter predicate of length {} is larger than target array of length {}",
537 predicate.filter.len(),
538 values.len()
539 )));
540 }
541
542 match predicate.strategy {
543 IterationStrategy::None => Ok(new_empty_array(values.data_type())),
544 IterationStrategy::All => Ok(values.slice(0, predicate.count)),
545 _ => downcast_primitive_array! {
547 values => Ok(Arc::new(filter_primitive(values, predicate))),
548 DataType::Boolean => {
549 let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
550 Ok(Arc::new(filter_boolean(values, predicate)))
551 }
552 DataType::Utf8 => {
553 Ok(Arc::new(filter_bytes(values.as_string::<i32>(), predicate)))
554 }
555 DataType::LargeUtf8 => {
556 Ok(Arc::new(filter_bytes(values.as_string::<i64>(), predicate)))
557 }
558 DataType::Utf8View => {
559 Ok(Arc::new(filter_byte_view(values.as_string_view(), predicate)))
560 }
561 DataType::Binary => {
562 Ok(Arc::new(filter_bytes(values.as_binary::<i32>(), predicate)))
563 }
564 DataType::LargeBinary => {
565 Ok(Arc::new(filter_bytes(values.as_binary::<i64>(), predicate)))
566 }
567 DataType::BinaryView => {
568 Ok(Arc::new(filter_byte_view(values.as_binary_view(), predicate)))
569 }
570 DataType::FixedSizeBinary(_) => {
571 Ok(Arc::new(filter_fixed_size_binary(values.as_fixed_size_binary(), predicate)))
572 }
573 DataType::ListView(_) => {
574 Ok(Arc::new(filter_list_view::<i32>(values.as_list_view(), predicate)))
575 }
576 DataType::LargeListView(_) => {
577 Ok(Arc::new(filter_list_view::<i64>(values.as_list_view(), predicate)))
578 }
579 DataType::RunEndEncoded(_, _) => {
580 downcast_run_array!{
581 values => Ok(Arc::new(filter_run_end_array(values, predicate)?)),
582 t => unimplemented!("Filter not supported for RunEndEncoded type {:?}", t)
583 }
584 }
585 DataType::Dictionary(_, _) => downcast_dictionary_array! {
586 values => Ok(Arc::new(filter_dict(values, predicate))),
587 t => unimplemented!("Filter not supported for dictionary type {:?}", t)
588 }
589 DataType::Struct(_) => {
590 Ok(Arc::new(filter_struct(values.as_struct(), predicate)?))
591 }
592 DataType::Union(_, UnionMode::Sparse) => {
593 Ok(Arc::new(filter_sparse_union(values.as_union(), predicate)?))
594 }
595 _ => {
596 let data = values.to_data();
597 let mut mutable = MutableArrayData::new(
599 vec![&data],
600 false,
601 predicate.count,
602 );
603
604 match &predicate.strategy {
605 IterationStrategy::Slices(slices) => {
606 for (start, end) in slices {
607 mutable.try_extend(0, *start, *end)?;
608 }
609 }
610 _ => {
611 let iter = SlicesIterator::new(&predicate.filter);
612 for (start, end) in iter {
613 mutable.try_extend(0, start, end)?;
614 }
615 }
616 }
617
618 let data = mutable.freeze();
619 Ok(make_array(data))
620 }
621 },
622 }
623}
624
625fn filter_run_end_array<R: RunEndIndexType>(
627 array: &RunArray<R>,
628 predicate: &FilterPredicate,
629) -> Result<RunArray<R>, ArrowError>
630where
631 R::Native: Into<i64> + From<bool>,
632 R::Native: AddAssign,
633{
634 let run_ends: &RunEndBuffer<R::Native> = array.run_ends();
635 let start_physical = run_ends.get_start_physical_index();
636 let end_physical = run_ends.get_end_physical_index();
637 let physical_len = end_physical - start_physical + 1;
638
639 let mut new_run_ends = vec![R::default_value(); physical_len];
640 let offset = run_ends.offset() as u64;
641
642 let mut start = 0u64;
643 let mut j = 0;
644 let mut count = R::default_value();
645 let filter_values = predicate.filter.values();
646 let run_ends = run_ends.inner();
647
648 let pred: BooleanArray = BooleanBuffer::collect_bool(physical_len, |i| {
649 let mut keep = false;
650 let mut end = (run_ends[i + start_physical].into() as u64).saturating_sub(offset);
651 let difference = end.saturating_sub(filter_values.len() as u64);
652 end -= difference;
653
654 for pred in (start..end).map(|i| unsafe { filter_values.value_unchecked(i as usize) }) {
656 count += R::Native::from(pred);
657 keep |= pred
658 }
659 new_run_ends[j] = count;
661 j += keep as usize;
662
663 start = end;
664 keep
665 })
666 .into();
667
668 new_run_ends.truncate(j);
669
670 let values = array.values_slice();
671 let values = filter(values.as_ref(), &pred)?;
672
673 let run_ends = PrimitiveArray::<R>::try_new(new_run_ends.into(), None)?;
674 RunArray::try_new(&run_ends, &values)
675}
676
677fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
679 let src = buffer.values();
680 let offset = buffer.offset();
681 assert!(buffer.len() >= predicate.filter.len());
682
683 match &predicate.strategy {
684 IterationStrategy::IndexIterator => {
685 let bits =
686 IndexIterator::new(&predicate.filter, predicate.count).map(|src_idx| unsafe {
688 bit_util::get_bit_raw(buffer.values().as_ptr(), src_idx + offset)
689 });
690
691 unsafe { MutableBuffer::from_trusted_len_iter_bool(bits).into() }
693 }
694 IterationStrategy::Indices(indices) => {
695 let bits = indices.iter().map(|src_idx| unsafe {
697 bit_util::get_bit_raw(buffer.values().as_ptr(), *src_idx + offset)
698 });
699 unsafe { MutableBuffer::from_trusted_len_iter_bool(bits).into() }
701 }
702 IterationStrategy::SlicesIterator => {
703 let mut builder = BooleanBufferBuilder::new(predicate.count);
704 for (start, end) in SlicesIterator::new(&predicate.filter) {
705 builder.append_packed_range(start + offset..end + offset, src)
706 }
707 builder.into()
708 }
709 IterationStrategy::Slices(slices) => {
710 let mut builder = BooleanBufferBuilder::new(predicate.count);
711 for (start, end) in slices {
712 builder.append_packed_range(*start + offset..*end + offset, src)
713 }
714 builder.into()
715 }
716 IterationStrategy::All | IterationStrategy::None => unreachable!(),
717 }
718}
719
720fn filter_boolean(array: &BooleanArray, predicate: &FilterPredicate) -> BooleanArray {
722 let buffer = filter_bits(array.values(), predicate);
723 let values = BooleanBuffer::new(buffer, 0, predicate.count);
724 let nulls = predicate.filter_nulls(array.nulls());
725
726 BooleanArray::new(values, nulls)
727}
728
729#[inline(never)]
730pub(crate) fn filter_native<T: ArrowNativeType>(
731 values: &[T],
732 predicate: &FilterPredicate,
733) -> Buffer {
734 assert!(values.len() >= predicate.filter.len());
735
736 match &predicate.strategy {
737 IterationStrategy::SlicesIterator => {
738 let mut buffer = Vec::with_capacity(predicate.count);
739 for (start, end) in SlicesIterator::new(&predicate.filter) {
740 buffer.extend_from_slice(unsafe { values.get_unchecked(start..end) });
742 }
743 buffer.into()
744 }
745 IterationStrategy::Slices(slices) => {
746 let mut buffer = Vec::with_capacity(predicate.count);
747 for (start, end) in slices {
748 buffer.extend_from_slice(unsafe { values.get_unchecked(*start..*end) });
750 }
751 buffer.into()
752 }
753 IterationStrategy::IndexIterator => {
754 let iter = IndexIterator::new(&predicate.filter, predicate.count)
756 .map(|x| unsafe { *values.get_unchecked(x) });
757
758 unsafe { MutableBuffer::from_trusted_len_iter(iter) }.into()
760 }
761 IterationStrategy::Indices(indices) => {
762 let iter = indices.iter().map(|x| unsafe { *values.get_unchecked(*x) });
764 iter.collect::<Vec<_>>().into()
765 }
766 IterationStrategy::All | IterationStrategy::None => unreachable!(),
767 }
768}
769
770fn filter_primitive<T>(array: &PrimitiveArray<T>, predicate: &FilterPredicate) -> PrimitiveArray<T>
772where
773 T: ArrowPrimitiveType,
774{
775 let buffer = filter_native(array.values(), predicate);
776 let values = ScalarBuffer::new(buffer, 0, predicate.count);
777 let nulls = predicate.filter_nulls(array.nulls());
778 let filtered = PrimitiveArray::new(values, nulls);
779
780 if array.data_type() == &T::DATA_TYPE {
782 filtered
783 } else {
784 filtered.with_data_type(array.data_type().clone())
785 }
786}
787
788struct FilterBytes<'a, OffsetSize> {
793 src_offsets: &'a [OffsetSize],
794 src_values: &'a [u8],
795 dst_offsets: Vec<OffsetSize>,
796 dst_values: Vec<u8>,
797 cur_offset: OffsetSize,
798}
799
800impl<'a, OffsetSize> FilterBytes<'a, OffsetSize>
801where
802 OffsetSize: OffsetSizeTrait,
803{
804 fn new<T>(capacity: usize, array: &'a GenericByteArray<T>) -> Self
805 where
806 T: ByteArrayType<Offset = OffsetSize>,
807 {
808 let dst_values = Vec::new();
809 let mut dst_offsets: Vec<OffsetSize> = Vec::with_capacity(capacity + 1);
810 let cur_offset = OffsetSize::from_usize(0).unwrap();
811
812 dst_offsets.push(cur_offset);
813
814 Self {
815 src_offsets: array.value_offsets(),
816 src_values: array.value_data(),
817 dst_offsets,
818 dst_values,
819 cur_offset,
820 }
821 }
822
823 #[inline]
825 fn get_value_offset(&self, idx: usize) -> usize {
826 self.src_offsets[idx].as_usize()
827 }
828
829 #[inline]
831 fn get_value_range(&self, idx: usize) -> (usize, usize, OffsetSize) {
832 let start = self.get_value_offset(idx);
834 let end = self.get_value_offset(idx + 1);
835 let len = OffsetSize::from_usize(end - start).expect("illegal offset range");
836 (start, end, len)
837 }
838
839 fn extend_offsets_idx(&mut self, iter: impl Iterator<Item = usize>) {
840 self.dst_offsets.extend(iter.map(|idx| {
841 let start = self.src_offsets[idx].as_usize();
842 let end = self.src_offsets[idx + 1].as_usize();
843 let len = OffsetSize::from_usize(end - start).expect("illegal offset range");
844 self.cur_offset += len;
845
846 self.cur_offset
847 }));
848 }
849
850 fn extend_idx(&mut self, iter: impl Iterator<Item = usize>) {
852 self.dst_values.reserve_exact(self.cur_offset.as_usize());
853
854 for idx in iter {
855 let start = self.src_offsets[idx].as_usize();
856 let end = self.src_offsets[idx + 1].as_usize();
857 self.dst_values
858 .extend_from_slice(&self.src_values[start..end]);
859 }
860 }
861
862 fn extend_offsets_slices(&mut self, iter: impl Iterator<Item = (usize, usize)>, count: usize) {
863 self.dst_offsets.reserve_exact(count);
864 for (start, end) in iter {
865 for idx in start..end {
867 let (_, _, len) = self.get_value_range(idx);
868 self.cur_offset += len;
869 self.dst_offsets.push(self.cur_offset);
870 }
871 }
872 }
873
874 fn extend_slices(&mut self, iter: impl Iterator<Item = (usize, usize)>) {
876 self.dst_values.reserve_exact(self.cur_offset.as_usize());
877
878 for (start, end) in iter {
879 let value_start = self.get_value_offset(start);
880 let value_end = self.get_value_offset(end);
881 self.dst_values
882 .extend_from_slice(&self.src_values[value_start..value_end]);
883 }
884 }
885}
886
887fn filter_bytes<T>(array: &GenericByteArray<T>, predicate: &FilterPredicate) -> GenericByteArray<T>
892where
893 T: ByteArrayType,
894{
895 let mut filter = FilterBytes::new(predicate.count, array);
896
897 match &predicate.strategy {
898 IterationStrategy::SlicesIterator => {
899 filter.extend_offsets_slices(SlicesIterator::new(&predicate.filter), predicate.count);
900 filter.extend_slices(SlicesIterator::new(&predicate.filter))
901 }
902 IterationStrategy::Slices(slices) => {
903 filter.extend_offsets_slices(slices.iter().cloned(), predicate.count);
904 filter.extend_slices(slices.iter().cloned())
905 }
906 IterationStrategy::IndexIterator => {
907 filter.extend_offsets_idx(IndexIterator::new(&predicate.filter, predicate.count));
908 filter.extend_idx(IndexIterator::new(&predicate.filter, predicate.count))
909 }
910 IterationStrategy::Indices(indices) => {
911 filter.extend_offsets_idx(indices.iter().cloned());
912 filter.extend_idx(indices.iter().cloned())
913 }
914 IterationStrategy::All | IterationStrategy::None => unreachable!(),
915 }
916
917 let offsets = unsafe { OffsetBuffer::new_unchecked(filter.dst_offsets.into()) };
920 let nulls = predicate.filter_nulls(array.nulls());
921
922 unsafe { GenericByteArray::new_unchecked(offsets, filter.dst_values.into(), nulls) }
926}
927
928fn filter_byte_view<T: ByteViewType>(
930 array: &GenericByteViewArray<T>,
931 predicate: &FilterPredicate,
932) -> GenericByteViewArray<T> {
933 let new_view_buffer = filter_native(array.views(), predicate);
934 let views = ScalarBuffer::new(new_view_buffer, 0, predicate.count);
935 let buffers = array.data_buffers().to_vec();
936 let nulls = predicate.filter_nulls(array.nulls());
937
938 unsafe { GenericByteViewArray::new_unchecked(views, buffers, nulls) }
942}
943
944fn filter_fixed_size_binary(
945 array: &FixedSizeBinaryArray,
946 predicate: &FilterPredicate,
947) -> FixedSizeBinaryArray {
948 let values: &[u8] = array.values();
949 let value_length = array.value_length() as usize;
950 let calculate_offset_from_index = |index: usize| index * value_length;
951 let buffer = match &predicate.strategy {
952 IterationStrategy::SlicesIterator => {
953 let mut buffer = MutableBuffer::with_capacity(predicate.count * value_length);
954 for (start, end) in SlicesIterator::new(&predicate.filter) {
955 buffer.extend_from_slice(
956 &values[calculate_offset_from_index(start)..calculate_offset_from_index(end)],
957 );
958 }
959 buffer
960 }
961 IterationStrategy::Slices(slices) => {
962 let mut buffer = MutableBuffer::with_capacity(predicate.count * value_length);
963 for (start, end) in slices {
964 buffer.extend_from_slice(
965 &values[calculate_offset_from_index(*start)..calculate_offset_from_index(*end)],
966 );
967 }
968 buffer
969 }
970 IterationStrategy::IndexIterator => {
971 let iter = IndexIterator::new(&predicate.filter, predicate.count).map(|x| {
972 &values[calculate_offset_from_index(x)..calculate_offset_from_index(x + 1)]
973 });
974
975 let mut buffer = MutableBuffer::new(predicate.count * value_length);
976 iter.for_each(|item| buffer.extend_from_slice(item));
977 buffer
978 }
979 IterationStrategy::Indices(indices) => {
980 let iter = indices.iter().map(|x| {
981 &values[calculate_offset_from_index(*x)..calculate_offset_from_index(*x + 1)]
982 });
983
984 let mut buffer = MutableBuffer::new(predicate.count * value_length);
985 iter.for_each(|item| buffer.extend_from_slice(item));
986 buffer
987 }
988 IterationStrategy::All | IterationStrategy::None => unreachable!(),
989 };
990
991 let nulls = predicate.filter_nulls(array.nulls());
992
993 FixedSizeBinaryArray::new(array.value_length(), buffer.into(), nulls)
994}
995
996fn filter_dict<K: ArrowDictionaryKeyType>(
998 array: &DictionaryArray<K>,
999 predicate: &FilterPredicate,
1000) -> DictionaryArray<K> {
1001 let new_keys = filter_primitive(array.keys(), predicate);
1004 unsafe { DictionaryArray::new_unchecked(new_keys, array.values().clone()) }
1005}
1006
1007fn filter_struct(
1009 array: &StructArray,
1010 predicate: &FilterPredicate,
1011) -> Result<StructArray, ArrowError> {
1012 let columns = array
1013 .columns()
1014 .iter()
1015 .map(|column| filter_array(column, predicate))
1016 .collect::<Result<_, _>>()?;
1017
1018 let nulls = predicate.filter_nulls(array.nulls());
1019
1020 Ok(unsafe {
1021 StructArray::new_unchecked_with_length(
1022 array.fields().clone(),
1023 columns,
1024 nulls,
1025 predicate.count(),
1026 )
1027 })
1028}
1029
1030fn filter_sparse_union(
1032 array: &UnionArray,
1033 predicate: &FilterPredicate,
1034) -> Result<UnionArray, ArrowError> {
1035 let DataType::Union(fields, UnionMode::Sparse) = array.data_type() else {
1036 unreachable!()
1037 };
1038
1039 let type_ids = filter_primitive(
1040 &Int8Array::try_new(array.type_ids().clone(), None)?,
1041 predicate,
1042 );
1043
1044 let children = fields
1045 .iter()
1046 .map(|(child_type_id, _)| filter_array(array.child(child_type_id), predicate))
1047 .collect::<Result<_, _>>()?;
1048
1049 Ok(unsafe {
1050 UnionArray::new_unchecked(fields.clone(), type_ids.into_parts().1, None, children)
1051 })
1052}
1053
1054fn filter_list_view<OffsetType: OffsetSizeTrait>(
1056 array: &GenericListViewArray<OffsetType>,
1057 predicate: &FilterPredicate,
1058) -> GenericListViewArray<OffsetType> {
1059 let filtered_offsets = filter_native::<OffsetType>(array.offsets(), predicate);
1060 let filtered_sizes = filter_native::<OffsetType>(array.sizes(), predicate);
1061
1062 let field = match array.data_type() {
1063 DataType::ListView(field) | DataType::LargeListView(field) => field.clone(),
1064 _ => unreachable!(),
1065 };
1066 let offsets = ScalarBuffer::new(filtered_offsets, 0, predicate.count);
1067 let sizes = ScalarBuffer::new(filtered_sizes, 0, predicate.count);
1068 let values = array.values().clone();
1069 let nulls = predicate.filter_nulls(array.nulls());
1070
1071 unsafe { GenericListViewArray::new_unchecked(field, offsets, sizes, values, nulls) }
1075}
1076
1077#[cfg(test)]
1078mod tests {
1079 use super::*;
1080 use arrow_array::builder::*;
1081 use arrow_array::cast::as_run_array;
1082 use arrow_array::types::*;
1083 use rand::distr::uniform::{UniformSampler, UniformUsize};
1084 use rand::distr::{Alphanumeric, StandardUniform};
1085 use rand::prelude::*;
1086 use rand::rng;
1087
1088 #[test]
1089 fn test_filter_nulls_all() {
1090 let predicate = FilterBuilder::new(&BooleanArray::from(vec![true, true])).build();
1091 let nulls = NullBuffer::from(vec![true, false]);
1092 assert_eq!(predicate.filter_nulls(Some(&nulls)), Some(nulls));
1093 }
1094
1095 #[test]
1096 fn test_filter_nulls_none() {
1097 let predicate = FilterBuilder::new(&BooleanArray::from(vec![false, false])).build();
1098 let nulls = NullBuffer::from(vec![true, false]);
1099 assert_eq!(predicate.filter_nulls(Some(&nulls)), None);
1100 }
1101
1102 #[test]
1103 fn test_filter_nulls_selection() {
1104 let filters = [
1107 vec![],
1108 vec![true],
1109 vec![true, true],
1110 vec![true; 6],
1111 vec![false; 6],
1112 vec![true, false, false, false, false, false],
1113 vec![true, true, true, true, true, false],
1114 ];
1115 for offset in [0, 3, 9] {
1116 let mut validity = vec![false; offset];
1117 validity.extend([true, false, true, false, true, false]);
1118 let nulls = NullBuffer::from(validity).slice(offset, 6);
1119 for filter in &filters {
1120 let expected: NullBuffer = filter
1121 .iter()
1122 .enumerate()
1123 .filter(|(_, selected)| **selected)
1124 .map(|(i, _)| nulls.is_valid(i))
1125 .collect();
1126 let expected = (expected.null_count() != 0).then_some(expected);
1127 let filter = BooleanArray::from(filter.clone());
1128 for optimize in [false, true] {
1129 let builder = FilterBuilder::new(&filter);
1130 let predicate = if optimize {
1131 builder.optimize()
1132 } else {
1133 builder
1134 }
1135 .build();
1136 assert_eq!(predicate.filter_nulls(Some(&nulls)), expected);
1137 assert_eq!(predicate.filter_nulls(None), None);
1138 assert_eq!(
1139 predicate.filter_nulls(Some(&NullBuffer::new_valid(6))),
1140 None
1141 );
1142 }
1143 }
1144 }
1145 }
1146
1147 #[test]
1148 fn test_filter_nulls_length_check() {
1149 for filter in [vec![true, true], vec![false, false], vec![true, false]] {
1150 let predicate = FilterBuilder::new(&BooleanArray::from(filter)).build();
1151 let nulls = NullBuffer::new_null(1);
1152 assert!(std::panic::catch_unwind(|| predicate.filter_nulls(Some(&nulls))).is_err());
1153 assert_eq!(
1155 predicate.filter_nulls(Some(&NullBuffer::new_valid(1))),
1156 None
1157 );
1158 assert_eq!(predicate.filter_nulls(None), None);
1159 }
1160 }
1161
1162 macro_rules! def_temporal_test {
1163 ($test:ident, $array_type: ident, $data: expr) => {
1164 #[test]
1165 fn $test() {
1166 let a = $data;
1167 let b = BooleanArray::from(vec![true, false, true, false]);
1168 let c = filter(&a, &b).unwrap();
1169 let d = c.as_ref().as_any().downcast_ref::<$array_type>().unwrap();
1170 assert_eq!(2, d.len());
1171 assert_eq!(1, d.value(0));
1172 assert_eq!(3, d.value(1));
1173 }
1174 };
1175 }
1176
1177 def_temporal_test!(
1178 test_filter_date32,
1179 Date32Array,
1180 Date32Array::from(vec![1, 2, 3, 4])
1181 );
1182 def_temporal_test!(
1183 test_filter_date64,
1184 Date64Array,
1185 Date64Array::from(vec![1, 2, 3, 4])
1186 );
1187 def_temporal_test!(
1188 test_filter_time32_second,
1189 Time32SecondArray,
1190 Time32SecondArray::from(vec![1, 2, 3, 4])
1191 );
1192 def_temporal_test!(
1193 test_filter_time32_millisecond,
1194 Time32MillisecondArray,
1195 Time32MillisecondArray::from(vec![1, 2, 3, 4])
1196 );
1197 def_temporal_test!(
1198 test_filter_time64_microsecond,
1199 Time64MicrosecondArray,
1200 Time64MicrosecondArray::from(vec![1, 2, 3, 4])
1201 );
1202 def_temporal_test!(
1203 test_filter_time64_nanosecond,
1204 Time64NanosecondArray,
1205 Time64NanosecondArray::from(vec![1, 2, 3, 4])
1206 );
1207 def_temporal_test!(
1208 test_filter_duration_second,
1209 DurationSecondArray,
1210 DurationSecondArray::from(vec![1, 2, 3, 4])
1211 );
1212 def_temporal_test!(
1213 test_filter_duration_millisecond,
1214 DurationMillisecondArray,
1215 DurationMillisecondArray::from(vec![1, 2, 3, 4])
1216 );
1217 def_temporal_test!(
1218 test_filter_duration_microsecond,
1219 DurationMicrosecondArray,
1220 DurationMicrosecondArray::from(vec![1, 2, 3, 4])
1221 );
1222 def_temporal_test!(
1223 test_filter_duration_nanosecond,
1224 DurationNanosecondArray,
1225 DurationNanosecondArray::from(vec![1, 2, 3, 4])
1226 );
1227 def_temporal_test!(
1228 test_filter_timestamp_second,
1229 TimestampSecondArray,
1230 TimestampSecondArray::from(vec![1, 2, 3, 4])
1231 );
1232 def_temporal_test!(
1233 test_filter_timestamp_millisecond,
1234 TimestampMillisecondArray,
1235 TimestampMillisecondArray::from(vec![1, 2, 3, 4])
1236 );
1237 def_temporal_test!(
1238 test_filter_timestamp_microsecond,
1239 TimestampMicrosecondArray,
1240 TimestampMicrosecondArray::from(vec![1, 2, 3, 4])
1241 );
1242 def_temporal_test!(
1243 test_filter_timestamp_nanosecond,
1244 TimestampNanosecondArray,
1245 TimestampNanosecondArray::from(vec![1, 2, 3, 4])
1246 );
1247
1248 #[test]
1249 fn test_filter_array_slice() {
1250 let a = Int32Array::from(vec![5, 6, 7, 8, 9]).slice(1, 4);
1251 let b = BooleanArray::from(vec![true, false, false, true]);
1252 let c = filter(&a, &b).unwrap();
1256 let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1257 assert_eq!(2, d.len());
1258 assert_eq!(6, d.value(0));
1259 assert_eq!(9, d.value(1));
1260 }
1261
1262 #[test]
1263 fn test_filter_array_low_density() {
1264 let mut data_values = (1..=65).collect::<Vec<i32>>();
1266 let mut filter_values = (1..=65).map(|i| matches!(i % 65, 0)).collect::<Vec<bool>>();
1267 data_values.extend_from_slice(&[66, 67]);
1269 filter_values.extend_from_slice(&[false, true]);
1270 let a = Int32Array::from(data_values);
1271 let b = BooleanArray::from(filter_values);
1272 let c = filter(&a, &b).unwrap();
1273 let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1274 assert_eq!(2, d.len());
1275 assert_eq!(65, d.value(0));
1276 assert_eq!(67, d.value(1));
1277 }
1278
1279 #[test]
1280 fn test_filter_array_high_density() {
1281 let mut data_values = (1..=65).map(Some).collect::<Vec<_>>();
1283 let mut filter_values = (1..=65)
1284 .map(|i| !matches!(i % 65, 0))
1285 .collect::<Vec<bool>>();
1286 data_values[1] = None;
1288 data_values.extend_from_slice(&[Some(66), None, Some(67), None]);
1290 filter_values.extend_from_slice(&[false, true, true, true]);
1291 let a = Int32Array::from(data_values);
1292 let b = BooleanArray::from(filter_values);
1293 let c = filter(&a, &b).unwrap();
1294 let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1295 assert_eq!(67, d.len());
1296 assert_eq!(3, d.null_count());
1297 assert_eq!(1, d.value(0));
1298 assert!(d.is_null(1));
1299 assert_eq!(64, d.value(63));
1300 assert!(d.is_null(64));
1301 assert_eq!(67, d.value(65));
1302 }
1303
1304 #[test]
1305 fn test_filter_string_array_simple() {
1306 let a = StringArray::from(vec!["hello", " ", "world", "!"]);
1307 let b = BooleanArray::from(vec![true, false, true, false]);
1308 let c = filter(&a, &b).unwrap();
1309 let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
1310 assert_eq!(2, d.len());
1311 assert_eq!("hello", d.value(0));
1312 assert_eq!("world", d.value(1));
1313 }
1314
1315 #[test]
1316 fn test_filter_primitive_array_with_null() {
1317 let a = Int32Array::from(vec![Some(5), None]);
1318 let b = BooleanArray::from(vec![false, true]);
1319 let c = filter(&a, &b).unwrap();
1320 let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1321 assert_eq!(1, d.len());
1322 assert!(d.is_null(0));
1323 }
1324
1325 #[test]
1326 fn test_filter_string_array_with_null() {
1327 let a = StringArray::from(vec![Some("hello"), None, Some("world"), None]);
1328 let b = BooleanArray::from(vec![true, false, false, true]);
1329 let c = filter(&a, &b).unwrap();
1330 let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
1331 assert_eq!(2, d.len());
1332 assert_eq!("hello", d.value(0));
1333 assert!(!d.is_null(0));
1334 assert!(d.is_null(1));
1335 }
1336
1337 #[test]
1338 fn test_filter_binary_array_with_null() {
1339 let data: Vec<Option<&[u8]>> = vec![Some(b"hello"), None, Some(b"world"), None];
1340 let a = BinaryArray::from(data);
1341 let b = BooleanArray::from(vec![true, false, false, true]);
1342 let c = filter(&a, &b).unwrap();
1343 let d = c.as_ref().as_any().downcast_ref::<BinaryArray>().unwrap();
1344 assert_eq!(2, d.len());
1345 assert_eq!(b"hello", d.value(0));
1346 assert!(!d.is_null(0));
1347 assert!(d.is_null(1));
1348 }
1349
1350 fn _test_filter_byte_view<T>()
1351 where
1352 T: ByteViewType,
1353 str: AsRef<T::Native>,
1354 T::Native: PartialEq,
1355 {
1356 let array = {
1357 let mut builder = GenericByteViewBuilder::<T>::new();
1359 builder.append_value("hello");
1360 builder.append_value("world");
1361 builder.append_null();
1362 builder.append_value("large payload over 12 bytes");
1363 builder.append_value("lulu");
1364 builder.finish()
1365 };
1366
1367 {
1368 let predicate = BooleanArray::from(vec![true, false, true, true, false]);
1369 let actual = filter(&array, &predicate).unwrap();
1370
1371 assert_eq!(actual.len(), 3);
1372
1373 let expected = {
1374 let mut builder = GenericByteViewBuilder::<T>::new();
1376 builder.append_value("hello");
1377 builder.append_null();
1378 builder.append_value("large payload over 12 bytes");
1379 builder.finish()
1380 };
1381
1382 assert_eq!(actual.as_ref(), &expected);
1383 }
1384
1385 {
1386 let predicate = BooleanArray::from(vec![true, false, false, false, true]);
1387 let actual = filter(&array, &predicate).unwrap();
1388
1389 assert_eq!(actual.len(), 2);
1390
1391 let expected = {
1392 let mut builder = GenericByteViewBuilder::<T>::new();
1394 builder.append_value("hello");
1395 builder.append_value("lulu");
1396 builder.finish()
1397 };
1398
1399 assert_eq!(actual.as_ref(), &expected);
1400 }
1401 }
1402
1403 #[test]
1404 fn test_filter_string_view() {
1405 _test_filter_byte_view::<StringViewType>()
1406 }
1407
1408 #[test]
1409 fn test_filter_binary_view() {
1410 _test_filter_byte_view::<BinaryViewType>()
1411 }
1412
1413 #[test]
1414 fn test_filter_fixed_binary() {
1415 let v1 = [1_u8, 2];
1416 let v2 = [3_u8, 4];
1417 let v3 = [5_u8, 6];
1418 let v = vec![&v1, &v2, &v3];
1419 let a = FixedSizeBinaryArray::try_from(v).unwrap();
1420 let b = BooleanArray::from(vec![true, false, true]);
1421 let c = filter(&a, &b).unwrap();
1422 let d = c
1423 .as_ref()
1424 .as_any()
1425 .downcast_ref::<FixedSizeBinaryArray>()
1426 .unwrap();
1427 assert_eq!(d.len(), 2);
1428 assert_eq!(d.value(0), &v1);
1429 assert_eq!(d.value(1), &v3);
1430 let c2 = FilterBuilder::new(&b)
1431 .optimize()
1432 .build()
1433 .filter(&a)
1434 .unwrap();
1435 let d2 = c2
1436 .as_ref()
1437 .as_any()
1438 .downcast_ref::<FixedSizeBinaryArray>()
1439 .unwrap();
1440 assert_eq!(d, d2);
1441
1442 let b = BooleanArray::from(vec![false, false, false]);
1443 let c = filter(&a, &b).unwrap();
1444 let d = c
1445 .as_ref()
1446 .as_any()
1447 .downcast_ref::<FixedSizeBinaryArray>()
1448 .unwrap();
1449 assert_eq!(d.len(), 0);
1450
1451 let b = BooleanArray::from(vec![true, true, true]);
1452 let c = filter(&a, &b).unwrap();
1453 let d = c
1454 .as_ref()
1455 .as_any()
1456 .downcast_ref::<FixedSizeBinaryArray>()
1457 .unwrap();
1458 assert_eq!(d.len(), 3);
1459 assert_eq!(d.value(0), &v1);
1460 assert_eq!(d.value(1), &v2);
1461 assert_eq!(d.value(2), &v3);
1462
1463 let b = BooleanArray::from(vec![false, false, true]);
1464 let c = filter(&a, &b).unwrap();
1465 let d = c
1466 .as_ref()
1467 .as_any()
1468 .downcast_ref::<FixedSizeBinaryArray>()
1469 .unwrap();
1470 assert_eq!(d.len(), 1);
1471 assert_eq!(d.value(0), &v3);
1472 let c2 = FilterBuilder::new(&b)
1473 .optimize()
1474 .build()
1475 .filter(&a)
1476 .unwrap();
1477 let d2 = c2
1478 .as_ref()
1479 .as_any()
1480 .downcast_ref::<FixedSizeBinaryArray>()
1481 .unwrap();
1482 assert_eq!(d, d2);
1483 }
1484
1485 #[test]
1486 fn test_filter_array_slice_with_null() {
1487 let a = Int32Array::from(vec![Some(5), None, Some(7), Some(8), Some(9)]).slice(1, 4);
1488 let b = BooleanArray::from(vec![true, false, false, true]);
1489 let c = filter(&a, &b).unwrap();
1493 let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1494 assert_eq!(2, d.len());
1495 assert!(d.is_null(0));
1496 assert!(!d.is_null(1));
1497 assert_eq!(9, d.value(1));
1498 }
1499
1500 #[test]
1501 fn test_filter_run_end_encoding_array() {
1502 let run_ends = Int64Array::from(vec![2, 3, 8]);
1503 let values = Int64Array::from(vec![7, -2, 9]);
1504 let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1505 let b = BooleanArray::from(vec![true, false, true, false, true, false, true, false]);
1506 let c = filter(&a, &b).unwrap();
1507 let actual: &RunArray<Int64Type> = as_run_array(&c);
1508 assert_eq!(4, actual.len());
1509
1510 let expected = RunArray::try_new(
1511 &Int64Array::from(vec![1, 2, 4]),
1512 &Int64Array::from(vec![7, -2, 9]),
1513 )
1514 .expect("Failed to make expected RunArray test is broken");
1515
1516 assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1517 assert_eq!(actual.values(), expected.values())
1518 }
1519
1520 #[test]
1521 fn test_filter_run_end_encoding_array_sliced() {
1522 let run_ends = Int64Array::from(vec![2, 3, 8]);
1523 let values = Int64Array::from(vec![7, -2, 9]);
1524 let a = RunArray::try_new(&run_ends, &values).unwrap(); let a = a.slice(2, 3); let b = BooleanArray::from(vec![true, false, true]);
1527 let result = filter(&a, &b).unwrap();
1528
1529 let result = result.as_run::<Int64Type>();
1530 let result = result.downcast::<Int64Array>().unwrap();
1531
1532 let expected = vec![-2, 9];
1533 let actual = result.into_iter().flatten().collect::<Vec<_>>();
1534 assert_eq!(expected, actual);
1535 }
1536
1537 #[test]
1538 fn test_filter_run_end_encoding_array_remove_value() {
1539 let run_ends = Int32Array::from(vec![2, 3, 8, 10]);
1540 let values = Int32Array::from(vec![7, -2, 9, -8]);
1541 let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1542 let b = BooleanArray::from(vec![
1543 false, true, false, false, true, false, true, false, false, false,
1544 ]);
1545 let c = filter(&a, &b).unwrap();
1546 let actual: &RunArray<Int32Type> = as_run_array(&c);
1547 assert_eq!(3, actual.len());
1548
1549 let expected =
1550 RunArray::try_new(&Int32Array::from(vec![1, 3]), &Int32Array::from(vec![7, 9]))
1551 .expect("Failed to make expected RunArray test is broken");
1552
1553 assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1554 assert_eq!(actual.values(), expected.values())
1555 }
1556
1557 #[test]
1558 fn test_filter_run_end_encoding_array_remove_all_but_one() {
1559 let run_ends = Int16Array::from(vec![2, 3, 8, 10]);
1560 let values = Int16Array::from(vec![7, -2, 9, -8]);
1561 let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1562 let b = BooleanArray::from(vec![
1563 false, false, false, false, false, false, true, false, false, false,
1564 ]);
1565 let c = filter(&a, &b).unwrap();
1566 let actual: &RunArray<Int16Type> = as_run_array(&c);
1567 assert_eq!(1, actual.len());
1568
1569 let expected = RunArray::try_new(&Int16Array::from(vec![1]), &Int16Array::from(vec![9]))
1570 .expect("Failed to make expected RunArray test is broken");
1571
1572 assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1573 assert_eq!(actual.values(), expected.values())
1574 }
1575
1576 #[test]
1577 fn test_filter_run_end_encoding_array_empty() {
1578 let run_ends = Int64Array::from(vec![2, 3, 8, 10]);
1579 let values = Int64Array::from(vec![7, -2, 9, -8]);
1580 let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1581 let b = BooleanArray::from(vec![
1582 false, false, false, false, false, false, false, false, false, false,
1583 ]);
1584 let c = filter(&a, &b).unwrap();
1585 let actual: &RunArray<Int64Type> = as_run_array(&c);
1586 assert_eq!(0, actual.len());
1587 }
1588
1589 #[test]
1590 fn test_filter_run_end_encoding_array_max_value_gt_predicate_len() {
1591 let run_ends = Int64Array::from(vec![2, 3, 8, 10]);
1592 let values = Int64Array::from(vec![7, -2, 9, -8]);
1593 let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1594 let b = BooleanArray::from(vec![false, true, true]);
1595 let c = filter(&a, &b).unwrap();
1596 let actual: &RunArray<Int64Type> = as_run_array(&c);
1597 assert_eq!(2, actual.len());
1598
1599 let expected = RunArray::try_new(
1600 &Int64Array::from(vec![1, 2]),
1601 &Int64Array::from(vec![7, -2]),
1602 )
1603 .expect("Failed to make expected RunArray test is broken");
1604
1605 assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1606 assert_eq!(actual.values(), expected.values())
1607 }
1608
1609 #[test]
1610 fn test_filter_dictionary_array() {
1611 let values = [Some("hello"), None, Some("world"), Some("!")];
1612 let a: Int8DictionaryArray = values.iter().copied().collect();
1613 let b = BooleanArray::from(vec![false, true, true, false]);
1614 let c = filter(&a, &b).unwrap();
1615 let d = c
1616 .as_ref()
1617 .as_any()
1618 .downcast_ref::<Int8DictionaryArray>()
1619 .unwrap();
1620 let value_array = d.values();
1621 let values = value_array.as_any().downcast_ref::<StringArray>().unwrap();
1622 assert_eq!(3, values.len());
1624 assert_eq!(2, d.len());
1626 assert!(d.is_null(0));
1627 assert_eq!("world", values.value(d.keys().value(1) as usize));
1628 }
1629
1630 #[test]
1631 fn test_filter_list_array() {
1632 let field = Arc::new(Field::new_list_field(DataType::Int32, false));
1633 let offsets = OffsetBuffer::new(vec![0i64, 3, 6, 8, 8].into());
1634 let value_array = Arc::new(Int32Array::from_iter_values(0..8));
1635 let nulls = Some(NullBuffer::from(vec![true, true, true, false]));
1636 let a = LargeListArray::new(field.clone(), offsets, value_array, nulls);
1638 let b = BooleanArray::from(vec![false, true, false, true]);
1639 let result = filter(&a, &b).unwrap();
1640
1641 let offsets = OffsetBuffer::new(vec![0i64, 3, 3].into());
1643 let value_array = Arc::new(Int32Array::from_iter_values([3, 4, 5]));
1644 let nulls = Some(NullBuffer::from(vec![true, false]));
1645 let expected: ArrayRef = Arc::new(LargeListArray::new(field, offsets, value_array, nulls));
1646
1647 assert_eq!(&expected, &result);
1648 }
1649
1650 fn test_case_filter_list_view<T: OffsetSizeTrait>() {
1651 let mut list_array = GenericListViewBuilder::<T, _>::new(Int32Builder::new());
1653 list_array.append_value([Some(1), Some(2)]);
1654 list_array.append_null();
1655 list_array.append_value([]);
1656 list_array.append_value([Some(3), Some(4)]);
1657
1658 let list_array = list_array.finish();
1659 let predicate = BooleanArray::from_iter([true, false, true, false]);
1660
1661 let filtered = filter(&list_array, &predicate)
1663 .unwrap()
1664 .as_list_view::<T>()
1665 .clone();
1666
1667 let mut expected =
1668 GenericListViewBuilder::<T, _>::with_capacity(Int32Builder::with_capacity(5), 3);
1669 expected.append_value([Some(1), Some(2)]);
1670 expected.append_value([]);
1671 let expected = expected.finish();
1672
1673 assert_eq!(&filtered, &expected);
1674 }
1675
1676 fn test_case_filter_sliced_list_view<T: OffsetSizeTrait>() {
1677 let mut list_array =
1679 GenericListViewBuilder::<T, _>::with_capacity(Int32Builder::with_capacity(6), 4);
1680 list_array.append_value([Some(1), Some(2)]);
1681 list_array.append_null();
1682 list_array.append_value([]);
1683 list_array.append_value([Some(3), Some(4)]);
1684
1685 let list_array = list_array.finish();
1686
1687 let sliced = list_array.slice(1, 3);
1689 let predicate = BooleanArray::from_iter([false, false, true]);
1690
1691 let filtered = filter(&sliced, &predicate)
1693 .unwrap()
1694 .as_list_view::<T>()
1695 .clone();
1696
1697 let mut expected = GenericListViewBuilder::<T, _>::new(Int32Builder::new());
1698 expected.append_value([Some(3), Some(4)]);
1699 let expected = expected.finish();
1700
1701 assert_eq!(&filtered, &expected);
1702 }
1703
1704 #[test]
1705 fn test_filter_list_view_array() {
1706 test_case_filter_list_view::<i32>();
1707 test_case_filter_list_view::<i64>();
1708
1709 test_case_filter_sliced_list_view::<i32>();
1710 test_case_filter_sliced_list_view::<i64>();
1711 }
1712
1713 #[test]
1714 fn test_slice_iterator_bits() {
1715 let filter_values = (0..64).map(|i| i == 1).collect::<Vec<bool>>();
1716 let filter = BooleanArray::from(filter_values);
1717 let filter_count = filter.true_count();
1718
1719 let iter = SlicesIterator::new(&filter);
1720 let chunks = iter.collect::<Vec<_>>();
1721
1722 assert_eq!(chunks, vec![(1, 2)]);
1723 assert_eq!(filter_count, 1);
1724 }
1725
1726 #[test]
1727 fn test_slice_iterator_bits1() {
1728 let filter_values = (0..64).map(|i| i != 1).collect::<Vec<bool>>();
1729 let filter = BooleanArray::from(filter_values);
1730 let filter_count = filter.true_count();
1731
1732 let iter = SlicesIterator::new(&filter);
1733 let chunks = iter.collect::<Vec<_>>();
1734
1735 assert_eq!(chunks, vec![(0, 1), (2, 64)]);
1736 assert_eq!(filter_count, 64 - 1);
1737 }
1738
1739 #[test]
1740 fn test_slice_iterator_chunk_and_bits() {
1741 let filter_values = (0..130).map(|i| i % 62 != 0).collect::<Vec<bool>>();
1742 let filter = BooleanArray::from(filter_values);
1743 let filter_count = filter.true_count();
1744
1745 let iter = SlicesIterator::new(&filter);
1746 let chunks = iter.collect::<Vec<_>>();
1747
1748 assert_eq!(chunks, vec![(1, 62), (63, 124), (125, 130)]);
1749 assert_eq!(filter_count, 61 + 61 + 5);
1750 }
1751
1752 #[test]
1753 fn test_filter_selection_iterators() {
1754 let slices = [(0, 2), (4, 5)];
1755 let mut ranges = Vec::new();
1756 let selection: FilterSlices<'_> = FilterIterator::Materialized(slices.iter().copied());
1757 selection.for_each(|range| ranges.push(range));
1758 assert_eq!(ranges, slices);
1759
1760 let filter = BooleanArray::from(vec![true, true, false, false, true]);
1761 let mut ranges = Vec::new();
1762 let selection: FilterSlices<'_> = FilterIterator::Lazy(SlicesIterator::new(&filter));
1763 selection
1764 .try_for_each(|range| {
1765 ranges.push(range);
1766 Ok::<(), ArrowError>(())
1767 })
1768 .unwrap();
1769 assert_eq!(ranges, vec![(0, 2), (4, 5)]);
1770
1771 let indices = [1, 3, 5];
1772 let mut selected = Vec::new();
1773 let selection: FilterIndices<'_> = FilterIterator::Materialized(indices.iter().copied());
1774 selection.for_each(|idx| selected.push(idx));
1775 assert_eq!(selected, indices);
1776
1777 let filter = BooleanArray::from(vec![false, true, false, true]);
1778 let mut selected = Vec::new();
1779 let selection: FilterIndices<'_> = FilterIterator::Lazy(IndexIterator::new(&filter, 2));
1780 selection
1781 .try_for_each(|idx| {
1782 selected.push(idx);
1783 Ok::<(), ArrowError>(())
1784 })
1785 .unwrap();
1786 assert_eq!(selected, vec![1, 3]);
1787 }
1788
1789 #[test]
1790 fn test_null_mask() {
1791 let a = Int64Array::from(vec![Some(1), Some(2), None]);
1792
1793 let mask1 = BooleanArray::from(vec![Some(true), Some(true), None]);
1794 let out = filter(&a, &mask1).unwrap();
1795 assert_eq!(out.as_ref(), &a.slice(0, 2));
1796 }
1797
1798 #[test]
1799 fn test_filter_record_batch_no_columns() {
1800 let pred = BooleanArray::from(vec![Some(true), Some(true), None]);
1801 let options = RecordBatchOptions::default().with_row_count(Some(100));
1802 let record_batch =
1803 RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap();
1804 let out = filter_record_batch(&record_batch, &pred).unwrap();
1805
1806 assert_eq!(out.num_rows(), 2);
1807 }
1808
1809 #[test]
1810 fn test_fast_path() {
1811 let a: PrimitiveArray<Int64Type> = PrimitiveArray::from(vec![Some(1), Some(2), None]);
1812
1813 let mask = BooleanArray::from(vec![true, true, true]);
1815 let out = filter(&a, &mask).unwrap();
1816 let b = out
1817 .as_any()
1818 .downcast_ref::<PrimitiveArray<Int64Type>>()
1819 .unwrap();
1820 assert_eq!(&a, b);
1821
1822 let mask = BooleanArray::from(vec![false, false, false]);
1824 let out = filter(&a, &mask).unwrap();
1825 assert_eq!(out.len(), 0);
1826 assert_eq!(out.data_type(), &DataType::Int64);
1827 }
1828
1829 #[test]
1830 fn test_slices() {
1831 let bools = std::iter::repeat_n(true, 10)
1833 .chain(std::iter::repeat_n(false, 30))
1834 .chain(std::iter::repeat_n(true, 20))
1835 .chain(std::iter::repeat_n(false, 17))
1836 .chain(std::iter::repeat_n(true, 4));
1837
1838 let bool_array: BooleanArray = bools.map(Some).collect();
1839
1840 let slices: Vec<_> = SlicesIterator::new(&bool_array).collect();
1841 let expected = vec![(0, 10), (40, 60), (77, 81)];
1842 assert_eq!(slices, expected);
1843
1844 let len = bool_array.len();
1846 let sliced_array = bool_array.slice(7, len - 10);
1847 let sliced_array = sliced_array
1848 .as_any()
1849 .downcast_ref::<BooleanArray>()
1850 .unwrap();
1851 let slices: Vec<_> = SlicesIterator::new(sliced_array).collect();
1852 let expected = vec![(0, 3), (33, 53), (70, 71)];
1853 assert_eq!(slices, expected);
1854 }
1855
1856 fn test_slices_fuzz(mask_len: usize, offset: usize, truncate: usize) {
1857 let mut rng = rng();
1858
1859 let bools: Vec<bool> = std::iter::from_fn(|| Some(rng.random()))
1860 .take(mask_len)
1861 .collect();
1862
1863 let buffer = Buffer::from_iter(bools.iter().cloned());
1864
1865 let truncated_length = mask_len - offset - truncate;
1866
1867 let filter = BooleanArray::new(BooleanBuffer::new(buffer, offset, truncated_length), None);
1868
1869 let slice_bits: Vec<_> = SlicesIterator::new(&filter)
1870 .flat_map(|(start, end)| start..end)
1871 .collect();
1872
1873 let count = filter.true_count();
1874 let index_bits: Vec<_> = IndexIterator::new(&filter, count).collect();
1875
1876 let expected_bits: Vec<_> = bools
1877 .iter()
1878 .skip(offset)
1879 .take(truncated_length)
1880 .enumerate()
1881 .flat_map(|(idx, v)| v.then(|| idx))
1882 .collect();
1883
1884 assert_eq!(slice_bits, expected_bits);
1885 assert_eq!(index_bits, expected_bits);
1886 }
1887
1888 #[test]
1889 #[cfg_attr(miri, ignore)]
1890 fn fuzz_test_slices_iterator() {
1891 let mut rng = rng();
1892
1893 let uusize = UniformUsize::new(usize::MIN, usize::MAX).unwrap();
1894 for _ in 0..100 {
1895 let mask_len = rng.random_range(0..1024);
1896 let max_offset = 64.min(mask_len);
1897 let offset = uusize.sample(&mut rng).checked_rem(max_offset).unwrap_or(0);
1898
1899 let max_truncate = 128.min(mask_len - offset);
1900 let truncate = uusize
1901 .sample(&mut rng)
1902 .checked_rem(max_truncate)
1903 .unwrap_or(0);
1904
1905 test_slices_fuzz(mask_len, offset, truncate);
1906 }
1907
1908 test_slices_fuzz(64, 0, 0);
1909 test_slices_fuzz(64, 8, 0);
1910 test_slices_fuzz(64, 8, 8);
1911 test_slices_fuzz(32, 8, 8);
1912 test_slices_fuzz(32, 5, 9);
1913 }
1914
1915 fn filter_rust<T>(values: impl IntoIterator<Item = T>, predicate: &[bool]) -> Vec<T> {
1917 values
1918 .into_iter()
1919 .zip(predicate)
1920 .filter(|(_, x)| **x)
1921 .map(|(a, _)| a)
1922 .collect()
1923 }
1924
1925 fn gen_primitive<T>(len: usize, valid_percent: f64) -> Vec<Option<T>>
1927 where
1928 StandardUniform: Distribution<T>,
1929 {
1930 let mut rng = rng();
1931 (0..len)
1932 .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
1933 .collect()
1934 }
1935
1936 fn gen_strings(
1938 len: usize,
1939 valid_percent: f64,
1940 str_len_range: std::ops::Range<usize>,
1941 ) -> Vec<Option<String>> {
1942 let mut rng = rng();
1943 (0..len)
1944 .map(|_| {
1945 rng.random_bool(valid_percent).then(|| {
1946 let len = rng.random_range(str_len_range.clone());
1947 (0..len)
1948 .map(|_| char::from(rng.sample(Alphanumeric)))
1949 .collect()
1950 })
1951 })
1952 .collect()
1953 }
1954
1955 fn as_deref<T: std::ops::Deref>(src: &[Option<T>]) -> impl Iterator<Item = Option<&T::Target>> {
1957 src.iter().map(|x| x.as_deref())
1958 }
1959
1960 #[test]
1961 #[cfg_attr(miri, ignore)]
1962 fn fuzz_filter() {
1963 let mut rng = rng();
1964
1965 for i in 0..100 {
1966 let filter_percent = match i {
1967 0..=4 => 1.,
1968 5..=10 => 0.,
1969 _ => rng.random_range(0.0..1.0),
1970 };
1971
1972 let valid_percent = rng.random_range(0.0..1.0);
1973
1974 let array_len = rng.random_range(32..256);
1975 let array_offset = rng.random_range(0..10);
1976
1977 let filter_offset = rng.random_range(0..10);
1979 let filter_truncate = rng.random_range(0..10);
1980 let bools: Vec<_> = std::iter::from_fn(|| Some(rng.random_bool(filter_percent)))
1981 .take(array_len + filter_offset - filter_truncate)
1982 .collect();
1983
1984 let predicate = BooleanArray::from_iter(bools.iter().cloned().map(Some));
1985
1986 let predicate = predicate.slice(filter_offset, array_len - filter_truncate);
1988 let predicate = predicate.as_any().downcast_ref::<BooleanArray>().unwrap();
1989 let bools = &bools[filter_offset..];
1990
1991 let values = gen_primitive(array_len + array_offset, valid_percent);
1993 let src = Int32Array::from_iter(values.iter().cloned());
1994
1995 let src = src.slice(array_offset, array_len);
1996 let src = src.as_any().downcast_ref::<Int32Array>().unwrap();
1997 let values = &values[array_offset..];
1998
1999 let filtered = filter(src, predicate).unwrap();
2000 let array = filtered.as_any().downcast_ref::<Int32Array>().unwrap();
2001 let actual: Vec<_> = array.iter().collect();
2002
2003 assert_eq!(actual, filter_rust(values.iter().cloned(), bools));
2004
2005 let strings = gen_strings(array_len + array_offset, valid_percent, 0..20);
2007 let src = StringArray::from_iter(as_deref(&strings));
2008
2009 let src = src.slice(array_offset, array_len);
2010 let src = src.as_any().downcast_ref::<StringArray>().unwrap();
2011
2012 let filtered = filter(src, predicate).unwrap();
2013 let array = filtered.as_any().downcast_ref::<StringArray>().unwrap();
2014 let actual: Vec<_> = array.iter().collect();
2015
2016 let expected_strings = filter_rust(as_deref(&strings[array_offset..]), bools);
2017 assert_eq!(actual, expected_strings);
2018
2019 let src = DictionaryArray::<Int32Type>::from_iter(as_deref(&strings));
2021
2022 let src = src.slice(array_offset, array_len);
2023 let src = src
2024 .as_any()
2025 .downcast_ref::<DictionaryArray<Int32Type>>()
2026 .unwrap();
2027
2028 let filtered = filter(src, predicate).unwrap();
2029
2030 let array = filtered
2031 .as_any()
2032 .downcast_ref::<DictionaryArray<Int32Type>>()
2033 .unwrap();
2034
2035 let values = array
2036 .values()
2037 .as_any()
2038 .downcast_ref::<StringArray>()
2039 .unwrap();
2040
2041 let actual: Vec<_> = array
2042 .keys()
2043 .iter()
2044 .map(|key| key.map(|key| values.value(key as usize)))
2045 .collect();
2046
2047 assert_eq!(actual, expected_strings);
2048 }
2049 }
2050
2051 #[test]
2052 fn test_filter_map() {
2053 let mut builder =
2054 MapBuilder::new(None, StringBuilder::new(), Int64Builder::with_capacity(4));
2055 builder.keys().append_value("key1");
2057 builder.values().append_value(1);
2058 builder.append(true).unwrap();
2059 builder.keys().append_value("key2");
2060 builder.keys().append_value("key3");
2061 builder.values().append_value(2);
2062 builder.values().append_value(3);
2063 builder.append(true).unwrap();
2064 builder.append(false).unwrap();
2065 builder.keys().append_value("key1");
2066 builder.values().append_value(1);
2067 builder.append(true).unwrap();
2068 let maparray = Arc::new(builder.finish()) as ArrayRef;
2069
2070 let indices = vec![Some(true), Some(false), Some(false), Some(true)]
2071 .into_iter()
2072 .collect::<BooleanArray>();
2073 let got = filter(&maparray, &indices).unwrap();
2074
2075 let mut builder =
2076 MapBuilder::new(None, StringBuilder::new(), Int64Builder::with_capacity(2));
2077 builder.keys().append_value("key1");
2078 builder.values().append_value(1);
2079 builder.append(true).unwrap();
2080 builder.keys().append_value("key1");
2081 builder.values().append_value(1);
2082 builder.append(true).unwrap();
2083 let expected = Arc::new(builder.finish()) as ArrayRef;
2084
2085 assert_eq!(&expected, &got);
2086 }
2087
2088 #[test]
2089 fn test_filter_fixed_size_list_arrays() {
2090 let field = Arc::new(Field::new_list_field(DataType::Int32, false));
2091 let value_array = Arc::new(Int32Array::from_iter_values(0..9));
2092 let array = FixedSizeListArray::new(field, 3, value_array, None);
2093
2094 let filter_array = BooleanArray::from(vec![true, false, false]);
2095
2096 let c = filter(&array, &filter_array).unwrap();
2097 let filtered = c.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
2098
2099 assert_eq!(filtered.len(), 1);
2100
2101 let list = filtered.value(0);
2102 assert_eq!(
2103 &[0, 1, 2],
2104 list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2105 );
2106
2107 let filter_array = BooleanArray::from(vec![true, false, true]);
2108
2109 let c = filter(&array, &filter_array).unwrap();
2110 let filtered = c.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
2111
2112 assert_eq!(filtered.len(), 2);
2113
2114 let list = filtered.value(0);
2115 assert_eq!(
2116 &[0, 1, 2],
2117 list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2118 );
2119 let list = filtered.value(1);
2120 assert_eq!(
2121 &[6, 7, 8],
2122 list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2123 );
2124 }
2125
2126 #[test]
2127 fn test_filter_fixed_size_list_arrays_with_null() {
2128 let field = Arc::new(Field::new_list_field(DataType::Int32, false));
2129 let value_array = Arc::new(Int32Array::from_iter_values(0..10));
2130 let nulls = Some(NullBuffer::from(vec![true, false, false, true, true]));
2131 let array = FixedSizeListArray::new(field, 2, value_array, nulls);
2132
2133 let filter_array = BooleanArray::from(vec![true, true, false, true, false]);
2134
2135 let c = filter(&array, &filter_array).unwrap();
2136 let filtered = c.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
2137
2138 assert_eq!(filtered.len(), 3);
2139
2140 let list = filtered.value(0);
2141 assert_eq!(
2142 &[0, 1],
2143 list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2144 );
2145 assert!(filtered.is_null(1));
2146 let list = filtered.value(2);
2147 assert_eq!(
2148 &[6, 7],
2149 list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2150 );
2151 }
2152
2153 fn test_filter_union_array(array: UnionArray) {
2154 let filter_array = BooleanArray::from(vec![true, false, false]);
2155 let c = filter(&array, &filter_array).unwrap();
2156 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2157
2158 let mut builder = UnionBuilder::new_dense();
2159 builder.append::<Int32Type>("A", 1).unwrap();
2160 let expected_array = builder.build().unwrap();
2161
2162 compare_union_arrays(filtered, &expected_array);
2163
2164 let filter_array = BooleanArray::from(vec![true, false, true]);
2165 let c = filter(&array, &filter_array).unwrap();
2166 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2167
2168 let mut builder = UnionBuilder::new_dense();
2169 builder.append::<Int32Type>("A", 1).unwrap();
2170 builder.append::<Int32Type>("A", 34).unwrap();
2171 let expected_array = builder.build().unwrap();
2172
2173 compare_union_arrays(filtered, &expected_array);
2174
2175 let filter_array = BooleanArray::from(vec![true, true, false]);
2176 let c = filter(&array, &filter_array).unwrap();
2177 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2178
2179 let mut builder = UnionBuilder::new_dense();
2180 builder.append::<Int32Type>("A", 1).unwrap();
2181 builder.append::<Float64Type>("B", 3.2).unwrap();
2182 let expected_array = builder.build().unwrap();
2183
2184 compare_union_arrays(filtered, &expected_array);
2185 }
2186
2187 #[test]
2188 fn test_filter_union_array_dense() {
2189 let mut builder = UnionBuilder::new_dense();
2190 builder.append::<Int32Type>("A", 1).unwrap();
2191 builder.append::<Float64Type>("B", 3.2).unwrap();
2192 builder.append::<Int32Type>("A", 34).unwrap();
2193 let array = builder.build().unwrap();
2194
2195 test_filter_union_array(array);
2196 }
2197
2198 #[test]
2199 fn test_filter_run_union_array_dense() {
2200 let mut builder = UnionBuilder::new_dense();
2201 builder.append::<Int32Type>("A", 1).unwrap();
2202 builder.append::<Int32Type>("A", 3).unwrap();
2203 builder.append::<Int32Type>("A", 34).unwrap();
2204 let array = builder.build().unwrap();
2205
2206 let filter_array = BooleanArray::from(vec![true, true, false]);
2207 let c = filter(&array, &filter_array).unwrap();
2208 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2209
2210 let mut builder = UnionBuilder::new_dense();
2211 builder.append::<Int32Type>("A", 1).unwrap();
2212 builder.append::<Int32Type>("A", 3).unwrap();
2213 let expected = builder.build().unwrap();
2214
2215 assert_eq!(filtered.to_data(), expected.to_data());
2216 }
2217
2218 #[test]
2219 fn test_filter_union_array_dense_with_nulls() {
2220 let mut builder = UnionBuilder::new_dense();
2221 builder.append::<Int32Type>("A", 1).unwrap();
2222 builder.append::<Float64Type>("B", 3.2).unwrap();
2223 builder.append_null::<Float64Type>("B").unwrap();
2224 builder.append::<Int32Type>("A", 34).unwrap();
2225 let array = builder.build().unwrap();
2226
2227 let filter_array = BooleanArray::from(vec![true, true, false, false]);
2228 let c = filter(&array, &filter_array).unwrap();
2229 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2230
2231 let mut builder = UnionBuilder::new_dense();
2232 builder.append::<Int32Type>("A", 1).unwrap();
2233 builder.append::<Float64Type>("B", 3.2).unwrap();
2234 let expected_array = builder.build().unwrap();
2235
2236 compare_union_arrays(filtered, &expected_array);
2237
2238 let filter_array = BooleanArray::from(vec![true, false, true, false]);
2239 let c = filter(&array, &filter_array).unwrap();
2240 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2241
2242 let mut builder = UnionBuilder::new_dense();
2243 builder.append::<Int32Type>("A", 1).unwrap();
2244 builder.append_null::<Float64Type>("B").unwrap();
2245 let expected_array = builder.build().unwrap();
2246
2247 compare_union_arrays(filtered, &expected_array);
2248 }
2249
2250 #[test]
2251 fn test_filter_union_array_sparse() {
2252 let mut builder = UnionBuilder::new_sparse();
2253 builder.append::<Int32Type>("A", 1).unwrap();
2254 builder.append::<Float64Type>("B", 3.2).unwrap();
2255 builder.append::<Int32Type>("A", 34).unwrap();
2256 let array = builder.build().unwrap();
2257
2258 test_filter_union_array(array);
2259 }
2260
2261 #[test]
2262 fn test_filter_union_array_sparse_with_nulls() {
2263 let mut builder = UnionBuilder::new_sparse();
2264 builder.append::<Int32Type>("A", 1).unwrap();
2265 builder.append::<Float64Type>("B", 3.2).unwrap();
2266 builder.append_null::<Float64Type>("B").unwrap();
2267 builder.append::<Int32Type>("A", 34).unwrap();
2268 let array = builder.build().unwrap();
2269
2270 let filter_array = BooleanArray::from(vec![true, false, true, false]);
2271 let c = filter(&array, &filter_array).unwrap();
2272 let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2273
2274 let mut builder = UnionBuilder::new_sparse();
2275 builder.append::<Int32Type>("A", 1).unwrap();
2276 builder.append_null::<Float64Type>("B").unwrap();
2277 let expected_array = builder.build().unwrap();
2278
2279 compare_union_arrays(filtered, &expected_array);
2280 }
2281
2282 fn compare_union_arrays(union1: &UnionArray, union2: &UnionArray) {
2283 assert_eq!(union1.len(), union2.len());
2284
2285 for i in 0..union1.len() {
2286 let type_id = union1.type_id(i);
2287
2288 let slot1 = union1.value(i);
2289 let slot2 = union2.value(i);
2290
2291 assert_eq!(slot1.is_null(0), slot2.is_null(0));
2292
2293 if !slot1.is_null(0) && !slot2.is_null(0) {
2294 match type_id {
2295 0 => {
2296 let slot1 = slot1.as_any().downcast_ref::<Int32Array>().unwrap();
2297 assert_eq!(slot1.len(), 1);
2298 let value1 = slot1.value(0);
2299
2300 let slot2 = slot2.as_any().downcast_ref::<Int32Array>().unwrap();
2301 assert_eq!(slot2.len(), 1);
2302 let value2 = slot2.value(0);
2303 assert_eq!(value1, value2);
2304 }
2305 1 => {
2306 let slot1 = slot1.as_any().downcast_ref::<Float64Array>().unwrap();
2307 assert_eq!(slot1.len(), 1);
2308 let value1 = slot1.value(0);
2309
2310 let slot2 = slot2.as_any().downcast_ref::<Float64Array>().unwrap();
2311 assert_eq!(slot2.len(), 1);
2312 let value2 = slot2.value(0);
2313 assert_eq!(value1, value2);
2314 }
2315 _ => unreachable!(),
2316 }
2317 }
2318 }
2319 }
2320
2321 #[test]
2322 fn test_filter_struct() {
2323 let predicate = BooleanArray::from(vec![true, false, true, false]);
2324
2325 let a = Arc::new(StringArray::from(vec!["hello", " ", "world", "!"]));
2326 let a_filtered = Arc::new(StringArray::from(vec!["hello", "world"]));
2327
2328 let b = Arc::new(Int32Array::from(vec![5, 6, 7, 8]));
2329 let b_filtered = Arc::new(Int32Array::from(vec![5, 7]));
2330
2331 let null_mask = NullBuffer::from(vec![true, false, false, true]);
2332 let null_mask_filtered = NullBuffer::from(vec![true, false]);
2333
2334 let a_field = Field::new("a", DataType::Utf8, false);
2335 let b_field = Field::new("b", DataType::Int32, false);
2336
2337 let array = StructArray::new(vec![a_field.clone()].into(), vec![a.clone()], None);
2338 let expected =
2339 StructArray::new(vec![a_field.clone()].into(), vec![a_filtered.clone()], None);
2340
2341 let result = filter(&array, &predicate).unwrap();
2342
2343 assert_eq!(result.to_data(), expected.to_data());
2344
2345 let array = StructArray::new(
2346 vec![a_field.clone()].into(),
2347 vec![a.clone()],
2348 Some(null_mask.clone()),
2349 );
2350 let expected = StructArray::new(
2351 vec![a_field.clone()].into(),
2352 vec![a_filtered.clone()],
2353 Some(null_mask_filtered.clone()),
2354 );
2355
2356 let result = filter(&array, &predicate).unwrap();
2357
2358 assert_eq!(result.to_data(), expected.to_data());
2359
2360 let array = StructArray::new(
2361 vec![a_field.clone(), b_field.clone()].into(),
2362 vec![a.clone(), b.clone()],
2363 None,
2364 );
2365 let expected = StructArray::new(
2366 vec![a_field.clone(), b_field.clone()].into(),
2367 vec![a_filtered.clone(), b_filtered.clone()],
2368 None,
2369 );
2370
2371 let result = filter(&array, &predicate).unwrap();
2372
2373 assert_eq!(result.to_data(), expected.to_data());
2374
2375 let array = StructArray::new(
2376 vec![a_field.clone(), b_field.clone()].into(),
2377 vec![a.clone(), b.clone()],
2378 Some(null_mask.clone()),
2379 );
2380
2381 let expected = StructArray::new(
2382 vec![a_field.clone(), b_field.clone()].into(),
2383 vec![a_filtered.clone(), b_filtered.clone()],
2384 Some(null_mask_filtered.clone()),
2385 );
2386
2387 let result = filter(&array, &predicate).unwrap();
2388
2389 assert_eq!(result.to_data(), expected.to_data());
2390 }
2391
2392 #[test]
2393 fn test_filter_empty_struct() {
2394 let fields = arrow_schema::Field::new(
2401 "a",
2402 arrow_schema::DataType::Struct(arrow_schema::Fields::from(vec![
2403 arrow_schema::Field::new("b", arrow_schema::DataType::Int64, true),
2404 arrow_schema::Field::new(
2405 "c",
2406 arrow_schema::DataType::Struct(arrow_schema::Fields::empty()),
2407 true,
2408 ),
2409 ])),
2410 true,
2411 );
2412
2413 let schema = Arc::new(Schema::new(vec![fields]));
2421
2422 let b = Arc::new(Int64Array::from(vec![None, None, None]));
2423 let c = Arc::new(StructArray::new_empty_fields(
2424 3,
2425 Some(NullBuffer::from(vec![true, true, true])),
2426 ));
2427 let a = StructArray::new(
2428 vec![
2429 Field::new("b", DataType::Int64, true),
2430 Field::new("c", DataType::Struct(Fields::empty()), true),
2431 ]
2432 .into(),
2433 vec![b.clone(), c.clone()],
2434 Some(NullBuffer::from(vec![true, true, true])),
2435 );
2436 let record_batch = RecordBatch::try_new(schema, vec![Arc::new(a)]).unwrap();
2437 println!("{record_batch:?}");
2438
2439 let predicate = BooleanArray::from(vec![true, false, true]);
2441 let filtered_batch = filter_record_batch(&record_batch, &predicate).unwrap();
2442
2443 assert_eq!(filtered_batch.num_rows(), 2);
2445 }
2446
2447 #[test]
2448 #[should_panic]
2449 fn test_filter_bits_too_large() {
2450 let buffer = BooleanBuffer::from(vec![false; 8]);
2451 let predicate = BooleanArray::from(vec![true; 9]);
2452 let filter = FilterBuilder::new(&predicate).build();
2453 filter_bits(&buffer, &filter);
2454 }
2455
2456 #[test]
2457 #[should_panic]
2458 fn test_filter_native_too_large() {
2459 let values = vec![1; 8];
2460 let predicate = BooleanArray::from(vec![false; 9]);
2461 let filter = FilterBuilder::new(&predicate).build();
2462 filter_native(&values, &filter);
2463 }
2464}