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