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 if batch_size > limit {
454 if self.buffered_rows == 0 {
457 self.completed.push_back(batch);
458 return Ok(());
459 }
460
461 if self.buffered_rows > limit {
466 self.finish_buffered_batch()?;
467 self.completed.push_back(batch);
468 return Ok(());
469 }
470
471 }
476 }
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> {
732 self.copy_rows_by_selection(filter.selection())
733 }
734
735 fn copy_rows_by_filter_from(
737 &mut self,
738 source: ArrayRef,
739 filter: &FilterPredicate,
740 ) -> Result<(), ArrowError> {
741 self.set_source(Some(source));
742 let result = self.copy_rows_by_filter(filter);
743 self.set_source(None);
744 result
745 }
746
747 fn copy_rows_by_selection(&mut self, selection: FilterSelection<'_>) -> Result<(), ArrowError> {
749 match selection {
750 FilterSelection::None => Ok(()),
751 FilterSelection::All { len } => self.copy_rows(0, len),
752 FilterSelection::Slices(slices) => {
753 slices.try_for_each(|(start, end)| self.copy_rows(start, end - start))
754 }
755 FilterSelection::Indices(indices) => indices.try_for_each(|idx| self.copy_rows(idx, 1)),
756 }
757 }
758
759 fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
761
762 fn size(&self) -> usize;
764}
765
766#[cfg(test)]
767mod tests {
768 use super::*;
769 use crate::concat::concat_batches;
770 use crate::filter::filter_record_batch;
771 use arrow_array::builder::StringViewBuilder;
772 use arrow_array::cast::AsArray;
773 use arrow_array::types::Int32Type;
774 use arrow_array::{
775 BinaryViewArray, Int32Array, Int64Array, RecordBatchOptions, StringArray, StringViewArray,
776 TimestampNanosecondArray, UInt32Array, UInt64Array, make_array,
777 };
778 use arrow_buffer::BooleanBufferBuilder;
779 use arrow_schema::{DataType, Field, Schema};
780 use rand::{Rng, SeedableRng};
781 use std::ops::Range;
782
783 #[test]
784 fn test_coalesce() {
785 let batch = uint32_batch(0..8);
786 Test::new("coalesce")
787 .with_batches(std::iter::repeat_n(batch, 10))
788 .with_batch_size(21)
790 .with_expected_output_sizes(vec![21, 21, 21, 17])
791 .run();
792 }
793
794 #[test]
795 fn test_coalesce_one_by_one() {
796 let batch = uint32_batch(0..1); Test::new("coalesce_one_by_one")
798 .with_batches(std::iter::repeat_n(batch, 97))
799 .with_batch_size(20)
801 .with_expected_output_sizes(vec![20, 20, 20, 20, 17])
802 .run();
803 }
804
805 #[test]
806 fn test_coalesce_empty() {
807 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
808
809 Test::new("coalesce_empty")
810 .with_batches(vec![])
811 .with_schema(schema)
812 .with_batch_size(21)
813 .with_expected_output_sizes(vec![])
814 .run();
815 }
816
817 #[test]
818 fn test_sparse_filter_copy_threshold() {
819 assert!(should_use_sparse_filter_copy(8192, 8));
820 assert!(should_use_sparse_filter_copy(8192, 81));
821 assert!(!should_use_sparse_filter_copy(8192, 819));
822 assert!(!should_use_sparse_filter_copy(8192, 6553));
823 }
824
825 #[test]
826 fn test_single_large_batch_greater_than_target() {
827 let batch = uint32_batch(0..4096);
829 Test::new("coalesce_single_large_batch_greater_than_target")
830 .with_batch(batch)
831 .with_batch_size(1000)
832 .with_expected_output_sizes(vec![1000, 1000, 1000, 1000, 96])
833 .run();
834 }
835
836 #[test]
837 fn test_single_large_batch_smaller_than_target() {
838 let batch = uint32_batch(0..4096);
840 Test::new("coalesce_single_large_batch_smaller_than_target")
841 .with_batch(batch)
842 .with_batch_size(8192)
843 .with_expected_output_sizes(vec![4096])
844 .run();
845 }
846
847 #[test]
848 fn test_single_large_batch_equal_to_target() {
849 let batch = uint32_batch(0..4096);
851 Test::new("coalesce_single_large_batch_equal_to_target")
852 .with_batch(batch)
853 .with_batch_size(4096)
854 .with_expected_output_sizes(vec![4096])
855 .run();
856 }
857
858 #[test]
859 fn test_single_large_batch_equally_divisible_in_target() {
860 let batch = uint32_batch(0..4096);
862 Test::new("coalesce_single_large_batch_equally_divisible_in_target")
863 .with_batch(batch)
864 .with_batch_size(1024)
865 .with_expected_output_sizes(vec![1024, 1024, 1024, 1024])
866 .run();
867 }
868
869 #[test]
870 fn test_empty_schema() {
871 let schema = Schema::empty();
872 let batch = RecordBatch::new_empty(schema.into());
873 Test::new("coalesce_empty_schema")
874 .with_batch(batch)
875 .with_expected_output_sizes(vec![])
876 .run();
877 }
878
879 #[test]
881 fn test_coalesce_filtered_001() {
882 let mut filter_builder = RandomFilterBuilder {
883 num_rows: 8000,
884 selectivity: 0.001,
885 seed: 0,
886 };
887
888 let mut test = Test::new("coalesce_filtered_001");
892 for _ in 0..10 {
893 test = test
894 .with_batch(multi_column_batch(0..8000))
895 .with_filter(filter_builder.next_filter())
896 }
897 test.with_batch_size(15)
898 .with_expected_output_sizes(vec![15, 15, 15, 13])
899 .run();
900 }
901
902 #[test]
904 fn test_coalesce_filtered_01() {
905 let mut filter_builder = RandomFilterBuilder {
906 num_rows: 8000,
907 selectivity: 0.01,
908 seed: 0,
909 };
910
911 let mut test = Test::new("coalesce_filtered_01");
915 for _ in 0..10 {
916 test = test
917 .with_batch(multi_column_batch(0..8000))
918 .with_filter(filter_builder.next_filter())
919 }
920 test.with_batch_size(128)
921 .with_expected_output_sizes(vec![128, 128, 128, 128, 128, 128, 15])
922 .run();
923 }
924
925 #[test]
927 fn test_coalesce_filtered_10() {
928 let mut filter_builder = RandomFilterBuilder {
929 num_rows: 8000,
930 selectivity: 0.1,
931 seed: 0,
932 };
933
934 let mut test = Test::new("coalesce_filtered_10");
938 for _ in 0..10 {
939 test = test
940 .with_batch(multi_column_batch(0..8000))
941 .with_filter(filter_builder.next_filter())
942 }
943 test.with_batch_size(1024)
944 .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 840])
945 .run();
946 }
947
948 #[test]
950 fn test_coalesce_filtered_90() {
951 let mut filter_builder = RandomFilterBuilder {
952 num_rows: 800,
953 selectivity: 0.90,
954 seed: 0,
955 };
956
957 let mut test = Test::new("coalesce_filtered_90");
961 for _ in 0..10 {
962 test = test
963 .with_batch(multi_column_batch(0..800))
964 .with_filter(filter_builder.next_filter())
965 }
966 test.with_batch_size(1024)
967 .with_expected_output_sizes(vec![1024, 1024, 1024, 1024, 1024, 1024, 1024, 13])
968 .run();
969 }
970
971 #[test]
973 fn test_coalesce_filtered_mixed() {
974 let mut filter_builder = RandomFilterBuilder {
975 num_rows: 800,
976 selectivity: 0.90,
977 seed: 0,
978 };
979
980 let mut test = Test::new("coalesce_filtered_mixed");
981 for _ in 0..3 {
982 let mut all_filter_builder = BooleanBufferBuilder::new(1000);
985 all_filter_builder.append_n(500, true);
986 all_filter_builder.append_n(1, false);
987 all_filter_builder.append_n(499, false);
988 let all_filter = all_filter_builder.build();
989
990 test = test
991 .with_batch(multi_column_batch(0..1000))
992 .with_filter(BooleanArray::from(all_filter))
993 .with_batch(multi_column_batch(0..800))
994 .with_filter(filter_builder.next_filter());
995 filter_builder.selectivity *= 0.6;
997 }
998
999 test.with_batch_size(250)
1002 .with_expected_output_sizes(vec![
1003 250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 250, 179,
1004 ])
1005 .run();
1006 }
1007
1008 #[test]
1009 fn test_coalesce_non_null() {
1010 Test::new("coalesce_non_null")
1011 .with_batch(uint32_batch_non_null(0..3000))
1013 .with_batch(uint32_batch_non_null(0..1040))
1014 .with_batch_size(1024)
1015 .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1016 .run();
1017 }
1018 #[test]
1019 fn test_utf8_split() {
1020 Test::new("coalesce_utf8")
1021 .with_batch(utf8_batch(0..3000))
1023 .with_batch(utf8_batch(0..1040))
1024 .with_batch_size(1024)
1025 .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1026 .run();
1027 }
1028
1029 #[test]
1030 fn test_string_view_no_views() {
1031 let output_batches = Test::new("coalesce_string_view_no_views")
1032 .with_batch(stringview_batch([Some("foo"), Some("bar")]))
1034 .with_batch(stringview_batch([Some("baz"), Some("qux")]))
1035 .with_expected_output_sizes(vec![4])
1036 .run();
1037
1038 expect_buffer_layout(
1039 col_as_string_view("c0", output_batches.first().unwrap()),
1040 vec![],
1041 );
1042 }
1043
1044 #[test]
1045 fn test_string_view_batch_small_no_compact() {
1046 let batch = stringview_batch_repeated(1000, [Some("a"), Some("b"), Some("c")]);
1048 let output_batches = Test::new("coalesce_string_view_batch_small_no_compact")
1049 .with_batch(batch.clone())
1050 .with_expected_output_sizes(vec![1000])
1051 .run();
1052
1053 let array = col_as_string_view("c0", &batch);
1054 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1055 assert_eq!(array.data_buffers().len(), 0);
1056 assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); expect_buffer_layout(gc_array, vec![]);
1059 }
1060
1061 #[test]
1062 fn test_string_view_batch_large_no_compact() {
1063 let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1065 let output_batches = Test::new("coalesce_string_view_batch_large_no_compact")
1066 .with_batch(batch.clone())
1067 .with_batch_size(1000)
1068 .with_expected_output_sizes(vec![1000])
1069 .run();
1070
1071 let array = col_as_string_view("c0", &batch);
1072 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1073 assert_eq!(array.data_buffers().len(), 5);
1074 assert_eq!(array.data_buffers().len(), gc_array.data_buffers().len()); expect_buffer_layout(
1077 gc_array,
1078 vec![
1079 ExpectedLayout {
1080 len: 8190,
1081 capacity: 8192,
1082 },
1083 ExpectedLayout {
1084 len: 8190,
1085 capacity: 8192,
1086 },
1087 ExpectedLayout {
1088 len: 8190,
1089 capacity: 8192,
1090 },
1091 ExpectedLayout {
1092 len: 8190,
1093 capacity: 8192,
1094 },
1095 ExpectedLayout {
1096 len: 2240,
1097 capacity: 8192,
1098 },
1099 ],
1100 );
1101 }
1102
1103 #[test]
1104 fn test_string_view_batch_small_with_buffers_no_compact() {
1105 let short_strings = std::iter::repeat(Some("SmallString"));
1107 let long_strings = std::iter::once(Some("This string is longer than 12 bytes"));
1108 let values = short_strings.take(20).chain(long_strings);
1110 let batch = stringview_batch_repeated(1000, values)
1111 .slice(5, 10);
1113 let output_batches = Test::new("coalesce_string_view_batch_small_with_buffers_no_compact")
1114 .with_batch(batch.clone())
1115 .with_batch_size(1000)
1116 .with_expected_output_sizes(vec![10])
1117 .run();
1118
1119 let array = col_as_string_view("c0", &batch);
1120 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1121 assert_eq!(array.data_buffers().len(), 1); assert_eq!(gc_array.data_buffers().len(), 0); }
1124
1125 #[test]
1126 fn test_string_view_batch_large_slice_compact() {
1127 let batch = stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")])
1129 .slice(11, 22);
1131
1132 let output_batches = Test::new("coalesce_string_view_batch_large_slice_compact")
1133 .with_batch(batch.clone())
1134 .with_batch_size(1000)
1135 .with_expected_output_sizes(vec![22])
1136 .run();
1137
1138 let array = col_as_string_view("c0", &batch);
1139 let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
1140 assert_eq!(array.data_buffers().len(), 5);
1141
1142 expect_buffer_layout(
1143 gc_array,
1144 vec![ExpectedLayout {
1145 len: 770,
1146 capacity: 8192,
1147 }],
1148 );
1149 }
1150
1151 #[test]
1152 fn test_string_view_mixed() {
1153 let large_view_batch =
1154 stringview_batch_repeated(1000, [Some("This string is longer than 12 bytes")]);
1155 let small_view_batch = stringview_batch_repeated(1000, [Some("SmallString")]);
1156 let mixed_batch = stringview_batch_repeated(
1157 1000,
1158 [Some("This string is longer than 12 bytes"), Some("Small")],
1159 );
1160 let mixed_batch_nulls = stringview_batch_repeated(
1161 1000,
1162 [
1163 Some("This string is longer than 12 bytes"),
1164 Some("Small"),
1165 None,
1166 ],
1167 );
1168
1169 let output_batches = Test::new("coalesce_string_view_mixed")
1172 .with_batch(large_view_batch.clone())
1173 .with_batch(small_view_batch)
1174 .with_batch(large_view_batch.slice(10, 20))
1176 .with_batch(mixed_batch_nulls)
1177 .with_batch(large_view_batch.slice(10, 20))
1179 .with_batch(mixed_batch)
1180 .with_expected_output_sizes(vec![1024, 1024, 1024, 968])
1181 .run();
1182
1183 expect_buffer_layout(
1184 col_as_string_view("c0", output_batches.first().unwrap()),
1185 vec![
1186 ExpectedLayout {
1187 len: 8190,
1188 capacity: 8192,
1189 },
1190 ExpectedLayout {
1191 len: 8190,
1192 capacity: 8192,
1193 },
1194 ExpectedLayout {
1195 len: 8190,
1196 capacity: 8192,
1197 },
1198 ExpectedLayout {
1199 len: 8190,
1200 capacity: 8192,
1201 },
1202 ExpectedLayout {
1203 len: 2240,
1204 capacity: 8192,
1205 },
1206 ],
1207 );
1208 }
1209
1210 #[test]
1211 fn test_string_view_many_small_compact() {
1212 let batch = stringview_batch_repeated(
1215 200,
1216 [Some("This string is 28 bytes long"), Some("small string")],
1217 );
1218 let output_batches = Test::new("coalesce_string_view_many_small_compact")
1219 .with_batch(batch.clone())
1222 .with_batch(batch.clone())
1223 .with_batch(batch.clone())
1224 .with_batch(batch.clone())
1225 .with_batch(batch.clone())
1226 .with_batch(batch.clone())
1227 .with_batch(batch.clone())
1228 .with_batch(batch.clone())
1229 .with_batch(batch.clone())
1230 .with_batch(batch.clone())
1231 .with_batch_size(8000)
1232 .with_expected_output_sizes(vec![2000]) .run();
1234
1235 expect_buffer_layout(
1237 col_as_string_view("c0", output_batches.first().unwrap()),
1238 vec![
1239 ExpectedLayout {
1240 len: 8176,
1241 capacity: 8192,
1242 },
1243 ExpectedLayout {
1244 len: 16380,
1245 capacity: 16384,
1246 },
1247 ExpectedLayout {
1248 len: 3444,
1249 capacity: 32768,
1250 },
1251 ],
1252 );
1253 }
1254
1255 #[test]
1256 fn test_string_view_many_small_boundary() {
1257 let batch = stringview_batch_repeated(100, [Some("This string is a power of two=32")]);
1259 let output_batches = Test::new("coalesce_string_view_many_small_boundary")
1260 .with_batches(std::iter::repeat_n(batch, 20))
1261 .with_batch_size(900)
1262 .with_expected_output_sizes(vec![900, 900, 200])
1263 .run();
1264
1265 expect_buffer_layout(
1267 col_as_string_view("c0", output_batches.first().unwrap()),
1268 vec![
1269 ExpectedLayout {
1270 len: 8192,
1271 capacity: 8192,
1272 },
1273 ExpectedLayout {
1274 len: 16384,
1275 capacity: 16384,
1276 },
1277 ExpectedLayout {
1278 len: 4224,
1279 capacity: 32768,
1280 },
1281 ],
1282 );
1283 }
1284
1285 #[test]
1286 fn test_string_view_large_small() {
1287 let mixed_batch = stringview_batch_repeated(
1289 200,
1290 [Some("This string is 28 bytes long"), Some("small string")],
1291 );
1292 let all_large = stringview_batch_repeated(
1294 50,
1295 [Some(
1296 "This buffer has only large strings in it so there are no buffer copies",
1297 )],
1298 );
1299
1300 let output_batches = Test::new("coalesce_string_view_large_small")
1301 .with_batch(mixed_batch.clone())
1304 .with_batch(mixed_batch.clone())
1305 .with_batch(all_large.clone())
1306 .with_batch(mixed_batch.clone())
1307 .with_batch(all_large.clone())
1308 .with_batch(mixed_batch.clone())
1309 .with_batch(mixed_batch.clone())
1310 .with_batch(all_large.clone())
1311 .with_batch(mixed_batch.clone())
1312 .with_batch(all_large.clone())
1313 .with_batch_size(8000)
1314 .with_expected_output_sizes(vec![1400])
1315 .run();
1316
1317 expect_buffer_layout(
1318 col_as_string_view("c0", output_batches.first().unwrap()),
1319 vec![
1320 ExpectedLayout {
1321 len: 8190,
1322 capacity: 8192,
1323 },
1324 ExpectedLayout {
1325 len: 16366,
1326 capacity: 16384,
1327 },
1328 ExpectedLayout {
1329 len: 6244,
1330 capacity: 32768,
1331 },
1332 ],
1333 );
1334 }
1335
1336 #[test]
1337 fn test_binary_view() {
1338 let values: Vec<Option<&[u8]>> = vec![
1339 Some(b"foo"),
1340 None,
1341 Some(b"A longer string that is more than 12 bytes"),
1342 ];
1343
1344 let binary_view =
1345 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1346 let batch =
1347 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1348
1349 Test::new("coalesce_binary_view")
1350 .with_batch(batch.clone())
1351 .with_batch(batch.clone())
1352 .with_batch_size(512)
1353 .with_expected_output_sizes(vec![512, 512, 512, 464])
1354 .run();
1355 }
1356
1357 #[test]
1358 fn test_binary_view_filtered() {
1359 let values: Vec<Option<&[u8]>> = vec![
1360 Some(b"foo"),
1361 None,
1362 Some(b"A longer string that is more than 12 bytes"),
1363 ];
1364
1365 let binary_view =
1366 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1367 let batch =
1368 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1369 let filter = sparse_filter(1000);
1370
1371 Test::new("coalesce_binary_view_filtered")
1372 .with_batch(batch.clone())
1373 .with_filter(filter.clone())
1374 .with_batch(batch)
1375 .with_filter(filter)
1376 .with_batch_size(256)
1377 .with_expected_output_sizes(vec![250])
1378 .run();
1379 }
1380
1381 #[test]
1382 fn test_binary_view_filtered_inline() {
1383 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1384
1385 let binary_view =
1386 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1387 let batch =
1388 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1389 let filter = sparse_filter(1000);
1390
1391 Test::new("coalesce_binary_view_filtered_inline")
1392 .with_batch(batch.clone())
1393 .with_filter(filter.clone())
1394 .with_batch(batch)
1395 .with_filter(filter)
1396 .with_batch_size(300)
1397 .with_expected_output_sizes(vec![250])
1398 .run();
1399 }
1400
1401 #[test]
1402 fn test_string_view_filtered_inline() {
1403 let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1404
1405 let string_view =
1406 StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1407 let batch =
1408 RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1409 let filter = sparse_filter(1000);
1410
1411 Test::new("coalesce_string_view_filtered_inline")
1412 .with_batch(batch.clone())
1413 .with_filter(filter.clone())
1414 .with_batch(batch)
1415 .with_filter(filter)
1416 .with_batch_size(300)
1417 .with_expected_output_sizes(vec![250])
1418 .run();
1419 }
1420
1421 #[test]
1422 fn test_mixed_inline_binary_view_filtered() {
1423 let int_values =
1424 Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1425 let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1426 let binary_values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1427 let binary_view = BinaryViewArray::from_iter(
1428 std::iter::repeat(binary_values.iter()).flatten().take(1000),
1429 );
1430
1431 let batch = RecordBatch::try_from_iter(vec![
1432 ("i", Arc::new(int_values) as ArrayRef),
1433 ("f", Arc::new(float_values) as ArrayRef),
1434 ("b", Arc::new(binary_view) as ArrayRef),
1435 ])
1436 .unwrap();
1437
1438 let filter = sparse_filter(1000);
1439
1440 Test::new("coalesce_mixed_inline_binary_view_filtered")
1441 .with_batch(batch.clone())
1442 .with_filter(filter.clone())
1443 .with_batch(batch)
1444 .with_filter(filter)
1445 .with_batch_size(300)
1446 .with_expected_output_sizes(vec![250])
1447 .run();
1448 }
1449
1450 #[test]
1451 fn test_mixed_inline_string_view_filtered() {
1452 let int_values =
1453 Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1454 let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1455 let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1456 let string_view = StringViewArray::from_iter(
1457 std::iter::repeat(string_values.iter()).flatten().take(1000),
1458 );
1459
1460 let batch = RecordBatch::try_from_iter(vec![
1461 ("i", Arc::new(int_values) as ArrayRef),
1462 ("f", Arc::new(float_values) as ArrayRef),
1463 ("s", Arc::new(string_view) as ArrayRef),
1464 ])
1465 .unwrap();
1466
1467 let filter = sparse_filter(1000);
1468
1469 Test::new("coalesce_mixed_inline_string_view_filtered")
1470 .with_batch(batch.clone())
1471 .with_filter(filter.clone())
1472 .with_batch(batch)
1473 .with_filter(filter)
1474 .with_batch_size(300)
1475 .with_expected_output_sizes(vec![250])
1476 .run();
1477 }
1478
1479 #[test]
1480 fn test_inline_binary_view_sparse() {
1481 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1483 let binary_view =
1484 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1485 let batch =
1486 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1487 let filter = very_sparse_filter(1000);
1488
1489 Test::new("inline_binary_view_sparse")
1490 .with_batch(batch.clone())
1491 .with_filter(filter.clone())
1492 .with_batch(batch)
1493 .with_filter(filter)
1494 .with_batch_size(1024)
1495 .with_expected_output_sizes(vec![100])
1496 .run();
1497 }
1498
1499 #[test]
1500 fn test_inline_string_view_sparse() {
1501 let values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1502 let string_view =
1503 StringViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1504 let batch =
1505 RecordBatch::try_from_iter(vec![("c0", Arc::new(string_view) as ArrayRef)]).unwrap();
1506 let filter = very_sparse_filter(1000);
1507
1508 Test::new("inline_string_view_sparse")
1509 .with_batch(batch.clone())
1510 .with_filter(filter.clone())
1511 .with_batch(batch)
1512 .with_filter(filter)
1513 .with_batch_size(1024)
1514 .with_expected_output_sizes(vec![100])
1515 .run();
1516 }
1517
1518 #[test]
1519 fn test_inline_mixed_sparse() {
1520 let int_values =
1523 Int32Array::from_iter((0..1000).map(|v| if v % 5 == 0 { None } else { Some(v) }));
1524 let float_values = arrow_array::Float64Array::from_iter((0..1000).map(|v| Some(v as f64)));
1525 let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1526 let string_view = StringViewArray::from_iter(
1527 std::iter::repeat(string_values.iter()).flatten().take(1000),
1528 );
1529 let binary_values: Vec<Option<&[u8]>> = vec![Some(b"x"), None, Some(b"abcdef")];
1530 let binary_view = BinaryViewArray::from_iter(
1531 std::iter::repeat(binary_values.iter()).flatten().take(1000),
1532 );
1533
1534 let batch = RecordBatch::try_from_iter(vec![
1535 ("i", Arc::new(int_values) as ArrayRef),
1536 ("f", Arc::new(float_values) as ArrayRef),
1537 ("s", Arc::new(string_view) as ArrayRef),
1538 ("b", Arc::new(binary_view) as ArrayRef),
1539 ])
1540 .unwrap();
1541 let filter = very_sparse_filter(1000);
1542
1543 Test::new("inline_mixed_sparse")
1544 .with_batch(batch.clone())
1545 .with_filter(filter.clone())
1546 .with_batch(batch)
1547 .with_filter(filter)
1548 .with_batch_size(1024)
1549 .with_expected_output_sizes(vec![100])
1550 .run();
1551 }
1552
1553 #[test]
1554 fn test_inline_crosses_target_batch_size() {
1555 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), None, Some(b"barbaz")];
1559 let make_batch = || {
1560 let binary_view =
1561 BinaryViewArray::from_iter(std::iter::repeat(values.iter()).flatten().take(1000));
1562 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap()
1563 };
1564 let filter = very_sparse_filter(1000);
1565
1566 Test::new("inline_crosses_target_batch_size")
1567 .with_batch(make_batch())
1568 .with_filter(filter.clone())
1569 .with_batch(make_batch())
1570 .with_filter(filter.clone())
1571 .with_batch(make_batch())
1572 .with_filter(filter)
1573 .with_batch_size(100)
1574 .with_expected_output_sizes(vec![100, 50])
1575 .run();
1576 }
1577
1578 #[test]
1579 fn test_inline_filter_rejects_filter_longer_than_batch() {
1580 let values: Vec<Option<&[u8]>> = vec![Some(b"foo"), Some(b"bar")];
1581 let binary_view = BinaryViewArray::from_iter(values);
1582 let batch =
1583 RecordBatch::try_from_iter(vec![("c0", Arc::new(binary_view) as ArrayRef)]).unwrap();
1584 let filter = BooleanArray::from(vec![true, false, true]);
1585
1586 let mut coalescer = BatchCoalescer::new(batch.schema(), 100);
1587 let result = coalescer.push_batch_with_filter(batch, &filter);
1588 assert!(result.is_err());
1589 let err = result.unwrap_err().to_string();
1590 assert!(
1591 err.contains("Filter predicate of length 3 is larger than target array of length 2"),
1592 "unexpected error: {err}"
1593 );
1594 }
1595
1596 #[test]
1597 fn test_mixed_boolean_inline_string_view_filtered() {
1598 let bool_values = BooleanArray::from_iter((0..1000).map(|v| Some(v % 3 == 0)));
1599 let string_values: Vec<Option<&str>> = vec![Some("foo"), None, Some("barbaz")];
1600 let string_view = StringViewArray::from_iter(
1601 std::iter::repeat(string_values.iter()).flatten().take(1000),
1602 );
1603
1604 let batch = RecordBatch::try_from_iter(vec![
1605 ("b", Arc::new(bool_values) as ArrayRef),
1606 ("s", Arc::new(string_view) as ArrayRef),
1607 ])
1608 .unwrap();
1609
1610 let filter = sparse_filter(1000);
1611
1612 Test::new("coalesce_mixed_boolean_inline_string_view_filtered")
1613 .with_batch(batch.clone())
1614 .with_filter(filter.clone())
1615 .with_batch(batch)
1616 .with_filter(filter)
1617 .with_batch_size(300)
1618 .with_expected_output_sizes(vec![250])
1619 .run();
1620 }
1621
1622 #[test]
1623 fn test_filter_fast_path_schema_capability() {
1624 let supported = Arc::new(Schema::new(vec![
1625 Field::new("primitive", DataType::UInt32, false),
1626 Field::new("utf8_view", DataType::Utf8View, true),
1627 Field::new("binary_view", DataType::BinaryView, true),
1628 ]));
1629 let coalescer = BatchCoalescer::new(supported, 100);
1630 assert!(!coalescer.has_non_specialized_filter_columns);
1631
1632 let utf8 = Arc::new(Schema::new(vec![Field::new("utf8", DataType::Utf8, true)]));
1633 let coalescer = BatchCoalescer::new(utf8, 100);
1634 assert!(coalescer.has_non_specialized_filter_columns);
1635
1636 let boolean = Arc::new(Schema::new(vec![Field::new(
1637 "boolean",
1638 DataType::Boolean,
1639 true,
1640 )]));
1641 let coalescer = BatchCoalescer::new(boolean, 100);
1642 assert!(coalescer.has_non_specialized_filter_columns);
1643 }
1644
1645 #[derive(Debug, Clone, PartialEq)]
1646 struct ExpectedLayout {
1647 len: usize,
1648 capacity: usize,
1649 }
1650
1651 fn expect_buffer_layout(array: &StringViewArray, expected: Vec<ExpectedLayout>) {
1653 let actual = array
1654 .data_buffers()
1655 .iter()
1656 .map(|b| ExpectedLayout {
1657 len: b.len(),
1658 capacity: b.capacity(),
1659 })
1660 .collect::<Vec<_>>();
1661
1662 assert_eq!(
1663 actual, expected,
1664 "Expected buffer layout {expected:#?} but got {actual:#?}"
1665 );
1666 }
1667
1668 #[derive(Debug, Clone)]
1675 struct Test {
1676 name: String,
1678 input_batches: Vec<RecordBatch>,
1680 filters: Vec<BooleanArray>,
1685 schema: Option<SchemaRef>,
1687 expected_output_sizes: Vec<usize>,
1689 target_batch_size: usize,
1691 }
1692
1693 impl Default for Test {
1694 fn default() -> Self {
1695 Self {
1696 name: "".to_string(),
1697 input_batches: vec![],
1698 filters: vec![],
1699 schema: None,
1700 expected_output_sizes: vec![],
1701 target_batch_size: 1024,
1702 }
1703 }
1704 }
1705
1706 impl Test {
1707 fn new(name: impl Into<String>) -> Self {
1708 Self {
1709 name: name.into(),
1710 ..Self::default()
1711 }
1712 }
1713
1714 fn with_description(mut self, description: &str) -> Self {
1716 self.name.push_str(": ");
1717 self.name.push_str(description);
1718 self
1719 }
1720
1721 fn with_batch_size(mut self, target_batch_size: usize) -> Self {
1723 self.target_batch_size = target_batch_size;
1724 self
1725 }
1726
1727 fn with_batch(mut self, batch: RecordBatch) -> Self {
1729 self.input_batches.push(batch);
1730 self
1731 }
1732
1733 fn with_filter(mut self, filter: BooleanArray) -> Self {
1735 self.filters.push(filter);
1736 self
1737 }
1738
1739 fn with_batches(mut self, batches: impl IntoIterator<Item = RecordBatch>) -> Self {
1741 self.input_batches = batches.into_iter().collect();
1742 self
1743 }
1744
1745 fn with_schema(mut self, schema: SchemaRef) -> Self {
1747 self.schema = Some(schema);
1748 self
1749 }
1750
1751 fn with_expected_output_sizes(mut self, sizes: impl IntoIterator<Item = usize>) -> Self {
1753 self.expected_output_sizes.extend(sizes);
1754 self
1755 }
1756
1757 fn run(self) -> Vec<RecordBatch> {
1761 let mut extra_tests = vec![];
1766 extra_tests.push(self.clone().make_half_non_nullable());
1767 extra_tests.push(self.clone().insert_empty_batches());
1768 let single_column_tests = self.make_single_column_tests();
1769 for test in single_column_tests {
1770 extra_tests.push(test.clone().make_half_non_nullable());
1771 extra_tests.push(test);
1772 }
1773
1774 let results = self.run_inner();
1777 for extra in extra_tests {
1779 extra.run_inner();
1780 }
1781
1782 results
1783 }
1784
1785 fn run_inner(self) -> Vec<RecordBatch> {
1787 let expected_output = self.expected_output();
1788 let schema = self.schema();
1789
1790 let Self {
1791 name,
1792 input_batches,
1793 filters,
1794 schema: _,
1795 target_batch_size,
1796 expected_output_sizes,
1797 } = self;
1798
1799 println!("Running test '{name}'");
1800
1801 let had_input = input_batches.iter().any(|b| b.num_rows() > 0);
1802
1803 let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size);
1804
1805 let mut filters = filters.into_iter();
1807 for batch in input_batches {
1808 if let Some(filter) = filters.next() {
1809 coalescer.push_batch_with_filter(batch, &filter).unwrap();
1810 } else {
1811 coalescer.push_batch(batch).unwrap();
1812 }
1813 }
1814 assert_eq!(schema, coalescer.schema());
1815
1816 if had_input {
1817 assert!(!coalescer.is_empty(), "Coalescer should not be empty");
1818 } else {
1819 assert!(coalescer.is_empty(), "Coalescer should be empty");
1820 }
1821
1822 coalescer.finish_buffered_batch().unwrap();
1823 if had_input {
1824 assert!(
1825 coalescer.has_completed_batch(),
1826 "Coalescer should have completed batches"
1827 );
1828 }
1829
1830 let mut output_batches = vec![];
1831 while let Some(batch) = coalescer.next_completed_batch() {
1832 output_batches.push(batch);
1833 }
1834
1835 let mut starting_idx = 0;
1837 let actual_output_sizes: Vec<usize> =
1838 output_batches.iter().map(|b| b.num_rows()).collect();
1839 assert_eq!(
1840 expected_output_sizes, actual_output_sizes,
1841 "Unexpected number of rows in output batches\n\
1842 Expected\n{expected_output_sizes:#?}\nActual:{actual_output_sizes:#?}"
1843 );
1844 let iter = expected_output_sizes
1845 .iter()
1846 .zip(output_batches.iter())
1847 .enumerate();
1848
1849 for (i, (expected_size, batch)) in iter {
1851 let expected_batch = expected_output.slice(starting_idx, *expected_size);
1854 let expected_batch = normalize_batch(expected_batch);
1855 let batch = normalize_batch(batch.clone());
1856 assert_eq!(
1857 expected_batch, batch,
1858 "Unexpected content in batch {i}:\
1859 \n\nExpected:\n{expected_batch:#?}\n\nActual:\n{batch:#?}"
1860 );
1861 starting_idx += *expected_size;
1862 }
1863 output_batches
1864 }
1865
1866 fn schema(&self) -> SchemaRef {
1869 self.schema
1870 .clone()
1871 .unwrap_or_else(|| Arc::clone(&self.input_batches[0].schema()))
1872 }
1873
1874 fn expected_output(&self) -> RecordBatch {
1876 let schema = self.schema();
1877 if self.filters.is_empty() {
1878 return concat_batches(&schema, &self.input_batches).unwrap();
1879 }
1880
1881 let mut filters = self.filters.iter();
1882 let filtered_batches = self
1883 .input_batches
1884 .iter()
1885 .map(|batch| {
1886 if let Some(filter) = filters.next() {
1887 filter_record_batch(batch, filter).unwrap()
1888 } else {
1889 batch.clone()
1890 }
1891 })
1892 .collect::<Vec<_>>();
1893 concat_batches(&schema, &filtered_batches).unwrap()
1894 }
1895
1896 fn make_half_non_nullable(mut self) -> Self {
1899 self.input_batches = self
1901 .input_batches
1902 .iter()
1903 .enumerate()
1904 .map(|(i, batch)| {
1905 if i % 2 == 1 {
1906 batch.clone()
1907 } else {
1908 Self::remove_nulls_from_batch(batch)
1909 }
1910 })
1911 .collect();
1912 self.with_description("non-nullable")
1913 }
1914
1915 fn insert_empty_batches(mut self) -> Self {
1917 let empty_batch = RecordBatch::new_empty(self.schema());
1918 self.input_batches = self
1919 .input_batches
1920 .into_iter()
1921 .flat_map(|batch| [empty_batch.clone(), batch])
1922 .collect();
1923 let empty_filters = BooleanArray::builder(0).finish();
1924 self.filters = self
1925 .filters
1926 .into_iter()
1927 .flat_map(|filter| [empty_filters.clone(), filter])
1928 .collect();
1929 self.with_description("empty batches inserted")
1930 }
1931
1932 fn remove_nulls_from_batch(batch: &RecordBatch) -> RecordBatch {
1934 let new_columns = batch
1935 .columns()
1936 .iter()
1937 .map(Self::remove_nulls_from_array)
1938 .collect::<Vec<_>>();
1939 let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
1940 RecordBatch::try_new_with_options(batch.schema(), new_columns, &options).unwrap()
1941 }
1942
1943 fn remove_nulls_from_array(array: &ArrayRef) -> ArrayRef {
1944 make_array(array.to_data().into_builder().nulls(None).build().unwrap())
1945 }
1946
1947 fn make_single_column_tests(&self) -> Vec<Self> {
1953 let original_schema = self.schema();
1954 let mut new_tests = vec![];
1955 for column in original_schema.fields() {
1956 let single_column_schema = Arc::new(Schema::new(vec![column.clone()]));
1957
1958 let single_column_batches = self.input_batches.iter().map(|batch| {
1959 let single_column = batch.column_by_name(column.name()).unwrap();
1960 RecordBatch::try_new(
1961 Arc::clone(&single_column_schema),
1962 vec![single_column.clone()],
1963 )
1964 .unwrap()
1965 });
1966
1967 let single_column_test = self
1968 .clone()
1969 .with_schema(Arc::clone(&single_column_schema))
1970 .with_batches(single_column_batches)
1971 .with_description("single column")
1972 .with_description(column.name());
1973
1974 new_tests.push(single_column_test);
1975 }
1976 new_tests
1977 }
1978 }
1979
1980 fn uint32_batch<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
1983 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, true)]));
1984
1985 let array = UInt32Array::from_iter(range.map(|i| if i % 3 == 0 { None } else { Some(i) }));
1986 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
1987 }
1988
1989 fn uint32_batch_non_null<T: std::iter::Iterator<Item = u32>>(range: T) -> RecordBatch {
1991 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]));
1992
1993 let array = UInt32Array::from_iter_values(range);
1994 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
1995 }
1996
1997 fn uint64_batch_non_null<T: std::iter::Iterator<Item = u64>>(range: T) -> RecordBatch {
1999 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt64, false)]));
2000
2001 let array = UInt64Array::from_iter_values(range);
2002 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2003 }
2004
2005 fn utf8_batch(range: Range<u32>) -> RecordBatch {
2008 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Utf8, true)]));
2009
2010 let array = StringArray::from_iter(range.map(|i| {
2011 if i % 3 == 0 {
2012 None
2013 } else {
2014 Some(format!("value{i}"))
2015 }
2016 }));
2017
2018 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2019 }
2020
2021 fn stringview_batch<'a>(values: impl IntoIterator<Item = Option<&'a str>>) -> RecordBatch {
2023 let schema = Arc::new(Schema::new(vec![Field::new(
2024 "c0",
2025 DataType::Utf8View,
2026 false,
2027 )]));
2028
2029 let array = StringViewArray::from_iter(values);
2030 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2031 }
2032
2033 fn stringview_batch_repeated<'a>(
2036 num_rows: usize,
2037 values: impl IntoIterator<Item = Option<&'a str>>,
2038 ) -> RecordBatch {
2039 let schema = Arc::new(Schema::new(vec![Field::new(
2040 "c0",
2041 DataType::Utf8View,
2042 true,
2043 )]));
2044
2045 let values: Vec<_> = values.into_iter().collect();
2047 let values_iter = std::iter::repeat(values.iter())
2048 .flatten()
2049 .cloned()
2050 .take(num_rows);
2051
2052 let mut builder = StringViewBuilder::with_capacity(100).with_fixed_block_size(8192);
2053 for val in values_iter {
2054 builder.append_option(val);
2055 }
2056
2057 let array = builder.finish();
2058 RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(array)]).unwrap()
2059 }
2060
2061 fn multi_column_batch(range: Range<i32>) -> RecordBatch {
2063 let int64_array = Int64Array::from_iter(
2064 range
2065 .clone()
2066 .map(|v| if v % 5 == 0 { None } else { Some(v as i64) }),
2067 );
2068 let string_view_array = StringViewArray::from_iter(range.clone().map(|v| {
2069 if v % 5 == 0 {
2070 None
2071 } else if v % 7 == 0 {
2072 Some(format!("This is a string longer than 12 bytes{v}"))
2073 } else {
2074 Some(format!("Short {v}"))
2075 }
2076 }));
2077 let string_array = StringArray::from_iter(range.clone().map(|v| {
2078 if v % 11 == 0 {
2079 None
2080 } else {
2081 Some(format!("Value {v}"))
2082 }
2083 }));
2084 let timestamp_array = TimestampNanosecondArray::from_iter(range.map(|v| {
2085 if v % 3 == 0 {
2086 None
2087 } else {
2088 Some(v as i64 * 1000) }
2090 }))
2091 .with_timezone("America/New_York");
2092
2093 RecordBatch::try_from_iter(vec![
2094 ("int64", Arc::new(int64_array) as ArrayRef),
2095 ("stringview", Arc::new(string_view_array) as ArrayRef),
2096 ("string", Arc::new(string_array) as ArrayRef),
2097 ("timestamp", Arc::new(timestamp_array) as ArrayRef),
2098 ])
2099 .unwrap()
2100 }
2101
2102 #[derive(Debug)]
2108 struct RandomFilterBuilder {
2109 num_rows: usize,
2111 selectivity: f64,
2114 seed: u64,
2117 }
2118 impl RandomFilterBuilder {
2119 fn next_filter(&mut self) -> BooleanArray {
2122 assert!(self.selectivity >= 0.0 && self.selectivity <= 1.0);
2123 let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed);
2124 self.seed += 1;
2125 BooleanArray::from_iter(
2126 (0..self.num_rows)
2127 .map(|_| rng.random_bool(self.selectivity))
2128 .map(Some),
2129 )
2130 }
2131 }
2132
2133 fn col_as_string_view<'b>(name: &str, batch: &'b RecordBatch) -> &'b StringViewArray {
2135 batch
2136 .column_by_name(name)
2137 .expect("column not found")
2138 .as_string_view_opt()
2139 .expect("column is not a string view")
2140 }
2141
2142 fn sparse_filter(len: usize) -> BooleanArray {
2144 BooleanArray::from_iter((0..len).map(|idx| Some(idx % 8 == 0)))
2145 }
2146
2147 fn very_sparse_filter(len: usize) -> BooleanArray {
2151 BooleanArray::from_iter((0..len).map(|idx| Some(idx % 20 == 0)))
2152 }
2153
2154 fn normalize_batch(batch: RecordBatch) -> RecordBatch {
2157 let (schema, mut columns, row_count) = batch.into_parts();
2159
2160 for column in columns.iter_mut() {
2161 if let Some(string_view) = column.as_string_view_opt() {
2162 let mut builder = StringViewBuilder::new();
2165 for s in string_view.iter() {
2166 builder.append_option(s);
2167 }
2168 *column = Arc::new(builder.finish());
2169 continue;
2170 }
2171
2172 if let Some(binary_view) = column.as_binary_view_opt() {
2173 *column = Arc::new(BinaryViewArray::from_iter(binary_view.iter()));
2174 }
2175 }
2176
2177 let options = RecordBatchOptions::new().with_row_count(Some(row_count));
2178 RecordBatch::try_new_with_options(schema, columns, &options).unwrap()
2179 }
2180
2181 fn create_test_batch(num_rows: usize) -> RecordBatch {
2183 let schema = Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)]));
2184 let array = Int32Array::from_iter_values(0..num_rows as i32);
2185 RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
2186 }
2187 #[test]
2188 fn test_biggest_coalesce_batch_size_none_default() {
2189 let mut coalescer = BatchCoalescer::new(
2191 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2192 100,
2193 );
2194
2195 let large_batch = create_test_batch(1000);
2197 coalescer.push_batch(large_batch).unwrap();
2198
2199 let mut output_batches = vec![];
2201 while let Some(batch) = coalescer.next_completed_batch() {
2202 output_batches.push(batch);
2203 }
2204
2205 coalescer.finish_buffered_batch().unwrap();
2206 while let Some(batch) = coalescer.next_completed_batch() {
2207 output_batches.push(batch);
2208 }
2209
2210 assert_eq!(output_batches.len(), 10);
2212 for batch in output_batches {
2213 assert_eq!(batch.num_rows(), 100);
2214 }
2215 }
2216
2217 #[test]
2218 fn test_biggest_coalesce_batch_size_bypass_large_batch() {
2219 let mut coalescer = BatchCoalescer::new(
2221 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2222 100,
2223 );
2224 coalescer.set_biggest_coalesce_batch_size(Some(500));
2225
2226 let large_batch = create_test_batch(1000);
2228 coalescer.push_batch(large_batch.clone()).unwrap();
2229
2230 assert!(coalescer.has_completed_batch());
2232 let output_batch = coalescer.next_completed_batch().unwrap();
2233 assert_eq!(output_batch.num_rows(), 1000);
2234
2235 assert!(!coalescer.has_completed_batch());
2237 assert_eq!(coalescer.get_buffered_rows(), 0);
2238 }
2239
2240 #[test]
2241 fn test_biggest_coalesce_batch_size_coalesce_small_batch() {
2242 let mut coalescer = BatchCoalescer::new(
2244 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2245 100,
2246 )
2247 .with_biggest_coalesce_batch_size(Some(500));
2248
2249 let small_batch = create_test_batch(50);
2251 coalescer.push_batch(small_batch.clone()).unwrap();
2252
2253 assert!(!coalescer.has_completed_batch());
2255 assert_eq!(coalescer.get_buffered_rows(), 50);
2256
2257 coalescer.push_batch(small_batch).unwrap();
2259
2260 assert!(coalescer.has_completed_batch());
2262 let output_batch = coalescer.next_completed_batch().unwrap();
2263 let size = output_batch
2264 .column(0)
2265 .as_primitive::<Int32Type>()
2266 .get_buffer_memory_size();
2267 assert_eq!(size, 400); assert_eq!(output_batch.num_rows(), 100);
2269
2270 assert_eq!(coalescer.get_buffered_rows(), 0);
2271 }
2272
2273 #[test]
2274 fn test_biggest_coalesce_batch_size_equal_boundary() {
2275 let mut coalescer = BatchCoalescer::new(
2277 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2278 100,
2279 );
2280 coalescer.set_biggest_coalesce_batch_size(Some(500));
2281
2282 let boundary_batch = create_test_batch(500);
2284 coalescer.push_batch(boundary_batch).unwrap();
2285
2286 let mut output_count = 0;
2288 while coalescer.next_completed_batch().is_some() {
2289 output_count += 1;
2290 }
2291
2292 coalescer.finish_buffered_batch().unwrap();
2293 while coalescer.next_completed_batch().is_some() {
2294 output_count += 1;
2295 }
2296
2297 assert_eq!(output_count, 5);
2299 }
2300
2301 #[test]
2302 fn test_biggest_coalesce_batch_size_first_large_then_consecutive_bypass() {
2303 let mut coalescer = BatchCoalescer::new(
2306 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2307 100,
2308 );
2309 coalescer.set_biggest_coalesce_batch_size(Some(200));
2310
2311 let small_batch = create_test_batch(50);
2312
2313 coalescer.push_batch(small_batch).unwrap();
2315 assert_eq!(coalescer.get_buffered_rows(), 50);
2316 assert!(!coalescer.has_completed_batch());
2317
2318 let large_batch1 = create_test_batch(250);
2320 coalescer.push_batch(large_batch1).unwrap();
2321
2322 let mut completed_batches = vec![];
2324 while let Some(batch) = coalescer.next_completed_batch() {
2325 completed_batches.push(batch);
2326 }
2327 assert_eq!(completed_batches.len(), 3);
2328 assert_eq!(coalescer.get_buffered_rows(), 0);
2329
2330 let large_batch2 = create_test_batch(300);
2332 let large_batch3 = create_test_batch(400);
2333
2334 coalescer.push_batch(large_batch2).unwrap();
2336 assert!(coalescer.has_completed_batch());
2337 let output = coalescer.next_completed_batch().unwrap();
2338 assert_eq!(output.num_rows(), 300); assert_eq!(coalescer.get_buffered_rows(), 0);
2340
2341 coalescer.push_batch(large_batch3).unwrap();
2343 assert!(coalescer.has_completed_batch());
2344 let output = coalescer.next_completed_batch().unwrap();
2345 assert_eq!(output.num_rows(), 400); assert_eq!(coalescer.get_buffered_rows(), 0);
2347 }
2348
2349 #[test]
2350 fn test_biggest_coalesce_batch_size_empty_batch() {
2351 let mut coalescer = BatchCoalescer::new(
2353 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2354 100,
2355 );
2356 coalescer.set_biggest_coalesce_batch_size(Some(50));
2357
2358 let empty_batch = create_test_batch(0);
2359 coalescer.push_batch(empty_batch).unwrap();
2360
2361 assert!(!coalescer.has_completed_batch());
2363 assert_eq!(coalescer.get_buffered_rows(), 0);
2364 }
2365
2366 #[test]
2367 fn test_biggest_coalesce_batch_size_with_buffered_data_no_bypass() {
2368 let mut coalescer = BatchCoalescer::new(
2370 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2371 100,
2372 );
2373 coalescer.set_biggest_coalesce_batch_size(Some(200));
2374
2375 let small_batch = create_test_batch(30);
2377 coalescer.push_batch(small_batch.clone()).unwrap();
2378 coalescer.push_batch(small_batch).unwrap();
2379 assert_eq!(coalescer.get_buffered_rows(), 60);
2380
2381 let large_batch = create_test_batch(250);
2383 coalescer.push_batch(large_batch).unwrap();
2384
2385 let mut completed_batches = vec![];
2390 while let Some(batch) = coalescer.next_completed_batch() {
2391 completed_batches.push(batch);
2392 }
2393
2394 assert_eq!(completed_batches.len(), 3);
2395 for batch in &completed_batches {
2396 assert_eq!(batch.num_rows(), 100);
2397 }
2398 assert_eq!(coalescer.get_buffered_rows(), 10);
2399 }
2400
2401 #[test]
2402 fn test_biggest_coalesce_batch_size_zero_limit() {
2403 let mut coalescer = BatchCoalescer::new(
2405 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2406 100,
2407 );
2408 coalescer.set_biggest_coalesce_batch_size(Some(0));
2409
2410 let tiny_batch = create_test_batch(1);
2412 coalescer.push_batch(tiny_batch).unwrap();
2413
2414 assert!(coalescer.has_completed_batch());
2415 let output = coalescer.next_completed_batch().unwrap();
2416 assert_eq!(output.num_rows(), 1);
2417 }
2418
2419 #[test]
2420 fn test_biggest_coalesce_batch_size_bypass_only_when_no_buffer() {
2421 let mut coalescer = BatchCoalescer::new(
2423 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2424 100,
2425 );
2426 coalescer.set_biggest_coalesce_batch_size(Some(200));
2427
2428 let large_batch = create_test_batch(300);
2430 coalescer.push_batch(large_batch.clone()).unwrap();
2431
2432 assert!(coalescer.has_completed_batch());
2433 let output = coalescer.next_completed_batch().unwrap();
2434 assert_eq!(output.num_rows(), 300); assert_eq!(coalescer.get_buffered_rows(), 0);
2436
2437 let small_batch = create_test_batch(50);
2439 coalescer.push_batch(small_batch).unwrap();
2440 assert_eq!(coalescer.get_buffered_rows(), 50);
2441
2442 coalescer.push_batch(large_batch).unwrap();
2444
2445 let mut completed_batches = vec![];
2448 while let Some(batch) = coalescer.next_completed_batch() {
2449 completed_batches.push(batch);
2450 }
2451
2452 assert_eq!(completed_batches.len(), 3);
2453 for batch in &completed_batches {
2454 assert_eq!(batch.num_rows(), 100);
2455 }
2456 assert_eq!(coalescer.get_buffered_rows(), 50);
2457 }
2458
2459 #[test]
2460 fn test_biggest_coalesce_batch_size_consecutive_large_batches_scenario() {
2461 let mut coalescer = BatchCoalescer::new(
2463 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2464 1000,
2465 );
2466 coalescer.set_biggest_coalesce_batch_size(Some(500));
2467
2468 coalescer.push_batch(create_test_batch(20)).unwrap();
2470 coalescer.push_batch(create_test_batch(20)).unwrap();
2471 coalescer.push_batch(create_test_batch(30)).unwrap();
2472
2473 assert_eq!(coalescer.get_buffered_rows(), 70);
2474 assert!(!coalescer.has_completed_batch());
2475
2476 coalescer.push_batch(create_test_batch(700)).unwrap();
2478
2479 assert_eq!(coalescer.get_buffered_rows(), 770);
2481 assert!(!coalescer.has_completed_batch());
2482
2483 coalescer.push_batch(create_test_batch(600)).unwrap();
2485
2486 let mut outputs = vec![];
2488 while let Some(batch) = coalescer.next_completed_batch() {
2489 outputs.push(batch);
2490 }
2491 assert_eq!(outputs.len(), 2); assert_eq!(outputs[0].num_rows(), 770);
2493 assert_eq!(outputs[1].num_rows(), 600);
2494 assert_eq!(coalescer.get_buffered_rows(), 0);
2495
2496 let remaining_batches = [700, 900, 700, 600];
2498 for &size in &remaining_batches {
2499 coalescer.push_batch(create_test_batch(size)).unwrap();
2500
2501 assert!(coalescer.has_completed_batch());
2502 let output = coalescer.next_completed_batch().unwrap();
2503 assert_eq!(output.num_rows(), size);
2504 assert_eq!(coalescer.get_buffered_rows(), 0);
2505 }
2506 }
2507
2508 #[test]
2509 fn test_biggest_coalesce_batch_size_truly_consecutive_large_bypass() {
2510 let mut coalescer = BatchCoalescer::new(
2513 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2514 100,
2515 );
2516 coalescer.set_biggest_coalesce_batch_size(Some(200));
2517
2518 let large_batches = vec![
2520 create_test_batch(300),
2521 create_test_batch(400),
2522 create_test_batch(350),
2523 create_test_batch(500),
2524 ];
2525
2526 let mut all_outputs = vec![];
2527
2528 for (i, large_batch) in large_batches.into_iter().enumerate() {
2529 let expected_size = large_batch.num_rows();
2530
2531 assert_eq!(
2533 coalescer.get_buffered_rows(),
2534 0,
2535 "Buffer should be empty before batch {}",
2536 i
2537 );
2538
2539 coalescer.push_batch(large_batch).unwrap();
2540
2541 assert!(
2543 coalescer.has_completed_batch(),
2544 "Should have completed batch after pushing batch {}",
2545 i
2546 );
2547
2548 let output = coalescer.next_completed_batch().unwrap();
2549 assert_eq!(
2550 output.num_rows(),
2551 expected_size,
2552 "Batch {} should have bypassed with original size",
2553 i
2554 );
2555
2556 assert!(
2558 !coalescer.has_completed_batch(),
2559 "Should have no more completed batches after batch {}",
2560 i
2561 );
2562 assert_eq!(
2563 coalescer.get_buffered_rows(),
2564 0,
2565 "Buffer should be empty after batch {}",
2566 i
2567 );
2568
2569 all_outputs.push(output);
2570 }
2571
2572 assert_eq!(all_outputs.len(), 4);
2574 assert_eq!(all_outputs[0].num_rows(), 300);
2575 assert_eq!(all_outputs[1].num_rows(), 400);
2576 assert_eq!(all_outputs[2].num_rows(), 350);
2577 assert_eq!(all_outputs[3].num_rows(), 500);
2578 }
2579
2580 #[test]
2581 fn test_biggest_coalesce_batch_size_reset_consecutive_on_small_batch() {
2582 let mut coalescer = BatchCoalescer::new(
2584 Arc::new(Schema::new(vec![Field::new("c0", DataType::Int32, false)])),
2585 100,
2586 );
2587 coalescer.set_biggest_coalesce_batch_size(Some(200));
2588
2589 coalescer.push_batch(create_test_batch(300)).unwrap();
2591 let output = coalescer.next_completed_batch().unwrap();
2592 assert_eq!(output.num_rows(), 300);
2593
2594 coalescer.push_batch(create_test_batch(400)).unwrap();
2596 let output = coalescer.next_completed_batch().unwrap();
2597 assert_eq!(output.num_rows(), 400);
2598
2599 coalescer.push_batch(create_test_batch(50)).unwrap();
2601 assert_eq!(coalescer.get_buffered_rows(), 50);
2602
2603 coalescer.push_batch(create_test_batch(350)).unwrap();
2605
2606 let mut outputs = vec![];
2608 while let Some(batch) = coalescer.next_completed_batch() {
2609 outputs.push(batch);
2610 }
2611 assert_eq!(outputs.len(), 4);
2612 for batch in outputs {
2613 assert_eq!(batch.num_rows(), 100);
2614 }
2615 assert_eq!(coalescer.get_buffered_rows(), 0);
2616 }
2617
2618 #[test]
2619 fn test_coalasce_push_batch_with_indices() {
2620 const MID_POINT: u32 = 2333;
2621 const TOTAL_ROWS: u32 = 23333;
2622 let batch1 = uint32_batch_non_null(0..MID_POINT);
2623 let batch2 = uint32_batch_non_null((MID_POINT..TOTAL_ROWS).rev());
2624
2625 let mut coalescer = BatchCoalescer::new(
2626 Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])),
2627 TOTAL_ROWS as usize,
2628 );
2629 coalescer.push_batch(batch1).unwrap();
2630
2631 let rev_indices = (0..((TOTAL_ROWS - MID_POINT) as u64)).rev();
2632 let reversed_indices_batch = uint64_batch_non_null(rev_indices);
2633
2634 let reverse_indices = UInt64Array::from(reversed_indices_batch.column(0).to_data());
2635 coalescer
2636 .push_batch_with_indices(batch2, &reverse_indices)
2637 .unwrap();
2638
2639 coalescer.finish_buffered_batch().unwrap();
2640 let actual = coalescer.next_completed_batch().unwrap();
2641
2642 let expected = uint32_batch_non_null(0..TOTAL_ROWS);
2643
2644 assert_eq!(expected, actual);
2645 }
2646
2647 #[test]
2648 fn test_push_batch_schema_mismatch_fewer_columns() {
2649 let empty_schema = Arc::new(Schema::empty());
2651 let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2652 let batch = uint32_batch(0..5);
2653 let result = coalescer.push_batch(batch);
2654 assert!(result.is_err());
2655 let err = result.unwrap_err().to_string();
2656 assert!(
2657 err.contains("Batch has 1 columns but BatchCoalescer expects 0"),
2658 "unexpected error: {err}"
2659 );
2660 }
2661
2662 #[test]
2663 fn test_push_batch_schema_mismatch_more_columns() {
2664 let schema = Arc::new(Schema::new(vec![
2666 Field::new("c0", DataType::UInt32, false),
2667 Field::new("c1", DataType::UInt32, false),
2668 ]));
2669 let mut coalescer = BatchCoalescer::new(schema, 100);
2670 let batch = uint32_batch(0..5);
2671 let result = coalescer.push_batch(batch);
2672 assert!(result.is_err());
2673 let err = result.unwrap_err().to_string();
2674 assert!(
2675 err.contains("Batch has 1 columns but BatchCoalescer expects 2"),
2676 "unexpected error: {err}"
2677 );
2678 }
2679
2680 #[test]
2681 fn test_push_batch_schema_mismatch_two_vs_zero() {
2682 let empty_schema = Arc::new(Schema::empty());
2684 let mut coalescer = BatchCoalescer::new(empty_schema, 100);
2685 let schema = Arc::new(Schema::new(vec![
2686 Field::new("c0", DataType::UInt32, false),
2687 Field::new("c1", DataType::UInt32, false),
2688 ]));
2689 let batch = RecordBatch::try_new(
2690 schema,
2691 vec![
2692 Arc::new(UInt32Array::from(vec![1, 2, 3])),
2693 Arc::new(UInt32Array::from(vec![4, 5, 6])),
2694 ],
2695 )
2696 .unwrap();
2697 let result = coalescer.push_batch(batch);
2698 assert!(result.is_err());
2699 let err = result.unwrap_err().to_string();
2700 assert!(
2701 err.contains("Batch has 2 columns but BatchCoalescer expects 0"),
2702 "unexpected error: {err}"
2703 );
2704 }
2705
2706 #[test]
2707 fn test_size_grows_with_buffering_and_shrinks_when_draining() {
2708 let batch = uint32_batch(0..8);
2709 let mut coalescer = BatchCoalescer::new(batch.schema(), 21);
2710 let baseline = coalescer.size();
2711 assert!(baseline > 0, "size includes container capacities");
2712
2713 coalescer.push_batch(batch.clone()).unwrap();
2715 assert!(coalescer.next_completed_batch().is_none());
2716 let buffered = coalescer.size();
2717 assert!(
2718 buffered > baseline,
2719 "buffering rows should grow size ({buffered} > {baseline})"
2720 );
2721
2722 for _ in 0..10 {
2724 coalescer.push_batch(batch.clone()).unwrap();
2725 }
2726 let peak = coalescer.size();
2727 assert!(peak > buffered);
2728
2729 let mut prev = peak;
2731 let mut drained_any = false;
2732 while coalescer.next_completed_batch().is_some() {
2733 drained_any = true;
2734 let now = coalescer.size();
2735 assert!(now <= prev, "size grew while draining: {now} > {prev}");
2736 prev = now;
2737 }
2738 assert!(drained_any);
2739 assert!(
2740 prev < peak,
2741 "draining completed batches should release memory"
2742 );
2743 }
2744
2745 #[test]
2746 fn test_size_string_view_buffers_released_after_drain() {
2747 let batch = stringview_batch_repeated(
2750 1000,
2751 [Some("this string is definitely longer than 12 bytes")],
2752 );
2753 let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2754 let baseline = coalescer.size();
2755
2756 for _ in 0..20 {
2757 coalescer.push_batch(batch.clone()).unwrap();
2758 }
2759 let peak = coalescer.size();
2760 assert!(
2761 peak > baseline,
2762 "buffered string view data should grow size ({peak} > {baseline})"
2763 );
2764
2765 coalescer.finish_buffered_batch().unwrap();
2766 while coalescer.next_completed_batch().is_some() {}
2767
2768 let drained = coalescer.size();
2770 assert!(
2771 drained < peak,
2772 "draining should release buffered data ({drained} < {peak})"
2773 );
2774 }
2775
2776 #[test]
2780 fn test_size_accounting_conserved_across_cycles() {
2781 let batch = uint32_batch(0..8);
2785 let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
2786
2787 let run_cycle = |coalescer: &mut BatchCoalescer| {
2788 for _ in 0..20 {
2789 coalescer.push_batch(batch.clone()).unwrap();
2790 }
2791 coalescer.finish_buffered_batch().unwrap();
2792 let peak = coalescer.size();
2793 while coalescer.next_completed_batch().is_some() {}
2794 (peak, coalescer.size())
2795 };
2796
2797 let (peak1, drained1) = run_cycle(&mut coalescer);
2798 let (peak2, drained2) = run_cycle(&mut coalescer);
2799
2800 assert_eq!(
2801 peak1, peak2,
2802 "identical work must report identical peak size"
2803 );
2804 assert_eq!(
2805 drained1, drained2,
2806 "fully-drained size must be stable across cycles (no accounting leak)"
2807 );
2808 assert!(
2809 drained1 < peak1,
2810 "draining must release the accounted memory"
2811 );
2812 }
2813}