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