1use crate::filter::{FilterBuilder, FilterPredicate, FilterSelection};
24use crate::take::take_record_batch;
25use arrow_array::types::{BinaryViewType, StringViewType};
26use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch, downcast_primitive};
27use arrow_schema::{ArrowError, DataType, SchemaRef};
28use std::collections::VecDeque;
29use std::sync::Arc;
30mod byte_view;
34mod generic;
35mod primitive;
36
37use byte_view::InProgressByteViewArray;
38use generic::GenericInProgressArray;
39use primitive::InProgressPrimitiveArray;
40
41fn has_sparse_filter_copy(data_type: &DataType) -> bool {
42 data_type.is_primitive() || matches!(data_type, DataType::Utf8View | DataType::BinaryView)
43}
44
45const SPARSE_FILTER_COPY_MAX_SELECTIVITY_DENOMINATOR: usize = 16;
51
52fn should_use_sparse_filter_copy(filter_len: usize, selected_count: usize) -> bool {
53 selected_count <= filter_len / SPARSE_FILTER_COPY_MAX_SELECTIVITY_DENOMINATOR
54}
55
56#[derive(Debug)]
148pub struct BatchCoalescer {
149 schema: SchemaRef,
151 target_batch_size: usize,
155 in_progress_arrays: Vec<Box<dyn InProgressArray>>,
157 has_non_specialized_filter_columns: bool,
159 buffered_rows: usize,
161 completed: VecDeque<RecordBatch>,
163 biggest_coalesce_batch_size: Option<usize>,
165}
166
167impl BatchCoalescer {
168 pub fn new(schema: SchemaRef, target_batch_size: usize) -> Self {
176 let has_non_specialized_filter_columns = schema
177 .fields()
178 .iter()
179 .any(|field| !has_sparse_filter_copy(field.data_type()));
180 let in_progress_arrays = schema
181 .fields()
182 .iter()
183 .map(|field| create_in_progress_array(field.data_type(), target_batch_size))
184 .collect::<Vec<_>>();
185
186 Self {
187 schema,
188 target_batch_size,
189 in_progress_arrays,
190 has_non_specialized_filter_columns,
191 completed: VecDeque::with_capacity(1),
193 buffered_rows: 0,
194 biggest_coalesce_batch_size: None,
195 }
196 }
197
198 pub fn with_biggest_coalesce_batch_size(mut self, limit: Option<usize>) -> Self {
211 self.biggest_coalesce_batch_size = limit;
212 self
213 }
214
215 pub fn biggest_coalesce_batch_size(&self) -> Option<usize> {
219 self.biggest_coalesce_batch_size
220 }
221
222 pub fn set_biggest_coalesce_batch_size(&mut self, limit: Option<usize>) {
226 self.biggest_coalesce_batch_size = limit;
227 }
228
229 pub fn schema(&self) -> SchemaRef {
231 Arc::clone(&self.schema)
232 }
233
234 pub fn push_batch_with_filter(
259 &mut self,
260 batch: RecordBatch,
261 filter: &BooleanArray,
262 ) -> Result<(), ArrowError> {
263 self.push_batch_with_filtered_columns(batch, filter)
264 }
265
266 pub fn push_batch_with_indices(
290 &mut self,
291 batch: RecordBatch,
292 indices: &dyn Array,
293 ) -> Result<(), ArrowError> {
294 let taken_batch = take_record_batch(&batch, indices)?;
296 self.push_batch(taken_batch)
297 }
298
299 pub fn push_batch(&mut self, batch: RecordBatch) -> Result<(), ArrowError> {
326 let batch_size = batch.num_rows();
445
446 if batch_size == 0 {
448 return Ok(());
449 }
450
451 if let Some(limit) = self.biggest_coalesce_batch_size
453 && batch_size > limit
454 {
455 if self.buffered_rows == 0 {
458 self.completed.push_back(batch);
459 return Ok(());
460 }
461
462 if self.buffered_rows > limit {
467 self.finish_buffered_batch()?;
468 self.completed.push_back(batch);
469 return Ok(());
470 }
471
472 }
477
478 let (_schema, arrays, mut num_rows) = batch.into_parts();
479
480 if arrays.len() != self.in_progress_arrays.len() {
482 return Err(ArrowError::InvalidArgumentError(format!(
483 "Batch has {} columns but BatchCoalescer expects {}",
484 arrays.len(),
485 self.in_progress_arrays.len()
486 )));
487 }
488 self.in_progress_arrays
489 .iter_mut()
490 .zip(arrays)
491 .for_each(|(in_progress, array)| {
492 in_progress.set_source(Some(array));
493 });
494
495 let mut offset = 0;
498 while num_rows > (self.target_batch_size - self.buffered_rows) {
499 let remaining_rows = self.target_batch_size - self.buffered_rows;
500 debug_assert!(remaining_rows > 0);
501
502 for in_progress in self.in_progress_arrays.iter_mut() {
504 in_progress.copy_rows(offset, remaining_rows)?;
505 }
506
507 self.buffered_rows += remaining_rows;
508 offset += remaining_rows;
509 num_rows -= remaining_rows;
510
511 self.finish_buffered_batch()?;
512 }
513
514 self.buffered_rows += num_rows;
516 if num_rows > 0 {
517 for in_progress in self.in_progress_arrays.iter_mut() {
518 in_progress.copy_rows(offset, num_rows)?;
519 }
520 }
521
522 if self.buffered_rows >= self.target_batch_size {
524 self.finish_buffered_batch()?;
525 }
526
527 for in_progress in self.in_progress_arrays.iter_mut() {
529 in_progress.set_source(None);
530 }
531
532 Ok(())
533 }
534
535 pub fn get_buffered_rows(&self) -> usize {
537 self.buffered_rows
538 }
539
540 pub fn finish_buffered_batch(&mut self) -> Result<(), ArrowError> {
548 if self.buffered_rows == 0 {
549 return Ok(());
550 }
551 let new_arrays = self
552 .in_progress_arrays
553 .iter_mut()
554 .map(|array| array.finish())
555 .collect::<Result<Vec<_>, ArrowError>>()?;
556
557 for (array, field) in new_arrays.iter().zip(self.schema.fields().iter()) {
558 debug_assert_eq!(array.data_type(), field.data_type());
559 debug_assert_eq!(array.len(), self.buffered_rows);
560 }
561
562 let batch = unsafe {
564 RecordBatch::new_unchecked(Arc::clone(&self.schema), new_arrays, self.buffered_rows)
565 };
566
567 self.buffered_rows = 0;
568 self.completed.push_back(batch);
569 Ok(())
570 }
571
572 pub fn is_empty(&self) -> bool {
574 self.buffered_rows == 0 && self.completed.is_empty()
575 }
576
577 pub fn has_completed_batch(&self) -> bool {
579 !self.completed.is_empty()
580 }
581
582 pub fn next_completed_batch(&mut self) -> Option<RecordBatch> {
584 self.completed.pop_front()
585 }
586
587 pub fn size(&self) -> usize {
589 self.in_progress_arrays.capacity() * size_of::<Box<dyn InProgressArray>>()
590 + self
591 .in_progress_arrays
592 .iter()
593 .map(|array| array.size())
594 .sum::<usize>()
595 + self.completed.capacity() * size_of::<RecordBatch>()
596 + self
597 .completed
598 .iter()
599 .map(|batch| batch.get_array_memory_size())
600 .sum::<usize>()
601 }
602}
603
604impl BatchCoalescer {
605 fn filter_predicate_for_batch(
606 batch: &RecordBatch,
607 filter: &BooleanArray,
608 selected_count: usize,
609 ) -> FilterPredicate {
610 let mut filter_builder = FilterBuilder::new_with_count(filter, selected_count);
611 if batch.num_columns() > 1
612 || (batch.num_columns() > 0
613 && FilterBuilder::is_optimize_beneficial(batch.schema_ref().field(0).data_type()))
614 {
615 filter_builder = filter_builder.optimize();
616 }
617 filter_builder.build()
618 }
619
620 fn push_batch_with_filtered_columns(
621 &mut self,
622 batch: RecordBatch,
623 filter: &BooleanArray,
624 ) -> Result<(), ArrowError> {
625 let filter_len = filter.len();
626 let batch_num_rows = batch.num_rows();
627 let batch_num_columns = batch.num_columns();
628
629 if filter_len > batch_num_rows {
630 return Err(ArrowError::InvalidArgumentError(format!(
631 "Filter predicate of length {} is larger than target array of length {}",
632 filter_len, batch_num_rows
633 )));
634 }
635
636 let selected_count = filter.true_count();
637 if selected_count == 0 {
638 return Ok(());
639 }
640
641 if selected_count == batch_num_rows && filter_len == batch_num_rows {
642 return self.push_batch(batch);
643 }
644
645 if batch_num_columns != self.in_progress_arrays.len() {
646 return Err(ArrowError::InvalidArgumentError(format!(
647 "Batch has {} columns but BatchCoalescer expects {}",
648 batch_num_columns,
649 self.in_progress_arrays.len()
650 )));
651 }
652
653 let exceeds_coalesce_limit = self
654 .biggest_coalesce_batch_size
655 .is_some_and(|limit| selected_count > limit);
656 let does_not_fit_buffer = selected_count > self.target_batch_size - self.buffered_rows;
657 let should_materialize_filter = exceeds_coalesce_limit
658 || self.has_non_specialized_filter_columns
659 || does_not_fit_buffer
660 || !should_use_sparse_filter_copy(filter_len, selected_count);
661
662 if should_materialize_filter {
663 let predicate = Self::filter_predicate_for_batch(&batch, filter, selected_count);
665 let filtered_batch = predicate.filter_record_batch(&batch)?;
666 return self.push_batch(filtered_batch);
667 }
668
669 let predicate = Self::filter_predicate_for_batch(&batch, filter, selected_count);
670 let (_schema, arrays, _num_rows) = batch.into_parts();
671
672 for (in_progress, array) in self.in_progress_arrays.iter_mut().zip(arrays) {
673 in_progress.copy_rows_by_filter_from(array, &predicate)?;
674 }
675
676 self.buffered_rows += selected_count;
677 if self.buffered_rows >= self.target_batch_size {
678 self.finish_buffered_batch()?;
679 }
680
681 Ok(())
682 }
683}
684
685fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box<dyn InProgressArray> {
687 macro_rules! instantiate_primitive {
688 ($t:ty) => {
689 Box::new(InProgressPrimitiveArray::<$t>::new(
690 batch_size,
691 data_type.clone(),
692 ))
693 };
694 }
695
696 downcast_primitive! {
697 data_type => (instantiate_primitive),
699 DataType::Utf8View => Box::new(InProgressByteViewArray::<StringViewType>::new(batch_size)),
700 DataType::BinaryView => {
701 Box::new(InProgressByteViewArray::<BinaryViewType>::new(batch_size))
702 }
703 _ => Box::new(GenericInProgressArray::new()),
704 }
705}
706
707trait InProgressArray: std::fmt::Debug + Send + Sync {
717 fn set_source(&mut self, source: Option<ArrayRef>);
722
723 fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError>;
729
730 fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(), ArrowError> {
734 self.copy_rows_by_selection(filter.selection())
735 }
736
737 fn copy_rows_by_filter_from(
744 &mut self,
745 source: ArrayRef,
746 filter: &FilterPredicate,
747 ) -> Result<(), ArrowError> {
748 self.set_source(Some(source));
749 let result = self.copy_rows_by_filter(filter);
750 self.set_source(None);
751 result
752 }
753
754 fn copy_rows_by_selection(&mut self, selection: FilterSelection<'_>) -> Result<(), ArrowError> {
760 match selection {
761 FilterSelection::None => Ok(()),
762 FilterSelection::All { len } => self.copy_rows(0, len),
763 FilterSelection::Slices(slices) => {
764 slices.try_for_each(|(start, end)| self.copy_rows(start, end - start))
765 }
766 FilterSelection::Indices(indices) => indices.try_for_each(|idx| self.copy_rows(idx, 1)),
767 }
768 }
769
770 fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
772
773 fn size(&self) -> usize;
775}
776
777#[cfg(test)]
778mod tests {
779 use super::*;
780 use crate::concat::concat_batches;
781 use crate::filter::filter_record_batch;
782 use arrow_array::builder::StringViewBuilder;
783 use arrow_array::cast::AsArray;
784 use arrow_array::types::Int32Type;
785 use arrow_array::{
786 BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray,
787 TimestampNanosecondArray, UInt32Array, UInt64Array, make_array,
788 };
789 use arrow_buffer::BooleanBufferBuilder;
790 use arrow_schema::{DataType, Field, Schema};
791 use rand::{RngExt, SeedableRng};
792 use std::ops::Range;
793
794 #[test]
795 fn test_coalesce() {
796 let batch = uint32_batch(0..8);
797 Test::new("coalesce")
798 .with_batches(std::iter::repeat_n(batch, 10))
799 .with_batch_size(21)
801 .with_expected_output_sizes(vec![21, 21, 21, 17])
802 .run();
803 }
804
805 #[test]
806 fn test_coalesce_one_by_one() {
807 let batch = uint32_batch(0..1); Test::new("coalesce_one_by_one")
809 .with_batches(std::iter::repeat_n(batch, 97))
810 .with_batch_size(20)
812 .with_expected_output_sizes(vec![20, 20, 20, 20, 17])
813 .run();
814 }
815
816 #[test]
817 fn test_coalesce_empty() {
818 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
819
820 Test::new("coalesce_empty")
821 .with_batches(vec![])
822 .with_schema(schema)
823 .with_batch_size(21)
824 .with_expected_output_sizes(vec![])
825 .run();
826 }
827
828 #[test]
829 fn test_sparse_filter_copy_threshold() {
830 assert!(should_use_sparse_filter_copy(8192, 8));
831 assert!(should_use_sparse_filter_copy(8192, 81));
832 assert!(!should_use_sparse_filter_copy(8192, 819));
833 assert!(!should_use_sparse_filter_copy(8192, 6553));
834 }
835
836 #[test]
837 fn test_single_large_batch_greater_than_target() {
838 let batch = uint32_batch(0..4096);
840 Test::new("coalesce_single_large_batch_greater_than_target")
841 .with_batch(batch)
842 .with_batch_size(1000)
843 .with_expected_output_sizes(vec![1000, 1000, 1000, 1000, 96])
844 .run();
845 }
846
847 #[test]
848 fn test_single_large_batch_smaller_than_target() {
849 let batch = uint32_batch(0..4096);
851 Test::new("coalesce_single_large_batch_smaller_than_target")
852 .with_batch(batch)
853 .with_batch_size(8192)
854 .with_expected_output_sizes(vec![4096])
855 .run();
856 }
857
858 #[test]
859 fn test_single_large_batch_equal_to_target() {
860 let batch = uint32_batch(0..4096);
862 Test::new("coalesce_single_large_batch_equal_to_target")
863 .with_batch(batch)
864 .with_batch_size(4096)
865 .with_expected_output_sizes(vec![4096])
866 .run();
867 }
868
869 #[test]
870 fn test_single_large_batch_equally_divisible_in_target() {
871 let batch = uint32_batch(0..4096);
873 Test::new("coalesce_single_large_batch_equally_divisible_in_target")
874 .with_batch(batch)
875 .with_batch_size(1024)
876 .with_expected_output_sizes(vec![1024, 1024, 1024, 1024])
877 .run();
878 }
879
880 #[test]
881 fn test_empty_schema() {
882 let schema = Schema::empty();
883 let batch = RecordBatch::new_empty(schema.into());
884 Test::new("coalesce_empty_schema")
885 .with_batch(batch)
886 .with_expected_output_sizes(vec![])
887 .run();
888 }
889
890 #[test]
892 #[cfg_attr(miri, ignore)] fn test_coalesce_filtered_001() {
894 let mut filter_builder = RandomFilterBuilder {
895 num_rows: 8000,
896 selectivity: 0.001,
897 seed: 0,
898 };
899
900 let mut test = Test::new("coalesce_filtered_001");
904 for _ in 0..10 {
905 test = test
906 .with_batch(multi_column_batch(0..8000))
907 .with_filter(filter_builder.next_filter())
908 }
909 test.with_batch_size(15)
910 .with_expected_output_sizes(vec![15, 15, 15, 13])
911 .run();
912 }
913
914 #[test]
916 #[cfg_attr(miri, ignore)] fn test_coalesce_filtered_01() {
918 let mut filter_builder = RandomFilterBuilder {
919 num_rows: 8000,
920 selectivity: 0.01,
921 seed: 0,
922 };
923
924 let mut test = Test::new("coalesce_filtered_01");
928 for _ in 0..10 {
929 test = test
930 .with_batch(multi_column_batch(0..8000))
931 .with_filter(filter_builder.next_filter())
932 }
933 test.with_batch_size(128)
934 .with_expected_output_sizes(vec![128, 128, 128, 128, 128, 128, 15])
935 .run();
936 }
937
938 #[test]
940 #[cfg_attr(miri, ignore)] fn test_coalesce_filtered_10() {
942 let mut filter_builder = RandomFilterBuilder {
943 num_rows: 8000,
944 selectivity: 0.1,
945 seed: 0,
946 };
947
948 let mut test = Test::new("coalesce_filtered_10");
952 for _ in 0..10 {
953 test = test
954 .with_batch(multi_column_batch(0..8000))
955 .with_filter(filter_builder.next_filter())
956 }
957 test.with_batch_size(1024)
958 .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 840])
959 .run();
960 }
961
962 #[test]
964 #[cfg_attr(miri, ignore)] fn test_coalesce_filtered_90() {
966 let mut filter_builder = RandomFilterBuilder {
967 num_rows: 800,
968 selectivity: 0.90,
969 seed: 0,
970 };
971
972 let mut test = Test::new("coalesce_filtered_90");
976 for _ in 0..10 {
977 test = test
978 .with_batch(multi_column_batch(0..800))
979 .with_filter(filter_builder.next_filter())
980 }
981 test.with_batch_size(1024)
982 .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 13])
983 .run();
984 }
985
986 #[test]
988 #[cfg_attr(miri, ignore)] fn test_coalesce_filtered_mixed() {
990 let mut filter_builder = RandomFilterBuilder {
991 num_rows: 800,
992 selectivity: 0.90,
993 seed: 0,
994 };
995
996 let mut test = Test::new("coalesce_filtered_mixed");
997 for _ in 0..3 {
998 let mut all_filter_builder = BooleanBufferBuilder::new(1000);
1001 all_filter_builder.append_n(500, true);
1002 all_filter_builder.append_n(1, false);
1003 all_filter_builder.append_n(499, false);
1004 let all_filter = all_filter_builder.build();
1005
1006 test = test
1007 .with_batch(multi_column_batch(0..1000))
1008 .with_filter(BooleanArray::from(all_filter))
1009 .with_batch(multi_column_batch(0..800))
1010 .with_filter(filter_builder.next_filter());
1011 filter_builder.selectivity *= 0.6;
1013 }
1014
1015 test.with_batch_size(250)
1018 .with_expected_output_sizes(vec![
1019 250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 179,
1020 ])
1021 .run();
1022 }
1023
1024 #[test]
1025 fn test_coalesce_non_null() {
1026 Test::new("coalesce_non_null")
1027 .with_batch(uint32_batch_non_null(0..3000))
1029 .with_batch(uint32_batch_non_null(0..1040))
1030 .with_batch_size(1024)
1031 .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1032 .run();
1033 }
1034 #[test]
1035 #[cfg_attr(miri, ignore)] fn test_utf8_split() {
1037 Test::new("coalesce_utf8")
1038 .with_batch(utf8_batch(0..3000))
1040 .with_batch(utf8_batch(0..1040))
1041 .with_batch_size(1024)
1042 .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1043 .run();
1044 }
1045
1046 #[test]
1047 fn test_string_view_no_views() {
1048 let output_batches = Test::new("coalesce_string_view_no_views")
1049 .with_batch(stringview_batch([Some("foo"), Some("bar")]))
1051 .with_batch(stringview_batch([Some("baz"), Some("qux")]))
1052 .with_expected_output_sizes(vec![4])
1053 .run();
1054
1055 expect_buffer_layout(
1056 col_as_string_view("c0", output_batches.first().unwrap()),
1057 vec![],
1058 );
1059 }
1060
1061 #[test]
1062 fn test_string_view_batch_small_no_compact() {
1063 let batch = stringview_batch_repeated(1000, [Some("a"), Some("b"), Some("c")]);
1065 let output_batches = Test::new("coalesce_string_view_batch_small_no_compact")
1066 .with_batch(batch.clone())
1067 .with_expected_output_sizes(vec![1000])
1068 .run();
1069
1070 let array = col_as_string_view("c0", &batch);
1071 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1072 assert_eq!(array.data_buffers().len(), 0);
1073 assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); expect_buffer_layout(gc_array, vec![]);
1076 }
1077
1078 #[test]
1079 #[cfg_attr(miri, ignore)] fn test_string_view_batch_large_no_compact() {
1081 let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1083 let output_batches = Test::new("coalesce_string_view_batch_large_no_compact")
1084 .with_batch(batch.clone())
1085 .with_batch_size(1000)
1086 .with_expected_output_sizes(vec![1000])
1087 .run();
1088
1089 let array = col_as_string_view("c0", &batch);
1090 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1091 assert_eq!(array.data_buffers().len(), 5);
1092 assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); expect_buffer_layout(
1095 gc_array,
1096 vec![
1097 ExpectedLayout {
1098 len: 8190,
1099 capacity: 8192,
1100 },
1101 ExpectedLayout {
1102 len: 8190,
1103 capacity: 8192,
1104 },
1105 ExpectedLayout {
1106 len: 8190,
1107 capacity: 8192,
1108 },
1109 ExpectedLayout {
1110 len: 8190,
1111 capacity: 8192,
1112 },
1113 ExpectedLayout {
1114 len: 2240,
1115 capacity: 8192,
1116 },
1117 ],
1118 );
1119 }
1120
1121 #[test]
1122 fn test_string_view_batch_small_with_buffers_no_compact() {
1123 let short_strings = std::iter::repeat(Some("SmallString"));
1125 let long_strings = std::iter::once(Some("This string is longer than 12 bytes"));
1126 let values = short_strings.take(20).chain(long_strings);
1128 let batch = stringview_batch_repeated(1000, values)
1129 .slice(5, 10);
1131 let output_batches = Test::new("coalesce_string_view_batch_small_with_buffers_no_compact")
1132 .with_batch(batch.clone())
1133 .with_batch_size(1000)
1134 .with_expected_output_sizes(vec![10])
1135 .run();
1136
1137 let array = col_as_string_view("c0", &batch);
1138 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1139 assert_eq!(array.data_buffers().len(), 1); assert_eq!(gc_array.data_buffers().len(), 0); }
1142
1143 #[test]
1144 fn test_string_view_batch_large_slice_compact() {
1145 let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")])
1147 .slice(11, 22);
1149
1150 let output_batches = Test::new("coalesce_string_view_batch_large_slice_compact")
1151 .with_batch(batch.clone())
1152 .with_batch_size(1000)
1153 .with_expected_output_sizes(vec![22])
1154 .run();
1155
1156 let array = col_as_string_view("c0", &batch);
1157 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1158 assert_eq!(array.data_buffers().len(), 5);
1159
1160 expect_buffer_layout(
1161 gc_array,
1162 vec![ExpectedLayout {
1163 len: 770,
1164 capacity: 8192,
1165 }],
1166 );
1167 }
1168
1169 #[test]
1170 #[cfg_attr(miri, ignore)] fn test_string_view_mixed() {
1172 let large_view_batch =
1173 stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1174 let small_view_batch = stringview_batch_repeated(1000, [Some("SmallString")]);
1175 let mixed_batch = stringview_batch_repeated(
1176 1000,
1177 [Some("This string is longer than 12 bytes"), Some("Small")],
1178 );
1179 let mixed_batch_nulls = stringview_batch_repeated(
1180 1000,
1181 [
1182 Some("This string is longer than 12 bytes"),
1183 Some("Small"),
1184 None,
1185 ],
1186 );
1187
1188 let output_batches = Test::new("coalesce_string_view_mixed")
1191 .with_batch(large_view_batch.clone())
1192 .with_batch(small_view_batch)
1193 .with_batch(large_view_batch.slice(10, 20))
1195 .with_batch(mixed_batch_nulls)
1196 .with_batch(large_view_batch.slice(10, 20))
1198 .with_batch(mixed_batch)
1199 .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1200 .run();
1201
1202 expect_buffer_layout(
1203 col_as_string_view("c0", output_batches.first().unwrap()),
1204 vec![
1205 ExpectedLayout {
1206 len: 8190,
1207 capacity: 8192,
1208 },
1209 ExpectedLayout {
1210 len: 8190,
1211 capacity: 8192,
1212 },
1213 ExpectedLayout {
1214 len: 8190,
1215 capacity: 8192,
1216 },
1217 ExpectedLayout {
1218 len: 8190,
1219 capacity: 8192,
1220 },
1221 ExpectedLayout {
1222 len: 2240,
1223 capacity: 8192,
1224 },
1225 ],
1226 );
1227 }
1228
1229 #[test]
1230 #[cfg_attr(miri, ignore)] fn test_string_view_many_small_compact() {
1232 let batch = stringview_batch_repeated(
1235 200,
1236 [Some("This string is 28 bytes long"), Some("small string")],
1237 );
1238 let output_batches = Test::new("coalesce_string_view_many_small_compact")
1239 .with_batch(batch.clone())
1242 .with_batch(batch.clone())
1243 .with_batch(batch.clone())
1244 .with_batch(batch.clone())
1245 .with_batch(batch.clone())
1246 .with_batch(batch.clone())
1247 .with_batch(batch.clone())
1248 .with_batch(batch.clone())
1249 .with_batch(batch.clone())
1250 .with_batch(batch.clone())
1251 .with_batch_size(8000)
1252 .with_expected_output_sizes(vec![2000]) .run();
1254
1255 expect_buffer_layout(
1257 col_as_string_view("c0", output_batches.first().unwrap()),
1258 vec![
1259 ExpectedLayout {
1260 len: 8176,
1261 capacity: 8192,
1262 },
1263 ExpectedLayout {
1264 len: 16380,
1265 capacity: 16384,
1266 },
1267 ExpectedLayout {
1268 len: 3444,
1269 capacity: 32768,
1270 },
1271 ],
1272 );
1273 }
1274
1275 #[test]
1276 #[cfg_attr(miri, ignore)] fn test_string_view_many_small_boundary() {
1278 let batch = stringview_batch_repeated(100, [Some("This string is a power of two=32")]);
1280 let output_batches = Test::new("coalesce_string_view_many_small_boundary")
1281 .with_batches(std::iter::repeat_n(batch, 20))
1282 .with_batch_size(900)
1283 .with_expected_output_sizes(vec![900, 900, 200])
1284 .run();
1285
1286 expect_buffer_layout(
1288 col_as_string_view("c0", output_batches.first().unwrap()),
1289 vec![
1290 ExpectedLayout {
1291 len: 8192,
1292 capacity: 8192,
1293 },
1294 ExpectedLayout {
1295 len: 16384,
1296 capacity: 16384,
1297 },
1298 ExpectedLayout {
1299 len: 4224,
1300 capacity: 32768,
1301 },
1302 ],
1303 );
1304 }
1305
1306 #[test]
1307 #[cfg_attr(miri, ignore)] fn test_string_view_large_small() {
1309 let mixed_batch = stringview_batch_repeated(
1311 200,
1312 [Some("This string is 28 bytes long"), Some("small string")],
1313 );
1314 let all_large = stringview_batch_repeated(
1316 50,
1317 [Some(
1318 "This buffer has only large strings in it so there are no buffer copies",
1319 )],
1320 );
1321
1322 let output_batches = Test::new("coalesce_string_view_large_small")
1323 .with_batch(mixed_batch.clone())
1326 .with_batch(mixed_batch.clone())
1327 .with_batch(all_large.clone())
1328 .with_batch(mixed_batch.clone())
1329 .with_batch(all_large.clone())
1330 .with_batch(mixed_batch.clone())
1331 .with_batch(mixed_batch.clone())
1332 .with_batch(all_large.clone())
1333 .with_batch(mixed_batch.clone())
1334 .with_batch(all_large.clone())
1335 .with_batch_size(8000)
1336 .with_expected_output_sizes(vec![1400])
1337 .run();
1338
1339 expect_buffer_layout(
1340 col_as_string_view("c0", output_batches.first().unwrap()),
1341 vec![
1342 ExpectedLayout {
1343 len: 8190,
1344 capacity: 8192,
1345 },
1346 ExpectedLayout {
1347 len: 16366,
1348 capacity: 16384,
1349 },
1350 ExpectedLayout {
1351 len: 6244,
1352 capacity: 32768,
1353 },
1354 ],
1355 );
1356 }
1357
1358 #[test]
1359 #[cfg_attr(miri, ignore)] fn test_binary_view() {
1361 let values: Vec<Option<&[u8]>> = vec![
1362 Some(b"foo"),
1363 None,
1364 Some(b"A longer string that is more than 12 bytes"),
1365 ];
1366
1367 let binary_view =
1368 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1369 let batch =
1370 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1371
1372 Test::new("coalesce_binary_view")
1373 .with_batch(batch.clone())
1374 .with_batch(batch.clone())
1375 .with_batch_size(512)
1376 .with_expected_output_sizes(vec![512, 512, 512, 464])
1377 .run();
1378 }
1379
1380 #[test]
1381 fn test_binary_view_filtered() {
1382 let values: Vec<Option<&[u8]>> = vec![
1383 Some(b"foo"),
1384 None,
1385 Some(b"A longer string that is more than 12 bytes"),
1386 ];
1387
1388 let binary_view =
1389 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1390 let batch =
1391 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1392 let filter = sparse_filter(1000);
1393
1394 Test::new("coalesce_binary_view_filtered")
1395 .with_batch(batch.clone())
1396 .with_filter(filter.clone())
1397 .with_batch(batch)
1398 .with_filter(filter)
1399 .with_batch_size(256)
1400 .with_expected_output_sizes(vec![250])
1401 .run();
1402 }
1403
1404 #[test]
1405 fn test_binary_view_filtered_inline() {
1406 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1407
1408 let binary_view =
1409 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1410 let batch =
1411 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1412 let filter = sparse_filter(1000);
1413
1414 Test::new("coalesce_binary_view_filtered_inline")
1415 .with_batch(batch.clone())
1416 .with_filter(filter.clone())
1417 .with_batch(batch)
1418 .with_filter(filter)
1419 .with_batch_size(300)
1420 .with_expected_output_sizes(vec![250])
1421 .run();
1422 }
1423
1424 #[test]
1425 fn test_string_view_filtered_inline() {
1426 let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1427
1428 let string_view =
1429 StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1430 let batch =
1431 RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1432 let filter = sparse_filter(1000);
1433
1434 Test::new("coalesce_string_view_filtered_inline")
1435 .with_batch(batch.clone())
1436 .with_filter(filter.clone())
1437 .with_batch(batch)
1438 .with_filter(filter)
1439 .with_batch_size(300)
1440 .with_expected_output_sizes(vec![250])
1441 .run();
1442 }
1443
1444 #[test]
1445 fn test_mixed_inline_binary_view_filtered() {
1446 let int_values =
1447 Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1448 let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1449 let binary_values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1450 let binary_view = BinaryViewArray::from_iter(
1451 std::iter::repeat(binary_values.iter()).flatten().take(1000),
1452 );
1453
1454 let batch = RecordBatch::try_from_iter(vec![
1455 ("i", Arc::new(int_values) as ArrayRef),
1456 ("f", Arc::new(float_values) as ArrayRef),
1457 ("b", Arc::new(binary_view) as ArrayRef),
1458 ])
1459 .unwrap();
1460
1461 let filter = sparse_filter(1000);
1462
1463 Test::new("coalesce_mixed_inline_binary_view_filtered")
1464 .with_batch(batch.clone())
1465 .with_filter(filter.clone())
1466 .with_batch(batch)
1467 .with_filter(filter)
1468 .with_batch_size(300)
1469 .with_expected_output_sizes(vec![250])
1470 .run();
1471 }
1472
1473 #[test]
1474 fn test_mixed_inline_string_view_filtered() {
1475 let int_values =
1476 Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1477 let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1478 let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1479 let string_view = StringViewArray::from_iter(
1480 std::iter::repeat(string_values.iter()).flatten().take(1000),
1481 );
1482
1483 let batch = RecordBatch::try_from_iter(vec![
1484 ("i", Arc::new(int_values) as ArrayRef),
1485 ("f", Arc::new(float_values) as ArrayRef),
1486 ("s", Arc::new(string_view) as ArrayRef),
1487 ])
1488 .unwrap();
1489
1490 let filter = sparse_filter(1000);
1491
1492 Test::new("coalesce_mixed_inline_string_view_filtered")
1493 .with_batch(batch.clone())
1494 .with_filter(filter.clone())
1495 .with_batch(batch)
1496 .with_filter(filter)
1497 .with_batch_size(300)
1498 .with_expected_output_sizes(vec![250])
1499 .run();
1500 }
1501
1502 #[test]
1503 fn test_inline_binary_view_sparse() {
1504 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1506 let binary_view =
1507 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1508 let batch =
1509 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1510 let filter = very_sparse_filter(1000);
1511
1512 Test::new("inline_binary_view_sparse")
1513 .with_batch(batch.clone())
1514 .with_filter(filter.clone())
1515 .with_batch(batch)
1516 .with_filter(filter)
1517 .with_batch_size(1024)
1518 .with_expected_output_sizes(vec![100])
1519 .run();
1520 }
1521
1522 #[test]
1523 fn test_inline_string_view_sparse() {
1524 let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1525 let string_view =
1526 StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1527 let batch =
1528 RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1529 let filter = very_sparse_filter(1000);
1530
1531 Test::new("inline_string_view_sparse")
1532 .with_batch(batch.clone())
1533 .with_filter(filter.clone())
1534 .with_batch(batch)
1535 .with_filter(filter)
1536 .with_batch_size(1024)
1537 .with_expected_output_sizes(vec![100])
1538 .run();
1539 }
1540
1541 #[test]
1542 fn test_inline_mixed_sparse() {
1543 let int_values =
1546 Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1547 let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1548 let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1549 let string_view = StringViewArray::from_iter(
1550 std::iter::repeat(string_values.iter()).flatten().take(1000),
1551 );
1552 let binary_values: Vec<Option<&[u8]>> = vec![Some(b"x"), None, Some(b"abcdef")];
1553 let binary_view = BinaryViewArray::from_iter(
1554 std::iter::repeat(binary_values.iter()).flatten().take(1000),
1555 );
1556
1557 let batch = RecordBatch::try_from_iter(vec![
1558 ("i", Arc::new(int_values) as ArrayRef),
1559 ("f", Arc::new(float_values) as ArrayRef),
1560 ("s", Arc::new(string_view) as ArrayRef),
1561 ("b", Arc::new(binary_view) as ArrayRef),
1562 ])
1563 .unwrap();
1564 let filter = very_sparse_filter(1000);
1565
1566 Test::new("inline_mixed_sparse")
1567 .with_batch(batch.clone())
1568 .with_filter(filter.clone())
1569 .with_batch(batch)
1570 .with_filter(filter)
1571 .with_batch_size(1024)
1572 .with_expected_output_sizes(vec![100])
1573 .run();
1574 }
1575
1576 #[test]
1577 fn test_inline_crosses_target_batch_size() {
1578 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1582 let make_batch = || {
1583 let binary_view =
1584 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1585 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap()
1586 };
1587 let filter = very_sparse_filter(1000);
1588
1589 Test::new("inline_crosses_target_batch_size")
1590 .with_batch(make_batch())
1591 .with_filter(filter.clone())
1592 .with_batch(make_batch())
1593 .with_filter(filter.clone())
1594 .with_batch(make_batch())
1595 .with_filter(filter)
1596 .with_batch_size(100)
1597 .with_expected_output_sizes(vec![100, 50])
1598 .run();
1599 }
1600
1601 #[test]
1602 fn test_inline_filter_rejects_filter_longer_than_batch() {
1603 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), Some(b"bar")];
1604 let binary_view = BinaryViewArray::from_iter(values);
1605 let batch =
1606 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1607 let filter = BooleanArray::from(vec![true, false, true]);
1608
1609 let mut coalescer = BatchCoalescer::new(batch.schema(), 100);
1610 let result = coalescer.push_batch_with_filter(batch, &filter);
1611 assert!(result.is_err());
1612 let err = result.unwrap_err().to_string();
1613 assert!(
1614 err.contains("Filter predicate of length 3 is larger than target array of length 2"),
1615 "unexpected error: {err}"
1616 );
1617 }
1618
1619 #[test]
1620 fn test_mixed_boolean_inline_string_view_filtered() {
1621 let bool_values = BooleanArray::from_iter((0..1000).map(|v| Some(v % 3 == 0)));
1622 let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1623 let string_view = StringViewArray::from_iter(
1624 std::iter::repeat(string_values.iter()).flatten().take(1000),
1625 );
1626
1627 let batch = RecordBatch::try_from_iter(vec![
1628 ("b", Arc::new(bool_values) as ArrayRef),
1629 ("s", Arc::new(string_view) as ArrayRef),
1630 ])
1631 .unwrap();
1632
1633 let filter = sparse_filter(1000);
1634
1635 Test::new("coalesce_mixed_boolean_inline_string_view_filtered")
1636 .with_batch(batch.clone())
1637 .with_filter(filter.clone())
1638 .with_batch(batch)
1639 .with_filter(filter)
1640 .with_batch_size(300)
1641 .with_expected_output_sizes(vec![250])
1642 .run();
1643 }
1644
1645 #[test]
1646 fn test_filter_fast_path_schema_capability() {
1647 let supported = Arc::new(Schema::new(vec![
1648 Field::new("primitive", DataType::UInt32, false),
1649 Field::new("utf8_view", DataType::Utf8View, true),
1650 Field::new("binary_view", DataType::BinaryView, true),
1651 ]));
1652 let coalescer = BatchCoalescer::new(supported, 100);
1653 assert!(!coalescer.has_non_specialized_filter_columns);
1654
1655 let utf8 = Arc::new(Schema::new(vec![Field::new("utf8", DataType::Utf8, true)]));
1656 let coalescer = BatchCoalescer::new(utf8, 100);
1657 assert!(coalescer.has_non_specialized_filter_columns);
1658
1659 let boolean = Arc::new(Schema::new(vec![Field::new(
1660 "boolean",
1661 DataType::Boolean,
1662 true,
1663 )]));
1664 let coalescer = BatchCoalescer::new(boolean, 100);
1665 assert!(coalescer.has_non_specialized_filter_columns);
1666 }
1667
1668 #[derive(Debug, Clone, PartialEq)]
1669 struct ExpectedLayout {
1670 len: usize,
1671 capacity: usize,
1672 }
1673
1674 fn expect_buffer_layout(array: &StringViewArray, expected: Vec<ExpectedLayout>) {
1676 let actual = array
1677 .data_buffers()
1678 .iter()
1679 .map(|b| ExpectedLayout {
1680 len: b.len(),
1681 capacity: b.capacity(),
1682 })
1683 .collect::<Vec<_>>();
1684
1685 assert_eq!(
1686 actual, expected,
1687 "Expected buffer layout {expected:#?} but got {actual:#?}"
1688 );
1689 }
1690
1691 #[derive(Debug, Clone)]
1698 struct Test {
1699 name: String,
1701 input_batches: Vec<RecordBatch>,
1703 filters: Vec<BooleanArray>,
1708 schema: Option<SchemaRef>,
1710 expected_output_sizes: Vec<usize>,
1712 target_batch_size: usize,
1714 }
1715
1716 impl Default for Test {
1717 fn default() -> Self {
1718 Self {
1719 name: String::new(),
1720 input_batches: vec![],
1721 filters: vec![],
1722 schema: None,
1723 expected_output_sizes: vec![],
1724 target_batch_size: 1024,
1725 }
1726 }
1727 }
1728
1729 impl Test {
1730 fn new(name: impl Into<String>) -> Self {
1731 Self {
1732 name: name.into(),
1733 ..Self::default()
1734 }
1735 }
1736
1737 fn with_description(mut self, description: &str) -> Self {
1739 self.name.push_str(": ");
1740 self.name.push_str(description);
1741 self
1742 }
1743
1744 fn with_batch_size(mut self, target_batch_size: usize) -> Self {
1746 self.target_batch_size = target_batch_size;
1747 self
1748 }
1749
1750 fn with_batch(mut self, batch: RecordBatch) -> Self {
1752 self.input_batches.push(batch);
1753 self
1754 }
1755
1756 fn with_filter(mut self, filter: BooleanArray) -> Self {
1758 self.filters.push(filter);
1759 self
1760 }
1761
1762 fn with_batches(mut self, batches: impl IntoIterator<Item = RecordBatch>) -> Self {
1764 self.input_batches = batches.into_iter().collect();
1765 self
1766 }
1767
1768 fn with_schema(mut self, schema: SchemaRef) -> Self {
1770 self.schema = Some(schema);
1771 self
1772 }
1773
1774 fn with_expected_output_sizes(mut self, sizes: impl IntoIterator<Item = usize>) -> Self {
1776 self.expected_output_sizes.extend(sizes);
1777 self
1778 }
1779
1780 fn run(self) -> Vec<RecordBatch> {
1784 let mut extra_tests = vec![];
1789 extra_tests.push(self.clone().make_half_non_nullable());
1790 extra_tests.push(self.clone().insert_empty_batches());
1791 let single_column_tests = self.make_single_column_tests();
1792 for test in single_column_tests {
1793 extra_tests.push(test.clone().make_half_non_nullable());
1794 extra_tests.push(test);
1795 }
1796
1797 let results = self.run_inner();
1800 for extra in extra_tests {
1802 extra.run_inner();
1803 }
1804
1805 results
1806 }
1807
1808 fn run_inner(self) -> Vec<RecordBatch> {
1810 let expected_output = self.expected_output();
1811 let schema = self.schema();
1812
1813 let Self {
1814 name,
1815 input_batches,
1816 filters,
1817 schema: _,
1818 target_batch_size,
1819 expected_output_sizes,
1820 } = self;
1821
1822 println!("Running test '{name}'");
1823
1824 let had_input = input_batches.iter().any(|b| b.num_rows() > 0);
1825
1826 let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size);
1827
1828 let mut filters = filters.into_iter();
1830 for batch in input_batches {
1831 if let Some(filter) = filters.next() {
1832 coalescer.push_batch_with_filter(batch, &filter).unwrap();
1833 } else {
1834 coalescer.push_batch(batch).unwrap();
1835 }
1836 }
1837 assert_eq!(schema, coalescer.schema());
1838
1839 if had_input {
1840 assert!(!coalescer.is_empty(), "Coalescer should not be empty");
1841 } else {
1842 assert!(coalescer.is_empty(), "Coalescer should be empty");
1843 }
1844
1845 coalescer.finish_buffered_batch().unwrap();
1846 if had_input {
1847 assert!(
1848 coalescer.has_completed_batch(),
1849 "Coalescer should have completed batches"
1850 );
1851 }
1852
1853 let mut output_batches = vec![];
1854 while let Some(batch) = coalescer.next_completed_batch() {
1855 output_batches.push(batch);
1856 }
1857
1858 let mut starting_idx = 0;
1860 let actual_output_sizes: Vec<usize> =
1861 output_batches.iter().map(|b| b.num_rows()).collect();
1862 assert_eq!(
1863 expected_output_sizes, actual_output_sizes,
1864 "Unexpected number of rows in output batches\n\
1865 Expected\n{expected_output_sizes:#?}\nActual:{actual_output_sizes:#?}"
1866 );
1867 let iter = expected_output_sizes
1868 .iter()
1869 .zip(output_batches.iter())
1870 .enumerate();
1871
1872 for (i, (expected_size, batch)) in iter {
1874 let expected_batch = expected_output.slice(starting_idx, *expected_size);
1877 let expected_batch = normalize_batch(expected_batch);
1878 let batch = normalize_batch(batch.clone());
1879 assert_eq!(
1880 expected_batch, batch,
1881 "Unexpected content in batch {i}:\
1882 \n\nExpected:\n{expected_batch:#?}\n\nActual:\n{batch:#?}"
1883 );
1884 starting_idx += *expected_size;
1885 }
1886 output_batches
1887 }
1888
1889 fn schema(&self) -> SchemaRef {
1892 self.schema
1893 .clone()
1894 .unwrap_or_else(|| Arc::clone(&self.input_batches[0].schema()))
1895 }
1896
1897 fn expected_output(&self) -> RecordBatch {
1899 let schema = self.schema();
1900 if self.filters.is_empty() {
1901 return concat_batches(&schema, &self.input_batches).unwrap();
1902 }
1903
1904 let mut filters = self.filters.iter();
1905 let filtered_batches = self
1906 .input_batches
1907 .iter()
1908 .map(|batch| {
1909 if let Some(filter) = filters.next() {
1910 filter_record_batch(batch, filter).unwrap()
1911 } else {
1912 batch.clone()
1913 }
1914 })
1915 .collect::<Vec<_>>();
1916 concat_batches(&schema, &filtered_batches).unwrap()
1917 }
1918
1919 fn make_half_non_nullable(mut self) -> Self {
1922 self.input_batches = self
1924 .input_batches
1925 .iter()
1926 .enumerate()
1927 .map(|(i, batch)| {
1928 if i % 2 == 1 {
1929 batch.clone()
1930 } else {
1931 Self::remove_nulls_from_batch(batch)
1932 }
1933 })
1934 .collect();
1935 self.with_description("non-nullable")
1936 }
1937
1938 fn insert_empty_batches(mut self) -> Self {
1940 let empty_batch = RecordBatch::new_empty(self.schema());
1941 self.input_batches = self
1942 .input_batches
1943 .into_iter()
1944 .flat_map(|batch| [empty_batch.clone(), batch])
1945 .collect();
1946 let empty_filters = BooleanArray::builder(0).finish();
1947 self.filters = self
1948 .filters
1949 .into_iter()
1950 .flat_map(|filter| [empty_filters.clone(), filter])
1951 .collect();
1952 self.with_description("empty batches inserted")
1953 }
1954
1955 fn remove_nulls_from_batch(batch: &RecordBatch) -> RecordBatch {
1957 let new_columns = batch
1958 .columns()
1959 .iter()
1960 .map(Self::remove_nulls_from_array)
1961 .collect::<Vec<_>>();
1962 let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
1963 RecordBatch::try_new_with_options(batch.schema(), new_columns, &options).unwrap()
1964 }
1965
1966 fn remove_nulls_from_array(array: &ArrayRef) -> ArrayRef {
1967 make_array(array.to_data().into_builder().nulls(None).build().unwrap())
1968 }
1969
1970 fn make_single_column_tests(&self) -> Vec<Self> {
1976 let original_schema = self.schema();
1977 let mut new_tests = vec![];
1978 for column in original_schema.fields() {
1979 let single_column_schema = Arc::new(Schema::new(vec![column.clone()]));
1980
1981 let single_column_batches = self.input_batches.iter().map(|batch| {
1982 let single_column = batch.column_by_name(column.name()).unwrap();
1983 RecordBatch::try_new(
1984 Arc::clone(&single_column_schema),
1985 vec![single_column.clone()],
1986 )
1987 .unwrap()
1988 });
1989
1990 let single_column_test = self
1991 .clone()
1992 .with_schema(Arc::clone(&single_column_schema))
1993 .with_batches(single_column_batches)
1994 .with_description("single column")
1995 .with_description(column.name());
1996
1997 new_tests.push(single_column_test);
1998 }
1999 new_tests
2000 }
2001 }
2002
2003 fn uint32_batch<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
2006 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, true)]));
2007
2008 let array = UInt32Array::from_iter(range.map(|i| if i % 3 == 0 { None } else { Some(i) }));
2009 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2010 }
2011
2012 fn uint32_batch_non_null<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
2014 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
2015
2016 let array = UInt32Array::from_iter_values(range);
2017 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2018 }
2019
2020 fn uint64_batch_non_null<T: std::iter::Iterator<Item = u64>>(range: T) -> RecordBatch {
2022 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt64, false)]));
2023
2024 let array = UInt64Array::from_iter_values(range);
2025 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2026 }
2027
2028 fn utf8_batch(range: Range<u32>) -> RecordBatch {
2031 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Utf8, true)]));
2032
2033 let array = StringArray::from_iter(range.map(|i| {
2034 if i % 3 == 0 {
2035 None
2036 } else {
2037 Some(format!("value{i}"))
2038 }
2039 }));
2040
2041 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2042 }
2043
2044 fn stringview_batch<'a>(values: impl IntoIterator<Item = Option<&'a str>>) -> RecordBatch {
2046 let schema = Arc::new(Schema::new(vec![Field::new(
2047 "c0",
2048 DataType::Utf8View,
2049 false,
2050 )]));
2051
2052 let array = StringViewArray::from_iter(values);
2053 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2054 }
2055
2056 fn stringview_batch_repeated<'a>(
2059 num_rows: usize,
2060 values: impl IntoIterator<Item = Option<&'a str>>,
2061 ) -> RecordBatch {
2062 let schema = Arc::new(Schema::new(vec![Field::new(
2063 "c0",
2064 DataType::Utf8View,
2065 true,
2066 )]));
2067
2068 let values: Vec<_> = values.into_iter().collect();
2070 let values_iter = std::iter::repeat(values.iter())
2071 .flatten()
2072 .cloned()
2073 .take(num_rows);
2074
2075 let mut builder = StringViewBuilder::with_capacity(100).with_fixed_block_size(8192);
2076 for val in values_iter {
2077 builder.append_option(val);
2078 }
2079
2080 let array = builder.finish();
2081 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2082 }
2083
2084 fn multi_column_batch(range: Range<i32>) -> RecordBatch {
2086 let int64_array = Int64Array::from_iter(
2087 range
2088 .clone()
2089 .map(|v| if v % 5 == 0 { None } else { Some(v as i64) }),
2090 );
2091 let string_view_array = StringViewArray::from_iter(range.clone().map(|v| {
2092 if v % 5 == 0 {
2093 None
2094 } else if v % 7 == 0 {
2095 Some(format!("This is a string longer than 12 bytes{v}"))
2096 } else {
2097 Some(format!("Short {v}"))
2098 }
2099 }));
2100 let string_array = StringArray::from_iter(range.clone().map(|v| {
2101 if v % 11 == 0 {
2102 None
2103 } else {
2104 Some(format!("Value {v}"))
2105 }
2106 }));
2107 let timestamp_array = TimestampNanosecondArray::from_iter(range.map(|v| {
2108 if v % 3 == 0 {
2109 None
2110 } else {
2111 Some(v as i64 * 1000) }
2113 }))
2114 .with_timezone("America/New_York");
2115
2116 RecordBatch::try_from_iter(vec![
2117 ("int64", Arc::new(int64_array) as ArrayRef),
2118 ("stringview", Arc::new(string_view_array) as ArrayRef),
2119 ("string", Arc::new(string_array) as ArrayRef),
2120 ("timestamp", Arc::new(timestamp_array) as ArrayRef),
2121 ])
2122 .unwrap()
2123 }
2124
2125 #[derive(Debug)]
2131 struct RandomFilterBuilder {
2132 num_rows: usize,
2134 selectivity: f64,
2137 seed: u64,
2140 }
2141 impl RandomFilterBuilder {
2142 fn next_filter(&mut self) -> BooleanArray {
2145 assert!(self.selectivity >= 0.0 && self.selectivity <= 1.0);
2146 let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed);
2147 self.seed += 1;
2148 BooleanArray::from_iter(
2149 (0..self.num_rows)
2150 .map(|_| rng.random_bool(self.selectivity))
2151 .map(Some),
2152 )
2153 }
2154 }
2155
2156 fn col_as_string_view<'b>(name: &str, batch: &'b RecordBatch) -> &'b StringViewArray {
2158 batch
2159 .column_by_name(name)
2160 .expect("column not found")
2161 .as_string_view_opt()
2162 .expect("column is not a string view")
2163 }
2164
2165 fn sparse_filter(len: usize) -> BooleanArray {
2167 BooleanArray::from_iter((0..len).map(|idx| Some(idx % 8 == 0)))
2168 }
2169
2170 fn very_sparse_filter(len: usize) -> BooleanArray {
2174 BooleanArray::from_iter((0..len).map(|idx| Some(idx % 20 == 0)))
2175 }
2176
2177 fn normalize_batch(batch: RecordBatch) -> RecordBatch {
2180 let (schema, mut columns, row_count) = batch.into_parts();
2182
2183 for column in columns.iter_mut() {
2184 if let Some(string_view) = column.as_string_view_opt() {
2185 let mut builder = StringViewBuilder::new();
2188 for s in string_view.iter() {
2189 builder.append_option(s);
2190 }
2191 *column = Arc::new(builder.finish());
2192 continue;
2193 }
2194
2195 if let Some(binary_view) = column.as_binary_view_opt() {
2196 *column = Arc::new(BinaryViewArray::from_iter(binary_view.iter()));
2197 }
2198 }
2199
2200 let options = RecordBatchOptions::new().with_row_count(Some(row_count));
2201 RecordBatch::try_new_with_options(schema, columns, &options).unwrap()
2202 }
2203
2204 fn create_test_batch(num_rows: usize) -> RecordBatch {
2206 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)]));
2207 let array = Int32Array::from_iter_values(0..num_rows as i32);
2208 RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
2209 }
2210 #[test]
2211 fn test_biggest_coalesce_batch_size_none_default() {
2212 let mut coalescer = BatchCoalescer::new(
2214 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2215 100,
2216 );
2217
2218 let large_batch = create_test_batch(1000);
2220 coalescer.push_batch(large_batch).unwrap();
2221
2222 let mut output_batches = vec![];
2224 while let Some(batch) = coalescer.next_completed_batch() {
2225 output_batches.push(batch);
2226 }
2227
2228 coalescer.finish_buffered_batch().unwrap();
2229 while let Some(batch) = coalescer.next_completed_batch() {
2230 output_batches.push(batch);
2231 }
2232
2233 assert_eq!(output_batches.len(), 10);
2235 for batch in output_batches {
2236 assert_eq!(batch.num_rows(), 100);
2237 }
2238 }
2239
2240 #[test]
2241 fn test_biggest_coalesce_batch_size_bypass_large_batch() {
2242 let mut coalescer = BatchCoalescer::new(
2244 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2245 100,
2246 );
2247 coalescer.set_biggest_coalesce_batch_size(Some(500));
2248
2249 let large_batch = create_test_batch(1000);
2251 coalescer.push_batch(large_batch.clone()).unwrap();
2252
2253 assert!(coalescer.has_completed_batch());
2255 let output_batch = coalescer.next_completed_batch().unwrap();
2256 assert_eq!(output_batch.num_rows(), 1000);
2257
2258 assert!(!coalescer.has_completed_batch());
2260 assert_eq!(coalescer.get_buffered_rows(), 0);
2261 }
2262
2263 #[test]
2264 fn test_biggest_coalesce_batch_size_coalesce_small_batch() {
2265 let mut coalescer = BatchCoalescer::new(
2267 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2268 100,
2269 )
2270 .with_biggest_coalesce_batch_size(Some(500));
2271
2272 let small_batch = create_test_batch(50);
2274 coalescer.push_batch(small_batch.clone()).unwrap();
2275
2276 assert!(!coalescer.has_completed_batch());
2278 assert_eq!(coalescer.get_buffered_rows(), 50);
2279
2280 coalescer.push_batch(small_batch).unwrap();
2282
2283 assert!(coalescer.has_completed_batch());
2285 let output_batch = coalescer.next_completed_batch().unwrap();
2286 let size = output_batch
2287 .column(0)
2288 .as_primitive::<Int32Type>()
2289 .get_buffer_memory_size();
2290 assert_eq!(size, 400); assert_eq!(output_batch.num_rows(), 100);
2292
2293 assert_eq!(coalescer.get_buffered_rows(), 0);
2294 }
2295
2296 #[test]
2297 fn test_biggest_coalesce_batch_size_equal_boundary() {
2298 let mut coalescer = BatchCoalescer::new(
2300 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2301 100,
2302 );
2303 coalescer.set_biggest_coalesce_batch_size(Some(500));
2304
2305 let boundary_batch = create_test_batch(500);
2307 coalescer.push_batch(boundary_batch).unwrap();
2308
2309 let mut output_count = 0;
2311 while coalescer.next_completed_batch().is_some() {
2312 output_count += 1;
2313 }
2314
2315 coalescer.finish_buffered_batch().unwrap();
2316 while coalescer.next_completed_batch().is_some() {
2317 output_count += 1;
2318 }
2319
2320 assert_eq!(output_count, 5);
2322 }
2323
2324 #[test]
2325 fn test_biggest_coalesce_batch_size_first_large_then_consecutive_bypass() {
2326 let mut coalescer = BatchCoalescer::new(
2329 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2330 100,
2331 );
2332 coalescer.set_biggest_coalesce_batch_size(Some(200));
2333
2334 let small_batch = create_test_batch(50);
2335
2336 coalescer.push_batch(small_batch).unwrap();
2338 assert_eq!(coalescer.get_buffered_rows(), 50);
2339 assert!(!coalescer.has_completed_batch());
2340
2341 let large_batch1 = create_test_batch(250);
2343 coalescer.push_batch(large_batch1).unwrap();
2344
2345 let mut completed_batches = vec![];
2347 while let Some(batch) = coalescer.next_completed_batch() {
2348 completed_batches.push(batch);
2349 }
2350 assert_eq!(completed_batches.len(), 3);
2351 assert_eq!(coalescer.get_buffered_rows(), 0);
2352
2353 let large_batch2 = create_test_batch(300);
2355 let large_batch3 = create_test_batch(400);
2356
2357 coalescer.push_batch(large_batch2).unwrap();
2359 assert!(coalescer.has_completed_batch());
2360 let output = coalescer.next_completed_batch().unwrap();
2361 assert_eq!(output.num_rows(), 300); assert_eq!(coalescer.get_buffered_rows(), 0);
2363
2364 coalescer.push_batch(large_batch3).unwrap();
2366 assert!(coalescer.has_completed_batch());
2367 let output = coalescer.next_completed_batch().unwrap();
2368 assert_eq!(output.num_rows(), 400); assert_eq!(coalescer.get_buffered_rows(), 0);
2370 }
2371
2372 #[test]
2373 fn test_biggest_coalesce_batch_size_empty_batch() {
2374 let mut coalescer = BatchCoalescer::new(
2376 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2377 100,
2378 );
2379 coalescer.set_biggest_coalesce_batch_size(Some(50));
2380
2381 let empty_batch = create_test_batch(0);
2382 coalescer.push_batch(empty_batch).unwrap();
2383
2384 assert!(!coalescer.has_completed_batch());
2386 assert_eq!(coalescer.get_buffered_rows(), 0);
2387 }
2388
2389 #[test]
2390 fn test_biggest_coalesce_batch_size_with_buffered_data_no_bypass() {
2391 let mut coalescer = BatchCoalescer::new(
2393 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2394 100,
2395 );
2396 coalescer.set_biggest_coalesce_batch_size(Some(200));
2397
2398 let small_batch = create_test_batch(30);
2400 coalescer.push_batch(small_batch.clone()).unwrap();
2401 coalescer.push_batch(small_batch).unwrap();
2402 assert_eq!(coalescer.get_buffered_rows(), 60);
2403
2404 let large_batch = create_test_batch(250);
2406 coalescer.push_batch(large_batch).unwrap();
2407
2408 let mut completed_batches = vec![];
2413 while let Some(batch) = coalescer.next_completed_batch() {
2414 completed_batches.push(batch);
2415 }
2416
2417 assert_eq!(completed_batches.len(), 3);
2418 for batch in &completed_batches {
2419 assert_eq!(batch.num_rows(), 100);
2420 }
2421 assert_eq!(coalescer.get_buffered_rows(), 10);
2422 }
2423
2424 #[test]
2425 fn test_biggest_coalesce_batch_size_zero_limit() {
2426 let mut coalescer = BatchCoalescer::new(
2428 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2429 100,
2430 );
2431 coalescer.set_biggest_coalesce_batch_size(Some(0));
2432
2433 let tiny_batch = create_test_batch(1);
2435 coalescer.push_batch(tiny_batch).unwrap();
2436
2437 assert!(coalescer.has_completed_batch());
2438 let output = coalescer.next_completed_batch().unwrap();
2439 assert_eq!(output.num_rows(), 1);
2440 }
2441
2442 #[test]
2443 fn test_biggest_coalesce_batch_size_bypass_only_when_no_buffer() {
2444 let mut coalescer = BatchCoalescer::new(
2446 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2447 100,
2448 );
2449 coalescer.set_biggest_coalesce_batch_size(Some(200));
2450
2451 let large_batch = create_test_batch(300);
2453 coalescer.push_batch(large_batch.clone()).unwrap();
2454
2455 assert!(coalescer.has_completed_batch());
2456 let output = coalescer.next_completed_batch().unwrap();
2457 assert_eq!(output.num_rows(), 300); assert_eq!(coalescer.get_buffered_rows(), 0);
2459
2460 let small_batch = create_test_batch(50);
2462 coalescer.push_batch(small_batch).unwrap();
2463 assert_eq!(coalescer.get_buffered_rows(), 50);
2464
2465 coalescer.push_batch(large_batch).unwrap();
2467
2468 let mut completed_batches = vec![];
2471 while let Some(batch) = coalescer.next_completed_batch() {
2472 completed_batches.push(batch);
2473 }
2474
2475 assert_eq!(completed_batches.len(), 3);
2476 for batch in &completed_batches {
2477 assert_eq!(batch.num_rows(), 100);
2478 }
2479 assert_eq!(coalescer.get_buffered_rows(), 50);
2480 }
2481
2482 #[test]
2483 fn test_biggest_coalesce_batch_size_consecutive_large_batches_scenario() {
2484 let mut coalescer = BatchCoalescer::new(
2486 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2487 1000,
2488 );
2489 coalescer.set_biggest_coalesce_batch_size(Some(500));
2490
2491 coalescer.push_batch(create_test_batch(20)).unwrap();
2493 coalescer.push_batch(create_test_batch(20)).unwrap();
2494 coalescer.push_batch(create_test_batch(30)).unwrap();
2495
2496 assert_eq!(coalescer.get_buffered_rows(), 70);
2497 assert!(!coalescer.has_completed_batch());
2498
2499 coalescer.push_batch(create_test_batch(700)).unwrap();
2501
2502 assert_eq!(coalescer.get_buffered_rows(), 770);
2504 assert!(!coalescer.has_completed_batch());
2505
2506 coalescer.push_batch(create_test_batch(600)).unwrap();
2508
2509 let mut outputs = vec![];
2511 while let Some(batch) = coalescer.next_completed_batch() {
2512 outputs.push(batch);
2513 }
2514 assert_eq!(outputs.len(), 2); assert_eq!(outputs[0].num_rows(), 770);
2516 assert_eq!(outputs[1].num_rows(), 600);
2517 assert_eq!(coalescer.get_buffered_rows(), 0);
2518
2519 let remaining_batches = [700, 900, 700, 600];
2521 for &size in &remaining_batches {
2522 coalescer.push_batch(create_test_batch(size)).unwrap();
2523
2524 assert!(coalescer.has_completed_batch());
2525 let output = coalescer.next_completed_batch().unwrap();
2526 assert_eq!(output.num_rows(), size);
2527 assert_eq!(coalescer.get_buffered_rows(), 0);
2528 }
2529 }
2530
2531 #[test]
2532 fn test_biggest_coalesce_batch_size_truly_consecutive_large_bypass() {
2533 let mut coalescer = BatchCoalescer::new(
2536 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2537 100,
2538 );
2539 coalescer.set_biggest_coalesce_batch_size(Some(200));
2540
2541 let large_batches = vec![
2543 create_test_batch(300),
2544 create_test_batch(400),
2545 create_test_batch(350),
2546 create_test_batch(500),
2547 ];
2548
2549 let mut all_outputs = vec![];
2550
2551 for (i, large_batch) in large_batches.into_iter().enumerate() {
2552 let expected_size = large_batch.num_rows();
2553
2554 assert_eq!(
2556 coalescer.get_buffered_rows(),
2557 0,
2558 "Buffer should be empty before batch {}",
2559 i
2560 );
2561
2562 coalescer.push_batch(large_batch).unwrap();
2563
2564 assert!(
2566 coalescer.has_completed_batch(),
2567 "Should have completed batch after pushing batch {}",
2568 i
2569 );
2570
2571 let output = coalescer.next_completed_batch().unwrap();
2572 assert_eq!(
2573 output.num_rows(),
2574 expected_size,
2575 "Batch {} should have bypassed with original size",
2576 i
2577 );
2578
2579 assert!(
2581 !coalescer.has_completed_batch(),
2582 "Should have no more completed batches after batch {}",
2583 i
2584 );
2585 assert_eq!(
2586 coalescer.get_buffered_rows(),
2587 0,
2588 "Buffer should be empty after batch {}",
2589 i
2590 );
2591
2592 all_outputs.push(output);
2593 }
2594
2595 assert_eq!(all_outputs.len(), 4);
2597 assert_eq!(all_outputs[0].num_rows(), 300);
2598 assert_eq!(all_outputs[1].num_rows(), 400);
2599 assert_eq!(all_outputs[2].num_rows(), 350);
2600 assert_eq!(all_outputs[3].num_rows(), 500);
2601 }
2602
2603 #[test]
2604 fn test_biggest_coalesce_batch_size_reset_consecutive_on_small_batch() {
2605 let mut coalescer = BatchCoalescer::new(
2607 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2608 100,
2609 );
2610 coalescer.set_biggest_coalesce_batch_size(Some(200));
2611
2612 coalescer.push_batch(create_test_batch(300)).unwrap();
2614 let output = coalescer.next_completed_batch().unwrap();
2615 assert_eq!(output.num_rows(), 300);
2616
2617 coalescer.push_batch(create_test_batch(400)).unwrap();
2619 let output = coalescer.next_completed_batch().unwrap();
2620 assert_eq!(output.num_rows(), 400);
2621
2622 coalescer.push_batch(create_test_batch(50)).unwrap();
2624 assert_eq!(coalescer.get_buffered_rows(), 50);
2625
2626 coalescer.push_batch(create_test_batch(350)).unwrap();
2628
2629 let mut outputs = vec![];
2631 while let Some(batch) = coalescer.next_completed_batch() {
2632 outputs.push(batch);
2633 }
2634 assert_eq!(outputs.len(), 4);
2635 for batch in outputs {
2636 assert_eq!(batch.num_rows(), 100);
2637 }
2638 assert_eq!(coalescer.get_buffered_rows(), 0);
2639 }
2640
2641 #[test]
2642 fn test_coalasce_push_batch_with_indices() {
2643 const MID_POINT: u32 = 2333;
2644 const TOTAL_ROWS: u32 = 23333;
2645 let batch1 = uint32_batch_non_null(0..MID_POINT);
2646 let batch2 = uint32_batch_non_null((MID_POINT..TOTAL_ROWS).rev());
2647
2648 let mut coalescer = BatchCoalescer::new(
2649 Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])),
2650 TOTAL_ROWS as usize,
2651 );
2652 coalescer.push_batch(batch1).unwrap();
2653
2654 let rev_indices = (0..((TOTAL_ROWS - MID_POINT) as u64)).rev();
2655 let reversed_indices_batch = uint64_batch_non_null(rev_indices);
2656
2657 let reverse_indices = UInt64Array::from(reversed_indices_batch.column(0).to_data());
2658 coalescer
2659 .push_batch_with_indices(batch2, &reverse_indices)
2660 .unwrap();
2661
2662 coalescer.finish_buffered_batch().unwrap();
2663 let actual = coalescer.next_completed_batch().unwrap();
2664
2665 let expected = uint32_batch_non_null(0..TOTAL_ROWS);
2666
2667 assert_eq!(expected, actual);
2668 }
2669
2670 #[test]
2671 fn test_push_batch_schema_mismatch_fewer_columns() {
2672 let empty_schema = Arc::new(Schema::empty());
2674 let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2675 let batch = uint32_batch(0..5);
2676 let result = coalescer.push_batch(batch);
2677 assert!(result.is_err());
2678 let err = result.unwrap_err().to_string();
2679 assert!(
2680 err.contains("Batch has 1 columns but BatchCoalescer expects 0"),
2681 "unexpected error: {err}"
2682 );
2683 }
2684
2685 #[test]
2686 fn test_push_batch_schema_mismatch_more_columns() {
2687 let schema = Arc::new(Schema::new(vec![
2689 Field::new("c0", DataType::UInt32, false),
2690 Field::new("c1", DataType::UInt32, false),
2691 ]));
2692 let mut coalescer = BatchCoalescer::new(schema, 100);
2693 let batch = uint32_batch(0..5);
2694 let result = coalescer.push_batch(batch);
2695 assert!(result.is_err());
2696 let err = result.unwrap_err().to_string();
2697 assert!(
2698 err.contains("Batch has 1 columns but BatchCoalescer expects 2"),
2699 "unexpected error: {err}"
2700 );
2701 }
2702
2703 #[test]
2704 fn test_push_batch_schema_mismatch_two_vs_zero() {
2705 let empty_schema = Arc::new(Schema::empty());
2707 let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2708 let schema = Arc::new(Schema::new(vec![
2709 Field::new("c0", DataType::UInt32, false),
2710 Field::new("c1", DataType::UInt32, false),
2711 ]));
2712 let batch = RecordBatch::try_new(
2713 schema,
2714 vec![
2715 Arc::new(UInt32Array::from(vec![1, 2, 3])),
2716 Arc::new(UInt32Array::from(vec![4, 5, 6])),
2717 ],
2718 )
2719 .unwrap();
2720 let result = coalescer.push_batch(batch);
2721 assert!(result.is_err());
2722 let err = result.unwrap_err().to_string();
2723 assert!(
2724 err.contains("Batch has 2 columns but BatchCoalescer expects 0"),
2725 "unexpected error: {err}"
2726 );
2727 }
2728
2729 #[test]
2730 fn test_size_grows_with_buffering_and_shrinks_when_draining() {
2731 let batch = uint32_batch(0..8);
2732 let mut coalescer = BatchCoalescer::new(batch.schema(), 21);
2733 let baseline = coalescer.size();
2734 assert!(baseline > 0, "size includes container capacities");
2735
2736 coalescer.push_batch(batch.clone()).unwrap();
2738 assert!(coalescer.next_completed_batch().is_none());
2739 let buffered = coalescer.size();
2740 assert!(
2741 buffered > baseline,
2742 "buffering rows should grow size ({buffered} > {baseline})"
2743 );
2744
2745 for _ in 0..10 {
2747 coalescer.push_batch(batch.clone()).unwrap();
2748 }
2749 let peak = coalescer.size();
2750 assert!(peak > buffered);
2751
2752 let mut prev = peak;
2754 let mut drained_any = false;
2755 while coalescer.next_completed_batch().is_some() {
2756 drained_any = true;
2757 let now = coalescer.size();
2758 assert!(now <= prev, "size grew while draining: {now} > {prev}");
2759 prev = now;
2760 }
2761 assert!(drained_any);
2762 assert!(
2763 prev < peak,
2764 "draining completed batches should release memory"
2765 );
2766 }
2767
2768 #[test]
2769 fn test_size_string_view_buffers_released_after_drain() {
2770 let batch = stringview_batch_repeated(
2773 1000,
2774 [Some("this string is definitely longer than 12 bytes")],
2775 );
2776 let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2777 let baseline = coalescer.size();
2778
2779 for _ in 0..20 {
2780 coalescer.push_batch(batch.clone()).unwrap();
2781 }
2782 let peak = coalescer.size();
2783 assert!(
2784 peak > baseline,
2785 "buffered string view data should grow size ({peak} > {baseline})"
2786 );
2787
2788 coalescer.finish_buffered_batch().unwrap();
2789 while coalescer.next_completed_batch().is_some() {}
2790
2791 let drained = coalescer.size();
2793 assert!(
2794 drained < peak,
2795 "draining should release buffered data ({drained} < {peak})"
2796 );
2797 }
2798
2799 #[test]
2803 fn test_size_accounting_conserved_across_cycles() {
2804 let batch = uint32_batch(0..8);
2808 let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2809
2810 let run_cycle = |coalescer: &mut BatchCoalescer| {
2811 for _ in 0..20 {
2812 coalescer.push_batch(batch.clone()).unwrap();
2813 }
2814 coalescer.finish_buffered_batch().unwrap();
2815 let peak = coalescer.size();
2816 while coalescer.next_completed_batch().is_some() {}
2817 (peak, coalescer.size())
2818 };
2819
2820 let (peak1, drained1) = run_cycle(&mut coalescer);
2821 let (peak2, drained2) = run_cycle(&mut coalescer);
2822
2823 assert_eq!(
2824 peak1, peak2,
2825 "identical work must report identical peak size"
2826 );
2827 assert_eq!(
2828 drained1, drained2,
2829 "fully-drained size must be stable across cycles (no accounting leak)"
2830 );
2831 assert!(
2832 drained1 < peak1,
2833 "draining must release the accounted memory"
2834 );
2835 }
2836}