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