1mod reader_builder;
22mod remaining;
23mod scan_plan;
24
25use crate::DecodeResult;
26pub use crate::arrow::arrow_reader::RowGroupSelection;
27use crate::arrow::arrow_reader::{
28 ArrowReaderBuilder, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader,
29};
30use crate::errors::ParquetError;
31use crate::file::metadata::ParquetMetaData;
32pub use crate::util::push_buffers::PushBuffers;
33use arrow_array::RecordBatch;
34use bytes::Bytes;
35use reader_builder::{RowGroupReaderBuilder, RowGroupReaderBuilderParts};
36use remaining::{RemainingRowGroups, RemainingRowGroupsParts};
37use scan_plan::RowBudget;
38pub use scan_plan::{PageKind, PlannedRange, ScanPlan};
39use std::ops::Range;
40use std::sync::Arc;
41
42pub type ParquetPushDecoderBuilder = ArrowReaderBuilder<PushDecoderInput>;
198
199#[derive(Debug, Default)]
206pub struct PushDecoderInput {
207 buffers: PushBuffers,
209}
210
211impl ParquetPushDecoderBuilder {
214 pub fn try_new_decoder(parquet_metadata: Arc<ParquetMetaData>) -> Result<Self, ParquetError> {
222 Self::try_new_decoder_with_options(parquet_metadata, ArrowReaderOptions::default())
223 }
224
225 pub fn try_new_decoder_with_options(
231 parquet_metadata: Arc<ParquetMetaData>,
232 arrow_reader_options: ArrowReaderOptions,
233 ) -> Result<Self, ParquetError> {
234 let arrow_reader_metadata =
235 ArrowReaderMetadata::try_new(parquet_metadata, arrow_reader_options)?;
236 Ok(Self::new_with_metadata(arrow_reader_metadata))
237 }
238
239 pub fn new_with_metadata(arrow_reader_metadata: ArrowReaderMetadata) -> Self {
244 Self::new_builder(PushDecoderInput::default(), arrow_reader_metadata)
245 }
246
247 pub fn with_buffers(self, buffers: PushBuffers) -> Self {
254 Self {
255 input: PushDecoderInput { buffers },
256 ..self
257 }
258 }
259
260 pub fn with_row_group_selections(
304 mut self,
305 row_group_selections: Vec<RowGroupSelection>,
306 ) -> Self {
307 self.row_group_plan
308 .set_row_group_selections(row_group_selections);
309 self
310 }
311
312 pub fn build(self) -> Result<ParquetPushDecoder, ParquetError> {
314 let Self {
315 input: PushDecoderInput { buffers },
316 metadata: parquet_metadata,
317 schema,
318 fields,
319 batch_size,
320 row_group_plan,
321 projection,
322 filter,
323 limit,
324 offset,
325 metrics,
326 row_selection_policy,
327 max_predicate_cache_size,
328 } = self;
329
330 let filter = filter.map(|filter| {
332 filter.fuse_same_projection(parquet_metadata.file_metadata().schema_descr())
333 });
334
335 let has_predicates = filter
336 .as_ref()
337 .is_some_and(|filter| !filter.predicates.is_empty());
338
339 let row_group_reader_builder = RowGroupReaderBuilder::new(
343 batch_size,
344 projection,
345 Arc::clone(&parquet_metadata),
346 fields,
347 filter,
348 metrics,
349 max_predicate_cache_size,
350 buffers,
351 row_selection_policy,
352 );
353
354 let mut remaining_row_groups = RemainingRowGroups::new(
356 schema,
357 parquet_metadata,
358 row_group_plan,
359 RowBudget::new(offset, limit),
360 has_predicates,
361 row_group_reader_builder,
362 )?;
363 remaining_row_groups.release_unread_bytes();
367
368 Ok(ParquetPushDecoder {
369 state: ParquetDecoderState::ReadingRowGroup {
370 remaining_row_groups: Box::new(remaining_row_groups),
371 },
372 })
373 }
374}
375
376fn builder_from_remaining(parts: RemainingRowGroupsParts) -> ParquetPushDecoderBuilder {
381 let RemainingRowGroupsParts {
382 metadata,
383 schema,
384 row_group_plan,
385 offset,
386 limit,
387 reader_builder,
388 } = parts;
389 let RowGroupReaderBuilderParts {
390 batch_size,
391 projection,
392 fields,
393 filter,
394 max_predicate_cache_size,
395 metrics,
396 row_selection_policy,
397 buffers,
398 } = reader_builder;
399
400 ArrowReaderBuilder {
401 input: PushDecoderInput::default(),
402 metadata,
403 schema,
404 fields,
405 batch_size,
406 row_group_plan,
407 projection,
408 filter,
409 row_selection_policy,
410 limit,
411 offset,
412 metrics,
413 max_predicate_cache_size,
414 }
415 .with_buffers(buffers)
418}
419
420#[derive(Debug)]
437pub struct ParquetPushDecoder {
438 state: ParquetDecoderState,
444}
445
446impl ParquetPushDecoder {
447 pub fn try_decode(&mut self) -> Result<DecodeResult<RecordBatch>, ParquetError> {
478 let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished);
479 let (new_state, decode_result) = current_state.try_next_batch()?;
480 self.state = new_state;
481 Ok(decode_result)
482 }
483
484 pub fn try_next_reader(
523 &mut self,
524 ) -> Result<DecodeResult<ParquetRecordBatchReader>, ParquetError> {
525 let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished);
526 let (new_state, decode_result) = current_state.try_next_reader()?;
527 self.state = new_state;
528 Ok(decode_result)
529 }
530
531 pub fn push_range(&mut self, range: Range<u64>, data: Bytes) -> Result<(), ParquetError> {
550 self.push_ranges(vec![range], vec![data])
551 }
552
553 pub fn push_ranges(
566 &mut self,
567 ranges: Vec<Range<u64>>,
568 data: Vec<Bytes>,
569 ) -> Result<(), ParquetError> {
570 let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished);
571 self.state = current_state.push_data(ranges, data)?;
572 Ok(())
573 }
574
575 pub fn buffered_bytes(&self) -> u64 {
586 self.state.buffered_bytes()
587 }
588
589 pub fn clear_all_ranges(&mut self) {
595 self.state.clear_all_ranges();
596 }
597
598 pub fn is_at_row_group_boundary(&self) -> bool {
609 self.state.is_at_row_group_boundary()
610 }
611
612 pub fn row_groups_remaining(&self) -> usize {
615 self.state.row_groups_remaining()
616 }
617
618 pub fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
639 self.state.peek_next_row_group()
640 }
641
642 pub fn scan_plan(&self) -> ScanPlan {
645 self.state.scan_plan()
646 }
647
648 pub fn into_builder(self) -> Result<ParquetPushDecoderBuilder, ParquetError> {
706 self.state.into_builder()
707 }
708}
709
710#[derive(Debug)]
712enum ParquetDecoderState {
713 ReadingRowGroup {
715 remaining_row_groups: Box<RemainingRowGroups>,
716 },
717 DecodingRowGroup {
719 record_batch_reader: Box<ParquetRecordBatchReader>,
721 remaining_row_groups: Box<RemainingRowGroups>,
722 },
723 Finished,
725}
726
727impl ParquetDecoderState {
728 fn try_next_reader(
731 self,
732 ) -> Result<(Self, DecodeResult<ParquetRecordBatchReader>), ParquetError> {
733 let mut current_state = self;
734 loop {
735 let (next_state, decode_result) = current_state.transition()?;
736 match decode_result {
738 DecodeResult::NeedsData(ranges) => {
739 return Ok((next_state, DecodeResult::NeedsData(ranges)));
740 }
741 DecodeResult::Data(()) | DecodeResult::Finished => {}
743 }
744 match next_state {
745 Self::ReadingRowGroup { .. } => current_state = next_state,
747 Self::DecodingRowGroup {
749 record_batch_reader,
750 remaining_row_groups,
751 } => {
752 let result = DecodeResult::Data(*record_batch_reader);
753 let next_state = Self::ReadingRowGroup {
754 remaining_row_groups,
755 };
756 return Ok((next_state, result));
757 }
758 Self::Finished => {
759 return Ok((Self::Finished, DecodeResult::Finished));
760 }
761 }
762 }
763 }
764
765 fn try_next_batch(self) -> Result<(Self, DecodeResult<RecordBatch>), ParquetError> {
772 let mut current_state = self;
773 loop {
774 let (new_state, decode_result) = current_state.transition()?;
775 match decode_result {
777 DecodeResult::NeedsData(ranges) => {
778 return Ok((new_state, DecodeResult::NeedsData(ranges)));
779 }
780 DecodeResult::Data(()) | DecodeResult::Finished => {}
782 }
783 match new_state {
784 Self::ReadingRowGroup { .. } => current_state = new_state,
786 Self::DecodingRowGroup {
788 mut record_batch_reader,
789 remaining_row_groups,
790 } => {
791 match record_batch_reader.next() {
792 Some(Ok(batch)) => {
794 let result = DecodeResult::Data(batch);
795 let next_state = Self::DecodingRowGroup {
796 record_batch_reader,
797 remaining_row_groups,
798 };
799 return Ok((next_state, result));
800 }
801 None => {
803 current_state = Self::ReadingRowGroup {
804 remaining_row_groups,
805 }
806 }
807 Some(Err(e)) => {
809 return Err(ParquetError::ArrowError(e.to_string()));
811 }
812 }
813 }
814 Self::Finished => {
815 return Ok((Self::Finished, DecodeResult::Finished));
816 }
817 }
818 }
819 }
820
821 fn transition(self) -> Result<(Self, DecodeResult<()>), ParquetError> {
826 let data_ready = DecodeResult::Data(());
828 match self {
829 Self::ReadingRowGroup {
830 mut remaining_row_groups,
831 } => {
832 match remaining_row_groups.try_next_reader()? {
833 DecodeResult::Data(record_batch_reader) => {
835 Ok((
837 Self::DecodingRowGroup {
838 record_batch_reader: Box::new(record_batch_reader),
839 remaining_row_groups,
840 },
841 data_ready,
842 ))
843 }
844 DecodeResult::NeedsData(ranges) => {
845 Ok((
848 Self::ReadingRowGroup {
849 remaining_row_groups,
850 },
851 DecodeResult::NeedsData(ranges),
852 ))
853 }
854 DecodeResult::Finished => {
856 Ok((Self::Finished, DecodeResult::Finished))
858 }
859 }
860 }
861 Self::DecodingRowGroup { .. } => Ok((self, data_ready)),
863 Self::Finished => Ok((self, DecodeResult::Finished)),
865 }
866 }
867
868 pub fn push_data(
872 self,
873 ranges: Vec<Range<u64>>,
874 data: Vec<Bytes>,
875 ) -> Result<Self, ParquetError> {
876 match self {
877 ParquetDecoderState::ReadingRowGroup {
878 mut remaining_row_groups,
879 } => {
880 remaining_row_groups.push_data(ranges, data)?;
882 Ok(ParquetDecoderState::ReadingRowGroup {
883 remaining_row_groups,
884 })
885 }
886 ParquetDecoderState::DecodingRowGroup {
888 record_batch_reader,
889 mut remaining_row_groups,
890 } => {
891 remaining_row_groups.push_data(ranges, data)?;
892 Ok(ParquetDecoderState::DecodingRowGroup {
893 record_batch_reader,
894 remaining_row_groups,
895 })
896 }
897 ParquetDecoderState::Finished => Err(ParquetError::General(
898 "Cannot push data to a finished decoder".to_string(),
899 )),
900 }
901 }
902
903 fn buffered_bytes(&self) -> u64 {
905 match self {
906 ParquetDecoderState::ReadingRowGroup {
907 remaining_row_groups,
908 } => remaining_row_groups.buffered_bytes(),
909 ParquetDecoderState::DecodingRowGroup {
910 record_batch_reader: _,
911 remaining_row_groups,
912 } => remaining_row_groups.buffered_bytes(),
913 ParquetDecoderState::Finished => 0,
914 }
915 }
916
917 fn clear_all_ranges(&mut self) {
919 match self {
920 ParquetDecoderState::ReadingRowGroup {
921 remaining_row_groups,
922 } => remaining_row_groups.clear_all_ranges(),
923 ParquetDecoderState::DecodingRowGroup {
924 record_batch_reader: _,
925 remaining_row_groups,
926 } => remaining_row_groups.clear_all_ranges(),
927 ParquetDecoderState::Finished => {}
928 }
929 }
930
931 fn is_at_row_group_boundary(&self) -> bool {
932 match self {
933 ParquetDecoderState::ReadingRowGroup {
934 remaining_row_groups,
935 } => remaining_row_groups.is_at_row_group_boundary(),
936 ParquetDecoderState::DecodingRowGroup { .. } => false,
940 ParquetDecoderState::Finished => false,
941 }
942 }
943
944 fn scan_plan(&self) -> ScanPlan {
946 match self {
947 ParquetDecoderState::ReadingRowGroup {
948 remaining_row_groups,
949 }
950 | ParquetDecoderState::DecodingRowGroup {
951 remaining_row_groups,
952 ..
953 } => remaining_row_groups.scan_plan(),
954 ParquetDecoderState::Finished => ScanPlan::empty(),
955 }
956 }
957
958 fn row_groups_remaining(&self) -> usize {
959 match self {
960 ParquetDecoderState::ReadingRowGroup {
961 remaining_row_groups,
962 } => remaining_row_groups.row_groups_remaining(),
963 ParquetDecoderState::DecodingRowGroup {
964 remaining_row_groups,
965 ..
966 } => remaining_row_groups.row_groups_remaining(),
967 ParquetDecoderState::Finished => 0,
968 }
969 }
970
971 fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
979 match self {
980 ParquetDecoderState::ReadingRowGroup {
981 remaining_row_groups,
982 } => remaining_row_groups.peek_next_row_group(),
983 ParquetDecoderState::DecodingRowGroup {
984 remaining_row_groups,
985 ..
986 } => remaining_row_groups.peek_next_row_group(),
987 ParquetDecoderState::Finished => Ok(None),
988 }
989 }
990
991 fn into_builder(self) -> Result<ParquetPushDecoderBuilder, ParquetError> {
992 let remaining_row_groups = match self {
993 ParquetDecoderState::ReadingRowGroup {
994 remaining_row_groups,
995 } => remaining_row_groups,
996 ParquetDecoderState::DecodingRowGroup { .. } => {
997 return Err(ParquetError::General(
998 "into_builder called while a row group is being decoded; \
999 check is_at_row_group_boundary() first"
1000 .to_string(),
1001 ));
1002 }
1003 ParquetDecoderState::Finished => {
1004 return Err(ParquetError::General(
1005 "into_builder called on a finished decoder".to_string(),
1006 ));
1007 }
1008 };
1009 if !remaining_row_groups.is_at_row_group_boundary() {
1010 return Err(ParquetError::General(
1011 "into_builder called mid-row-group; check is_at_row_group_boundary() first"
1012 .to_string(),
1013 ));
1014 }
1015 Ok(builder_from_remaining(remaining_row_groups.into_parts()))
1016 }
1017}
1018
1019#[cfg(test)]
1020mod test {
1021 use super::*;
1022 use crate::DecodeResult;
1023 use crate::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector};
1024 use crate::arrow::push_decoder::reader_builder::column_chunk_range;
1025 use crate::arrow::push_decoder::{
1026 ParquetPushDecoder, ParquetPushDecoderBuilder, RowGroupSelection,
1027 };
1028 use crate::arrow::{ArrowWriter, ProjectionMask};
1029 use crate::errors::ParquetError;
1030 use crate::file::metadata::page_index::PageIndexBuilder;
1031 use crate::file::metadata::{PageIndexPolicy, ParquetMetaDataPushDecoder};
1032 use crate::file::properties::WriterProperties;
1033 use arrow::compute::kernels::cmp::{gt, lt};
1034 use arrow_array::cast::AsArray;
1035 use arrow_array::types::Int64Type;
1036 use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray};
1037 use arrow_buffer::BooleanBuffer;
1038 use arrow_select::concat::concat_batches;
1039 use bytes::Bytes;
1040 use std::fmt::Debug;
1041 use std::ops::Range;
1042 use std::sync::{Arc, LazyLock};
1043
1044 #[test]
1047 fn test_decoder_size() {
1048 assert_eq!(std::mem::size_of::<ParquetDecoderState>(), 24);
1049 }
1050
1051 #[test]
1054 fn test_decoder_all_data() {
1055 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1056 .unwrap()
1057 .build()
1058 .unwrap();
1059
1060 decoder
1061 .push_range(test_file_range(), TEST_FILE_DATA.clone())
1062 .unwrap();
1063
1064 let results = vec![
1065 expect_data(decoder.try_decode()),
1067 expect_data(decoder.try_decode()),
1069 ];
1070 expect_finished(decoder.try_decode());
1071
1072 let all_output = concat_batches(&TEST_BATCH.schema(), &results).unwrap();
1073 assert_eq!(all_output, *TEST_BATCH);
1075 }
1076
1077 #[test]
1080 fn test_decoder_incremental() {
1081 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1082 .unwrap()
1083 .build()
1084 .unwrap();
1085
1086 let mut results = vec![];
1087
1088 let ranges = expect_needs_data(decoder.try_decode());
1090 let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum();
1091 push_ranges_to_decoder(&mut decoder, ranges);
1092 assert_eq!(decoder.buffered_bytes(), num_bytes_requested);
1094 results.push(expect_data(decoder.try_decode()));
1095 assert_eq!(decoder.buffered_bytes(), 0);
1097
1098 let ranges = expect_needs_data(decoder.try_decode());
1100 let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum();
1101 push_ranges_to_decoder(&mut decoder, ranges);
1102 assert_eq!(decoder.buffered_bytes(), num_bytes_requested);
1104 results.push(expect_data(decoder.try_decode()));
1105 assert_eq!(decoder.buffered_bytes(), 0);
1107 expect_finished(decoder.try_decode());
1108
1109 let all_output = concat_batches(&TEST_BATCH.schema(), &results).unwrap();
1111 assert_eq!(all_output, *TEST_BATCH);
1112 }
1113
1114 #[test]
1117 fn test_decoder_clear_all_ranges() {
1118 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1119 .unwrap()
1120 .with_batch_size(100)
1121 .build()
1122 .unwrap();
1123
1124 decoder
1125 .push_range(test_file_range(), TEST_FILE_DATA.clone())
1126 .unwrap();
1127 assert_eq!(decoder.buffered_bytes(), test_file_len());
1128
1129 let batch1 = expect_data(decoder.try_decode());
1133 assert_eq!(batch1, TEST_BATCH.slice(0, 100));
1134 assert_eq!(
1135 decoder.buffered_bytes(),
1136 test_file_len() - test_file_row_group_bytes(0)
1137 );
1138
1139 decoder.clear_all_ranges();
1141 assert_eq!(decoder.buffered_bytes(), 0);
1142
1143 let batch2 = expect_data(decoder.try_decode());
1146 assert_eq!(batch2, TEST_BATCH.slice(100, 100));
1147 assert_eq!(decoder.buffered_bytes(), 0);
1148
1149 let ranges = expect_needs_data(decoder.try_decode());
1152 let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum();
1153 push_ranges_to_decoder(&mut decoder, ranges);
1154 assert_eq!(decoder.buffered_bytes(), num_bytes_requested);
1155
1156 let batch3 = expect_data(decoder.try_decode());
1157 assert_eq!(batch3, TEST_BATCH.slice(200, 100));
1158 assert_eq!(decoder.buffered_bytes(), 0);
1159
1160 let batch4 = expect_data(decoder.try_decode());
1161 assert_eq!(batch4, TEST_BATCH.slice(300, 100));
1162 assert_eq!(decoder.buffered_bytes(), 0);
1163
1164 expect_finished(decoder.try_decode());
1165 }
1166
1167 #[test]
1169 fn test_decoder_partial() {
1170 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1171 .unwrap()
1172 .build()
1173 .unwrap();
1174
1175 let ranges = expect_needs_data(decoder.try_decode());
1177 push_ranges_to_decoder(&mut decoder, ranges);
1178
1179 let batch1 = expect_data(decoder.try_decode());
1180 let expected1 = TEST_BATCH.slice(0, 200);
1181 assert_eq!(batch1, expected1);
1182
1183 let ranges = expect_needs_data(decoder.try_decode());
1185 let (ranges1, ranges2) = ranges.split_at(ranges.len() / 2);
1186 assert!(!ranges1.is_empty());
1187 assert!(!ranges2.is_empty());
1188 push_ranges_to_decoder(&mut decoder, ranges1.to_vec());
1190
1191 let ranges = expect_needs_data(decoder.try_decode());
1193 assert_eq!(ranges, ranges2); push_ranges_to_decoder(&mut decoder, vec![]);
1196 let ranges = expect_needs_data(decoder.try_decode());
1197 assert_eq!(ranges, ranges2); push_ranges_to_decoder(&mut decoder, ranges);
1199
1200 let batch2 = expect_data(decoder.try_decode());
1201 let expected2 = TEST_BATCH.slice(200, 200);
1202 assert_eq!(batch2, expected2);
1203
1204 expect_finished(decoder.try_decode());
1205 }
1206
1207 #[test]
1211 fn test_decoder_first_pages_only() {
1212 let metadata = test_file_parquet_metadata_with_offset_index();
1213 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata))
1214 .unwrap()
1215 .with_batch_size(100)
1218 .build()
1219 .unwrap();
1220
1221 let ranges = expect_needs_data(decoder.try_decode());
1224 assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]);
1225
1226 let page_index = metadata.page_index_for_row_group(0);
1229 let row_group = metadata.row_group(0);
1230 let mut first_page_ranges = vec![];
1231 let mut second_page_ranges = vec![];
1232 for (idx, column) in row_group.columns().iter().enumerate() {
1233 let (start, len) = column.byte_range();
1234 let locations = page_index.page_locations(idx).unwrap();
1235 assert_eq!(locations.len(), 2, "expected 2 data pages per column chunk");
1236 let second_page_start = locations[1].offset as u64;
1237 first_page_ranges.push(start..second_page_start);
1238 second_page_ranges.push(second_page_start..start + len);
1239 }
1240 assert_eq!(first_page_ranges, vec![4..1734, 1860..3590, 3716..10936]);
1242 assert_eq!(
1243 second_page_ranges,
1244 vec![1734..1860, 3590..3716, 10936..11062]
1245 );
1246
1247 push_ranges_to_decoder(&mut decoder, first_page_ranges);
1250
1251 let ranges = expect_needs_data(decoder.try_decode());
1257 assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]);
1258
1259 push_ranges_to_decoder(&mut decoder, second_page_ranges);
1263 let ranges = expect_needs_data(decoder.try_decode());
1264 assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]);
1265
1266 push_ranges_to_decoder(&mut decoder, ranges);
1269 let batch = expect_data(decoder.try_decode());
1270 assert_eq!(batch, TEST_BATCH.slice(0, 100));
1271 let batch = expect_data(decoder.try_decode());
1272 assert_eq!(batch, TEST_BATCH.slice(100, 100));
1273 }
1274
1275 fn decode_with_coalesced_pushes(mut decoder: ParquetPushDecoder) -> Vec<RecordBatch> {
1281 let mut batches = vec![];
1282 loop {
1283 match decoder.try_decode().unwrap() {
1284 DecodeResult::NeedsData(ranges) => {
1285 let start = ranges.iter().map(|r| r.start).min().unwrap();
1286 let end = ranges.iter().map(|r| r.end).max().unwrap();
1287 let coalesced = start..end;
1288 push_ranges_to_decoder(&mut decoder, vec![coalesced]);
1289 }
1290 DecodeResult::Data(batch) => {
1291 assert_eq!(decoder.buffered_bytes(), 0);
1293 batches.push(batch);
1294 }
1295 DecodeResult::Finished => return batches,
1296 }
1297 }
1298 }
1299
1300 #[test]
1303 fn test_decoder_releases_coalesced_push() {
1304 let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1305 .unwrap()
1306 .build()
1307 .unwrap();
1308 let batches = decode_with_coalesced_pushes(decoder);
1309 let all_output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
1310 assert_eq!(all_output, *TEST_BATCH);
1311 }
1312
1313 #[test]
1316 fn test_decoder_releases_unrequested_bytes_of_row_group() {
1317 let builder =
1318 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1319 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1320 let decoder = builder
1321 .with_projection(ProjectionMask::columns(&schema_descr, ["a", "c"]))
1322 .build()
1323 .unwrap();
1324 let batches = decode_with_coalesced_pushes(decoder);
1325 let all_output = concat_batches(&batches[0].schema(), &batches).unwrap();
1326 assert_eq!(all_output, TEST_BATCH.project(&[0, 2]).unwrap());
1327 }
1328
1329 #[test]
1332 fn test_decoder_keeps_bytes_of_a_row_group_that_is_read_again() {
1333 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1334 .unwrap()
1335 .with_row_groups(vec![0, 1, 0])
1336 .build()
1337 .unwrap();
1338 prefetch_test_file(&mut decoder);
1339 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 200));
1341 assert_eq!(decoder.buffered_bytes(), test_file_len());
1342
1343 assert_eq!(
1344 expect_data(decoder.try_decode()),
1345 TEST_BATCH.slice(200, 200)
1346 );
1347 assert_eq!(
1348 decoder.buffered_bytes(),
1349 test_file_len() - test_file_row_group_bytes(1)
1350 );
1351
1352 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 200));
1353 assert_eq!(
1354 decoder.buffered_bytes(),
1355 test_file_len() - test_file_row_group_bytes(0) - test_file_row_group_bytes(1)
1356 );
1357 expect_finished(decoder.try_decode());
1358 }
1359
1360 #[test]
1363 fn test_decoder_releases_coalesced_push_with_predicate() {
1364 let builder =
1365 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1366 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1367 let row_filter_a = ArrowPredicateFn::new(
1368 ProjectionMask::columns(&schema_descr, ["a"]),
1369 |batch: RecordBatch| {
1370 let scalar_250 = Int64Array::new_scalar(250);
1371 let column = batch.column(0).as_primitive::<Int64Type>();
1372 gt(column, &scalar_250)
1373 },
1374 );
1375 let decoder = builder
1376 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1377 .build()
1378 .unwrap();
1379 let batches = decode_with_coalesced_pushes(decoder);
1380 let all_output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
1381 assert_eq!(all_output, TEST_BATCH.slice(251, 149));
1382 }
1383
1384 #[test]
1387 fn test_decoder_selection_does_one_request() {
1388 let builder =
1389 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1390
1391 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1392
1393 let mut decoder = builder
1394 .with_projection(
1395 ProjectionMask::columns(&schema_descr, ["a", "b"]), )
1397 .build()
1398 .unwrap();
1399
1400 let ranges = expect_needs_data(decoder.try_decode());
1402 push_ranges_to_decoder(&mut decoder, ranges);
1403
1404 let batch1 = expect_data(decoder.try_decode());
1405 let expected1 = TEST_BATCH.slice(0, 200).project(&[0, 1]).unwrap();
1406 assert_eq!(batch1, expected1);
1407
1408 let ranges = expect_needs_data(decoder.try_decode());
1410 push_ranges_to_decoder(&mut decoder, ranges);
1411
1412 let batch2 = expect_data(decoder.try_decode());
1413 let expected2 = TEST_BATCH.slice(200, 200).project(&[0, 1]).unwrap();
1414 assert_eq!(batch2, expected2);
1415
1416 expect_finished(decoder.try_decode());
1417 }
1418
1419 #[test]
1422 fn test_decoder_single_filter_partial() {
1423 let builder =
1424 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1425
1426 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1429
1430 let row_filter_a = ArrowPredicateFn::new(
1432 ProjectionMask::columns(&schema_descr, ["a", "b"]),
1434 |batch: RecordBatch| {
1435 let scalar_250 = Int64Array::new_scalar(250);
1436 let column = batch.column(0).as_primitive::<Int64Type>();
1437 gt(column, &scalar_250)
1438 },
1439 );
1440
1441 let mut decoder = builder
1442 .with_projection(
1443 ProjectionMask::columns(&schema_descr, ["a"]), )
1446 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1447 .build()
1448 .unwrap();
1449
1450 let ranges = expect_needs_data(decoder.try_decode());
1452 let (ranges1, ranges2) = ranges.split_at(ranges.len() / 2);
1454 assert!(!ranges1.is_empty());
1455 assert!(!ranges2.is_empty());
1456 push_ranges_to_decoder(&mut decoder, ranges1.to_vec());
1457 let ranges = expect_needs_data(decoder.try_decode());
1459 assert_eq!(ranges, ranges2); let ranges = expect_needs_data(decoder.try_decode());
1461 assert_eq!(ranges, ranges2); push_ranges_to_decoder(&mut decoder, ranges2.to_vec());
1463
1464 let ranges = expect_needs_data(decoder.try_decode());
1469 push_ranges_to_decoder(&mut decoder, ranges);
1470
1471 let batch = expect_data(decoder.try_decode());
1472 let expected = TEST_BATCH.slice(251, 149).project(&[0]).unwrap();
1473 assert_eq!(batch, expected);
1474
1475 expect_finished(decoder.try_decode());
1476 }
1477
1478 #[test]
1480 fn test_decoder_single_filter_and_row_selection() {
1481 let builder =
1482 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1483
1484 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1487
1488 let row_filter_a = ArrowPredicateFn::new(
1490 ProjectionMask::columns(&schema_descr, ["a"]),
1491 |batch: RecordBatch| {
1492 let scalar_250 = Int64Array::new_scalar(250);
1493 let column = batch.column(0).as_primitive::<Int64Type>();
1494 gt(column, &scalar_250)
1495 },
1496 );
1497
1498 let mut decoder = builder
1499 .with_projection(
1500 ProjectionMask::columns(&schema_descr, ["b"]), )
1503 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1504 .with_row_selection(RowSelection::from(vec![
1505 RowSelector::skip(200), RowSelector::select(100), RowSelector::skip(100),
1508 ]))
1509 .build()
1510 .unwrap();
1511
1512 let ranges = expect_needs_data(decoder.try_decode());
1516 push_ranges_to_decoder(&mut decoder, ranges);
1517
1518 let ranges = expect_needs_data(decoder.try_decode());
1520 push_ranges_to_decoder(&mut decoder, ranges);
1521
1522 let batch = expect_data(decoder.try_decode());
1523 let expected = TEST_BATCH.slice(251, 49).project(&[1]).unwrap();
1524 assert_eq!(batch, expected);
1525
1526 expect_finished(decoder.try_decode());
1527 }
1528
1529 #[test]
1531 fn test_decoder_multi_filters() {
1532 let builder =
1534 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1535
1536 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1541
1542 let row_filter_a = ArrowPredicateFn::new(
1544 ProjectionMask::columns(&schema_descr, ["a"]),
1545 |batch: RecordBatch| {
1546 let scalar_175 = Int64Array::new_scalar(175);
1547 let column = batch.column(0).as_primitive::<Int64Type>();
1548 gt(column, &scalar_175)
1549 },
1550 );
1551
1552 let row_filter_b = ArrowPredicateFn::new(
1554 ProjectionMask::columns(&schema_descr, ["b"]),
1555 |batch: RecordBatch| {
1556 let scalar_625 = Int64Array::new_scalar(625);
1557 let column = batch.column(0).as_primitive::<Int64Type>();
1558 lt(column, &scalar_625)
1559 },
1560 );
1561
1562 let mut decoder = builder
1563 .with_projection(
1564 ProjectionMask::columns(&schema_descr, ["c"]), )
1566 .with_row_filter(RowFilter::new(vec![
1567 Box::new(row_filter_a),
1568 Box::new(row_filter_b),
1569 ]))
1570 .build()
1571 .unwrap();
1572
1573 let ranges = expect_needs_data(decoder.try_decode());
1575 push_ranges_to_decoder(&mut decoder, ranges);
1576
1577 let ranges = expect_needs_data(decoder.try_decode());
1579 push_ranges_to_decoder(&mut decoder, ranges);
1580
1581 let ranges = expect_needs_data(decoder.try_decode());
1583 push_ranges_to_decoder(&mut decoder, ranges);
1584
1585 let batch1 = expect_data(decoder.try_decode());
1587 let expected1 = TEST_BATCH.slice(176, 24).project(&[2]).unwrap();
1588 assert_eq!(batch1, expected1);
1589
1590 let ranges = expect_needs_data(decoder.try_decode());
1592 push_ranges_to_decoder(&mut decoder, ranges);
1593
1594 let ranges = expect_needs_data(decoder.try_decode());
1596 push_ranges_to_decoder(&mut decoder, ranges);
1597
1598 let ranges = expect_needs_data(decoder.try_decode());
1600 push_ranges_to_decoder(&mut decoder, ranges);
1601
1602 let batch2 = expect_data(decoder.try_decode());
1604 let expected2 = TEST_BATCH.slice(200, 25).project(&[2]).unwrap();
1605 assert_eq!(batch2, expected2);
1606
1607 expect_finished(decoder.try_decode());
1608 }
1609
1610 #[test]
1612 fn test_decoder_same_projection_filters() {
1613 let builder =
1614 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1615 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1616 let projection_a = ProjectionMask::columns(&schema_descr, ["a"]);
1617
1618 let row_filter_gt = ArrowPredicateFn::new(projection_a.clone(), |batch: RecordBatch| {
1619 let column = batch.column(0).as_primitive::<Int64Type>();
1620 gt(column, &Int64Array::new_scalar(175))
1621 });
1622 let row_filter_lt = ArrowPredicateFn::new(projection_a, |batch: RecordBatch| {
1623 let column = batch.column(0).as_primitive::<Int64Type>();
1624 assert!(column.iter().flatten().all(|value| value > 175));
1625 lt(column, &Int64Array::new_scalar(190))
1626 });
1627
1628 let mut decoder = builder
1629 .with_projection(ProjectionMask::columns(&schema_descr, ["c"]))
1630 .with_row_filter(RowFilter::new(vec![
1631 Box::new(row_filter_gt),
1632 Box::new(row_filter_lt),
1633 ]))
1634 .with_batch_size(50)
1635 .build()
1636 .unwrap();
1637
1638 let ranges = expect_needs_data(decoder.try_decode());
1640 push_ranges_to_decoder(&mut decoder, ranges);
1641 let ranges = expect_needs_data(decoder.try_decode());
1642 push_ranges_to_decoder(&mut decoder, ranges);
1643
1644 let batch = expect_data(decoder.try_decode());
1645 let expected = TEST_BATCH.slice(176, 14).project(&[2]).unwrap();
1646 assert_eq!(batch, expected);
1647
1648 let ranges = expect_needs_data(decoder.try_decode());
1651 push_ranges_to_decoder(&mut decoder, ranges);
1652 expect_finished(decoder.try_decode());
1653 }
1654
1655 #[test]
1658 fn test_decoder_same_projection_filters_with_limit() {
1659 use std::sync::atomic::{AtomicUsize, Ordering};
1660
1661 let builder =
1662 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1663 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1664 let projection_a = ProjectionMask::columns(&schema_descr, ["a"]);
1665
1666 let rows_evaluated = Arc::new(AtomicUsize::new(0));
1667 let rows_evaluated_for_predicate = Arc::clone(&rows_evaluated);
1668 let row_filter_gt =
1669 ArrowPredicateFn::new(projection_a.clone(), move |batch: RecordBatch| {
1670 rows_evaluated_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed);
1671 let column = batch.column(0).as_primitive::<Int64Type>();
1672 gt(column, &Int64Array::new_scalar(175))
1673 });
1674 let row_filter_lt = ArrowPredicateFn::new(projection_a, |batch: RecordBatch| {
1675 let column = batch.column(0).as_primitive::<Int64Type>();
1676 lt(column, &Int64Array::new_scalar(190))
1677 });
1678
1679 let mut decoder = builder
1680 .with_projection(ProjectionMask::columns(&schema_descr, ["c"]))
1681 .with_row_filter(RowFilter::new(vec![
1682 Box::new(row_filter_gt),
1683 Box::new(row_filter_lt),
1684 ]))
1685 .with_batch_size(10)
1686 .with_limit(5)
1687 .build()
1688 .unwrap();
1689
1690 let ranges = expect_needs_data(decoder.try_decode());
1692 push_ranges_to_decoder(&mut decoder, ranges);
1693 let ranges = expect_needs_data(decoder.try_decode());
1694 push_ranges_to_decoder(&mut decoder, ranges);
1695
1696 let batch = expect_data(decoder.try_decode());
1697 let expected = TEST_BATCH.slice(176, 5).project(&[2]).unwrap();
1698 assert_eq!(batch, expected);
1699
1700 expect_finished(decoder.try_decode());
1702
1703 let evaluated = rows_evaluated.load(Ordering::Relaxed);
1706 assert!(evaluated <= 190, "evaluated {evaluated} rows");
1707 }
1708
1709 #[test]
1712 fn test_decoder_reuses_filter_pages() {
1713 let builder =
1715 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1716
1717 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1720
1721 let row_filter_a = ArrowPredicateFn::new(
1723 ProjectionMask::columns(&schema_descr, ["a"]),
1724 |batch: RecordBatch| {
1725 let scalar_250 = Int64Array::new_scalar(250);
1726 let column = batch.column(0).as_primitive::<Int64Type>();
1727 gt(column, &scalar_250)
1728 },
1729 );
1730
1731 let mut decoder = builder
1732 .with_projection(
1733 ProjectionMask::columns(&schema_descr, ["a"]), )
1736 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1737 .build()
1738 .unwrap();
1739
1740 let ranges = expect_needs_data(decoder.try_decode());
1742 push_ranges_to_decoder(&mut decoder, ranges);
1743
1744 let ranges = expect_needs_data(decoder.try_decode());
1748 push_ranges_to_decoder(&mut decoder, ranges);
1749
1750 let batch = expect_data(decoder.try_decode());
1754 let expected = TEST_BATCH.slice(251, 149).project(&[0]).unwrap();
1755 assert_eq!(batch, expected);
1756
1757 expect_finished(decoder.try_decode());
1758 }
1759
1760 #[test]
1761 fn test_decoder_empty_filters() {
1762 let builder =
1763 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1764 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1765
1766 let mut decoder = builder
1768 .with_projection(
1769 ProjectionMask::columns(&schema_descr, ["c"]), )
1771 .with_row_filter(RowFilter::new(vec![
1772 ]))
1774 .build()
1775 .unwrap();
1776
1777 let ranges = expect_needs_data(decoder.try_decode());
1779 push_ranges_to_decoder(&mut decoder, ranges);
1780
1781 let batch1 = expect_data(decoder.try_decode());
1783 let expected1 = TEST_BATCH.slice(0, 200).project(&[2]).unwrap();
1784 assert_eq!(batch1, expected1);
1785
1786 let ranges = expect_needs_data(decoder.try_decode());
1788 push_ranges_to_decoder(&mut decoder, ranges);
1789
1790 let batch2 = expect_data(decoder.try_decode());
1792 let expected2 = TEST_BATCH.slice(200, 200).project(&[2]).unwrap();
1793
1794 assert_eq!(batch2, expected2);
1795
1796 expect_finished(decoder.try_decode());
1797 }
1798
1799 #[test]
1807 fn test_decoder_filter_with_limit_short_circuits_within_row_group() {
1808 use std::sync::atomic::{AtomicUsize, Ordering};
1809
1810 let builder =
1811 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1812 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1813
1814 let rows_filtered = Arc::new(AtomicUsize::new(0));
1815 let rows_filtered_for_predicate = Arc::clone(&rows_filtered);
1816
1817 let row_filter_a = ArrowPredicateFn::new(
1818 ProjectionMask::columns(&schema_descr, ["a"]),
1819 move |batch: RecordBatch| {
1820 rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed);
1821 let scalar_175 = Int64Array::new_scalar(175);
1822 let column = batch.column(0).as_primitive::<Int64Type>();
1823 gt(column, &scalar_175)
1824 },
1825 );
1826
1827 let mut decoder = builder
1834 .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
1835 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1836 .with_batch_size(10)
1837 .with_limit(10)
1838 .build()
1839 .unwrap();
1840
1841 let ranges = expect_needs_data(decoder.try_decode());
1843 push_ranges_to_decoder(&mut decoder, ranges);
1844
1845 let batch = expect_data(decoder.try_decode());
1847 let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap();
1848 assert_eq!(batch, expected);
1849
1850 expect_finished(decoder.try_decode());
1854
1855 let evaluated = rows_filtered.load(Ordering::Relaxed);
1859 assert!(
1860 evaluated <= 190,
1861 "predicate evaluated {evaluated} rows; expected ≤ 190 (stop within batch containing 11th match)"
1862 );
1863 }
1864
1865 #[test]
1869 fn test_decoder_filter_with_limit_skips_later_row_groups() {
1870 let builder =
1871 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1872 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1873
1874 let row_filter_a = ArrowPredicateFn::new(
1878 ProjectionMask::columns(&schema_descr, ["a"]),
1879 |batch: RecordBatch| {
1880 let scalar_175 = Int64Array::new_scalar(175);
1881 let column = batch.column(0).as_primitive::<Int64Type>();
1882 gt(column, &scalar_175)
1883 },
1884 );
1885
1886 let mut decoder = builder
1887 .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
1888 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1889 .with_limit(5)
1890 .build()
1891 .unwrap();
1892
1893 let ranges = expect_needs_data(decoder.try_decode());
1895 push_ranges_to_decoder(&mut decoder, ranges);
1896
1897 let batch = expect_data(decoder.try_decode());
1899 let expected = TEST_BATCH.slice(176, 5).project(&[0]).unwrap();
1900 assert_eq!(batch, expected);
1901
1902 expect_finished(decoder.try_decode());
1905 }
1906
1907 #[test]
1918 fn test_decoder_filter_with_offset_and_limit() {
1919 let builder =
1920 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1921 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1922
1923 let row_filter_a = ArrowPredicateFn::new(
1924 ProjectionMask::columns(&schema_descr, ["a"]),
1925 |batch: RecordBatch| {
1926 let scalar_175 = Int64Array::new_scalar(175);
1927 let column = batch.column(0).as_primitive::<Int64Type>();
1928 gt(column, &scalar_175)
1929 },
1930 );
1931
1932 let mut decoder = builder
1933 .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
1934 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1935 .with_offset(10)
1936 .with_limit(5)
1937 .build()
1938 .unwrap();
1939
1940 let ranges = expect_needs_data(decoder.try_decode());
1941 push_ranges_to_decoder(&mut decoder, ranges);
1942
1943 let batch = expect_data(decoder.try_decode());
1944 let expected = TEST_BATCH.slice(186, 5).project(&[0]).unwrap();
1945 assert_eq!(batch, expected);
1946
1947 expect_finished(decoder.try_decode());
1948 }
1949
1950 #[test]
1959 fn test_decoder_multi_filters_with_limit() {
1960 use std::sync::atomic::{AtomicUsize, Ordering};
1961
1962 let builder =
1963 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1964 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1965
1966 let first_predicate_rows = Arc::new(AtomicUsize::new(0));
1967 let second_predicate_rows = Arc::new(AtomicUsize::new(0));
1968
1969 let first_predicate_rows_for_filter = Arc::clone(&first_predicate_rows);
1970 let row_filter_a = ArrowPredicateFn::new(
1971 ProjectionMask::columns(&schema_descr, ["a"]),
1972 move |batch: RecordBatch| {
1973 first_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed);
1974 let scalar_175 = Int64Array::new_scalar(175);
1975 let column = batch.column(0).as_primitive::<Int64Type>();
1976 gt(column, &scalar_175)
1977 },
1978 );
1979
1980 let second_predicate_rows_for_filter = Arc::clone(&second_predicate_rows);
1981 let row_filter_b = ArrowPredicateFn::new(
1982 ProjectionMask::columns(&schema_descr, ["b"]),
1983 move |batch: RecordBatch| {
1984 second_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed);
1985 let scalar_625 = Int64Array::new_scalar(625);
1986 let column = batch.column(0).as_primitive::<Int64Type>();
1987 lt(column, &scalar_625)
1988 },
1989 );
1990
1991 let mut decoder = builder
1992 .with_projection(ProjectionMask::columns(&schema_descr, ["c"]))
1993 .with_row_filter(RowFilter::new(vec![
1994 Box::new(row_filter_a),
1995 Box::new(row_filter_b),
1996 ]))
1997 .with_batch_size(10)
1998 .with_limit(10)
1999 .build()
2000 .unwrap();
2001
2002 let ranges = expect_needs_data(decoder.try_decode());
2004 push_ranges_to_decoder(&mut decoder, ranges);
2005
2006 let ranges = expect_needs_data(decoder.try_decode());
2008 push_ranges_to_decoder(&mut decoder, ranges);
2009
2010 let ranges = expect_needs_data(decoder.try_decode());
2012 push_ranges_to_decoder(&mut decoder, ranges);
2013
2014 let batch = expect_data(decoder.try_decode());
2015 let expected = TEST_BATCH.slice(176, 10).project(&[2]).unwrap();
2016 assert_eq!(batch, expected);
2017
2018 expect_finished(decoder.try_decode());
2020
2021 assert_eq!(first_predicate_rows.load(Ordering::Relaxed), 200);
2022 assert!(
2023 second_predicate_rows.load(Ordering::Relaxed) < 24,
2024 "final predicate should short-circuit before consuming all 24 rows selected by the first predicate"
2025 );
2026 }
2027
2028 #[test]
2036 fn test_decoder_filter_with_row_selection_and_limit() {
2037 use std::sync::atomic::{AtomicUsize, Ordering};
2038
2039 let builder =
2040 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2041 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2042
2043 let rows_filtered = Arc::new(AtomicUsize::new(0));
2044 let rows_filtered_for_predicate = Arc::clone(&rows_filtered);
2045
2046 let row_filter_a = ArrowPredicateFn::new(
2047 ProjectionMask::columns(&schema_descr, ["a"]),
2048 move |batch: RecordBatch| {
2049 rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed);
2050 let scalar_175 = Int64Array::new_scalar(175);
2051 let column = batch.column(0).as_primitive::<Int64Type>();
2052 gt(column, &scalar_175)
2053 },
2054 );
2055
2056 let mut decoder = builder
2057 .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
2058 .with_row_selection(RowSelection::from(vec![
2059 RowSelector::skip(150),
2060 RowSelector::select(50),
2061 ]))
2062 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
2063 .with_batch_size(10)
2064 .with_limit(10)
2065 .build()
2066 .unwrap();
2067
2068 let ranges = expect_needs_data(decoder.try_decode());
2069 push_ranges_to_decoder(&mut decoder, ranges);
2070
2071 let batch = expect_data(decoder.try_decode());
2072 let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap();
2073 assert_eq!(batch, expected);
2074
2075 expect_finished(decoder.try_decode());
2076
2077 assert!(
2078 rows_filtered.load(Ordering::Relaxed) < 50,
2079 "predicate should short-circuit before consuming all 50 rows from the explicit row selection"
2080 );
2081 }
2082
2083 #[test]
2084 fn test_decoder_offset_limit() {
2085 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2086 .unwrap()
2087 .with_offset(225)
2089 .with_limit(20)
2091 .build()
2092 .unwrap();
2093
2094 let ranges = expect_needs_data(decoder.try_decode());
2098 push_ranges_to_decoder(&mut decoder, ranges);
2099
2100 let batch1 = expect_data(decoder.try_decode());
2102 let expected1 = TEST_BATCH.slice(225, 20);
2103 assert_eq!(batch1, expected1);
2104
2105 expect_finished(decoder.try_decode());
2106 }
2107
2108 #[test]
2109 fn test_decoder_try_next_reader_offset_limit() {
2110 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2111 .unwrap()
2112 .with_offset(225)
2113 .with_limit(20)
2114 .build()
2115 .unwrap();
2116
2117 let ranges = expect_needs_data(decoder.try_next_reader());
2118 push_ranges_to_decoder(&mut decoder, ranges);
2119
2120 let reader = expect_data(decoder.try_next_reader());
2121 let batches = reader
2122 .map(|batch| batch.expect("expected decoded batch"))
2123 .collect::<Vec<_>>();
2124 let output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
2125 assert_eq!(output, TEST_BATCH.slice(225, 20));
2126
2127 expect_finished(decoder.try_next_reader());
2128 }
2129
2130 #[test]
2131 fn test_decoder_row_group_selection() {
2132 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2134 .unwrap()
2135 .with_row_groups(vec![1])
2136 .build()
2137 .unwrap();
2138
2139 let ranges = expect_needs_data(decoder.try_decode());
2143 push_ranges_to_decoder(&mut decoder, ranges);
2144
2145 let batch1 = expect_data(decoder.try_decode());
2147 let expected1 = TEST_BATCH.slice(200, 200);
2148 assert_eq!(batch1, expected1);
2149
2150 expect_finished(decoder.try_decode());
2151 }
2152
2153 #[test]
2154 fn test_decoder_row_selection() {
2155 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2157 .unwrap()
2158 .with_row_selection(RowSelection::from(vec![
2159 RowSelector::skip(225), RowSelector::select(20), ]))
2162 .build()
2163 .unwrap();
2164
2165 let ranges = expect_needs_data(decoder.try_decode());
2169 push_ranges_to_decoder(&mut decoder, ranges);
2170
2171 let batch1 = expect_data(decoder.try_decode());
2173 let expected1 = TEST_BATCH.slice(225, 20);
2174 assert_eq!(batch1, expected1);
2175
2176 expect_finished(decoder.try_decode());
2177 }
2178
2179 #[test]
2180 fn test_decoder_row_group_local_selections() {
2181 let bitmap_selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(
2182 (0..45)
2183 .map(|row_idx| (25..45).contains(&row_idx))
2184 .collect::<Vec<_>>(),
2185 ));
2186 let rle_selection =
2187 RowSelection::from(vec![RowSelector::skip(190), RowSelector::select(10)]);
2188
2189 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2190 .unwrap()
2191 .with_row_group_selections(vec![
2192 RowGroupSelection::new(1, Some(bitmap_selection)),
2193 RowGroupSelection::new(0, Some(rle_selection)),
2194 ])
2195 .build()
2196 .unwrap();
2197 prefetch_test_file(&mut decoder);
2198
2199 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(225, 20));
2203 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(190, 10));
2204 expect_finished(decoder.try_decode());
2205 }
2206
2207 #[test]
2208 fn test_row_group_local_selections_respect_offset_and_limit() {
2209 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2210 .unwrap()
2211 .with_row_group_selections(vec![
2212 RowGroupSelection::new(1, None),
2213 RowGroupSelection::new(0, None),
2214 ])
2215 .with_offset(195)
2216 .with_limit(10)
2217 .build()
2218 .unwrap();
2219 prefetch_test_file(&mut decoder);
2220
2221 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(395, 5));
2222 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 5));
2223 expect_finished(decoder.try_decode());
2224 }
2225
2226 #[test]
2227 fn test_row_group_local_selections_allow_duplicate_row_groups() {
2228 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2229 .unwrap()
2230 .with_row_group_selections(vec![
2231 RowGroupSelection::new(0, None),
2232 RowGroupSelection::new(0, None),
2233 ])
2234 .build()
2235 .unwrap();
2236 prefetch_test_file(&mut decoder);
2237
2238 let expected = TEST_BATCH.slice(0, 200);
2239 assert_eq!(expect_data(decoder.try_decode()), expected);
2240 assert_eq!(expect_data(decoder.try_decode()), expected);
2241 expect_finished(decoder.try_decode());
2242 }
2243
2244 #[test]
2245 fn test_short_row_group_local_selection_skips_trailing_rows() {
2246 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2247 .unwrap()
2248 .with_row_group_selections(vec![RowGroupSelection::new(
2249 1,
2250 Some(RowSelection::from(vec![
2251 RowSelector::skip(5),
2252 RowSelector::select(3),
2253 ])),
2254 )])
2255 .build()
2256 .unwrap();
2257 prefetch_test_file(&mut decoder);
2258
2259 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(205, 3));
2260 expect_finished(decoder.try_decode());
2261 }
2262
2263 #[test]
2264 fn test_empty_row_group_local_selections_read_nothing() {
2265 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2266 .unwrap()
2267 .with_row_group_selections(vec![])
2268 .build()
2269 .unwrap();
2270
2271 expect_finished(decoder.try_decode());
2272 }
2273
2274 fn row_group_local_selections_with_filter() -> ParquetPushDecoderBuilder {
2283 let builder =
2284 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2285 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2286
2287 let row_filter_a = ArrowPredicateFn::new(
2289 ProjectionMask::columns(&schema_descr, ["a"]),
2290 |batch: RecordBatch| {
2291 let scalar_195 = Int64Array::new_scalar(195);
2292 let column = batch.column(0).as_primitive::<Int64Type>();
2293 gt(column, &scalar_195)
2294 },
2295 );
2296
2297 builder
2298 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
2299 .with_row_group_selections(vec![
2300 RowGroupSelection::new(
2301 1,
2302 Some(RowSelection::from(vec![
2303 RowSelector::skip(10),
2304 RowSelector::select(10),
2305 ])),
2306 ),
2307 RowGroupSelection::new(
2308 0,
2309 Some(RowSelection::from(vec![
2310 RowSelector::skip(190),
2311 RowSelector::select(10),
2312 ])),
2313 ),
2314 ])
2315 }
2316
2317 #[test]
2318 fn test_row_group_local_selections_with_row_filter() {
2319 let mut decoder = row_group_local_selections_with_filter().build().unwrap();
2320 prefetch_test_file(&mut decoder);
2321
2322 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(210, 10));
2323 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(196, 4));
2324 expect_finished(decoder.try_decode());
2325 }
2326
2327 #[test]
2331 fn test_into_builder_preserves_local_selections_with_row_filter() {
2332 let mut decoder = row_group_local_selections_with_filter().build().unwrap();
2333 prefetch_test_file(&mut decoder);
2334
2335 let reader1 = expect_data(decoder.try_next_reader());
2336 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2337 let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2338 assert_eq!(batch1, TEST_BATCH.slice(210, 10));
2339
2340 assert!(decoder.is_at_row_group_boundary());
2342 assert_eq!(decoder.row_groups_remaining(), 1);
2343 let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2344
2345 let reader0 = expect_data(decoder.try_next_reader());
2346 let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2347 let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2348 assert_eq!(batch0, TEST_BATCH.slice(196, 4));
2349 expect_finished(decoder.try_next_reader());
2350 }
2351
2352 #[test]
2353 fn test_row_group_local_selection_replaces_previous_configuration() {
2354 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2355 .unwrap()
2356 .with_row_group_selections(vec![RowGroupSelection::new(0, None)])
2357 .with_row_group_selections(vec![RowGroupSelection::new(1, None)])
2358 .build()
2359 .unwrap();
2360 prefetch_test_file(&mut decoder);
2361
2362 assert_eq!(
2365 expect_data(decoder.try_decode()),
2366 TEST_BATCH.slice(200, 200)
2367 );
2368 expect_finished(decoder.try_decode());
2369 }
2370
2371 #[test]
2372 fn test_legacy_row_groups_and_row_selection_remain_composable() {
2373 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2374 .unwrap()
2375 .with_row_groups(vec![1])
2376 .with_row_selection(RowSelection::from(vec![
2377 RowSelector::skip(25),
2378 RowSelector::select(20),
2379 ]))
2380 .build()
2381 .unwrap();
2382 prefetch_test_file(&mut decoder);
2383
2384 assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(225, 20));
2385 expect_finished(decoder.try_decode());
2386 }
2387
2388 #[test]
2389 fn test_row_group_local_selection_is_mutually_exclusive_with_legacy_configuration() {
2390 let metadata = test_file_parquet_metadata();
2391 let new_builder =
2392 || ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata)).unwrap();
2393 let local_selection = || vec![RowGroupSelection::new(0, None)];
2394 let global_selection = || RowSelection::from(vec![RowSelector::select(1)]);
2395 let assert_conflict = |builder: ParquetPushDecoderBuilder| {
2396 let error = builder.build().unwrap_err();
2397 assert_eq!(
2398 error.to_string(),
2399 "Parquet error: with_row_group_selections cannot be combined with with_row_groups or with_row_selection"
2400 );
2401 };
2402
2403 assert_conflict(
2404 new_builder()
2405 .with_row_groups(vec![0])
2406 .with_row_group_selections(local_selection()),
2407 );
2408 assert_conflict(
2409 new_builder()
2410 .with_row_group_selections(local_selection())
2411 .with_row_groups(vec![0]),
2412 );
2413 assert_conflict(
2414 new_builder()
2415 .with_row_selection(global_selection())
2416 .with_row_group_selections(local_selection()),
2417 );
2418 assert_conflict(
2419 new_builder()
2420 .with_row_group_selections(local_selection())
2421 .with_row_selection(global_selection()),
2422 );
2423 }
2424
2425 #[test]
2426 fn test_row_group_local_selection_validates_row_group_and_length_at_build() {
2427 let metadata = test_file_parquet_metadata();
2428
2429 let error = ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata))
2430 .unwrap()
2431 .with_row_group_selections(vec![RowGroupSelection::new(
2432 0,
2433 Some(RowSelection::from(vec![RowSelector::select(201)])),
2434 )])
2435 .build()
2436 .unwrap_err();
2437 assert!(
2438 error.to_string().contains(
2439 "Row selection for row group 0 contains 201 rows, but the row group has 200"
2440 ),
2441 "unexpected error: {error}"
2442 );
2443
2444 let error = ParquetPushDecoderBuilder::try_new_decoder(metadata)
2445 .unwrap()
2446 .with_row_group_selections(vec![RowGroupSelection::new(2, None)])
2447 .build()
2448 .unwrap_err();
2449 assert!(
2450 error
2451 .to_string()
2452 .contains("Row group index 2 out of bounds for file with 2 row groups"),
2453 "unexpected error: {error}"
2454 );
2455 }
2456
2457 #[test]
2461 fn test_peek_next_row_group_basic() {
2462 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2463 .unwrap()
2464 .build()
2465 .unwrap();
2466
2467 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(0));
2470 assert!(decoder.is_at_row_group_boundary());
2471
2472 let ranges = expect_needs_data(decoder.try_next_reader());
2473 push_ranges_to_decoder(&mut decoder, ranges);
2474 let reader = expect_data(decoder.try_next_reader());
2475 assert!(decoder.is_at_row_group_boundary());
2479 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2480
2481 for batch in reader {
2483 let _ = batch.unwrap();
2484 }
2485 let ranges = expect_needs_data(decoder.try_next_reader());
2486 push_ranges_to_decoder(&mut decoder, ranges);
2487 let reader = expect_data(decoder.try_next_reader());
2488 for batch in reader {
2489 let _ = batch.unwrap();
2490 }
2491
2492 assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2494 }
2495
2496 #[test]
2500 fn test_peek_next_row_group_respects_with_row_groups() {
2501 let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2502 .unwrap()
2503 .with_row_groups(vec![1])
2504 .build()
2505 .unwrap();
2506
2507 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2508 }
2509
2510 #[test]
2515 fn test_peek_next_row_group_skips_empty_selection() {
2516 let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2519 .unwrap()
2520 .with_row_selection(RowSelection::from(vec![
2521 RowSelector::skip(250),
2522 RowSelector::select(100),
2523 ]))
2524 .build()
2525 .unwrap();
2526
2527 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2528 }
2529
2530 #[test]
2532 fn test_peek_next_row_group_finished() {
2533 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2534 .unwrap()
2535 .with_row_groups(vec![])
2536 .build()
2537 .unwrap();
2538
2539 expect_finished(decoder.try_next_reader());
2541 assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2542 }
2543
2544 #[test]
2549 fn test_peek_next_row_group_skips_for_budget() {
2550 let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2553 .unwrap()
2554 .with_offset(200)
2555 .build()
2556 .unwrap();
2557
2558 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2559 }
2560
2561 #[test]
2565 fn test_peek_next_row_group_budget_drains_all() {
2566 let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2569 .unwrap()
2570 .with_offset(500)
2571 .build()
2572 .unwrap();
2573
2574 assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2575 }
2576
2577 #[test]
2583 fn test_peek_next_row_group_during_decoding_row_group() {
2584 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2588 .unwrap()
2589 .with_batch_size(100)
2590 .build()
2591 .unwrap();
2592 decoder
2593 .push_range(test_file_range(), TEST_FILE_DATA.clone())
2594 .unwrap();
2595
2596 let _batch0 = expect_data(decoder.try_decode());
2601 assert!(!decoder.is_at_row_group_boundary());
2602 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2603
2604 let _batch1 = expect_data(decoder.try_decode());
2606 assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2607
2608 let _batch2 = expect_data(decoder.try_decode());
2610 assert!(!decoder.is_at_row_group_boundary());
2611 assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2612
2613 let _batch3 = expect_data(decoder.try_decode());
2614 expect_finished(decoder.try_decode());
2615 }
2616
2617 #[test]
2622 fn test_peek_next_row_group_does_not_mutate_state() {
2623 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2624 .unwrap()
2625 .build()
2626 .unwrap();
2627
2628 for expected_rg in [0usize, 1usize] {
2630 assert!(decoder.is_at_row_group_boundary());
2631
2632 let first_peek = decoder.peek_next_row_group().unwrap();
2635 let second_peek = decoder.peek_next_row_group().unwrap();
2636 let third_peek = decoder.peek_next_row_group().unwrap();
2637 assert_eq!(first_peek, Some(expected_rg));
2638 assert_eq!(first_peek, second_peek);
2639 assert_eq!(second_peek, third_peek);
2640
2641 let ranges = expect_needs_data(decoder.try_next_reader());
2644 push_ranges_to_decoder(&mut decoder, ranges);
2645 let reader = expect_data(decoder.try_next_reader());
2646 for batch in reader {
2647 let _ = batch.unwrap();
2648 }
2649 }
2650
2651 assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2654 expect_finished(decoder.try_next_reader());
2655 }
2656
2657 #[test]
2665 fn test_into_builder_installs_filter_between_row_groups() {
2666 let schema_descr = test_file_parquet_metadata()
2667 .file_metadata()
2668 .schema_descr_ptr();
2669 let mut decoder = prefetched_decoder(1024);
2670
2671 let reader0 = expect_data(decoder.try_next_reader());
2673 let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2674 let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2675 assert_eq!(batch0, TEST_BATCH.slice(0, 200));
2676
2677 assert!(decoder.is_at_row_group_boundary());
2679 assert_eq!(decoder.row_groups_remaining(), 1);
2680 let filter =
2681 ArrowPredicateFn::new(ProjectionMask::columns(&schema_descr, ["a"]), |batch| {
2682 gt(batch.column(0), &Int64Array::new_scalar(250))
2683 });
2684 let mut decoder = decoder
2685 .into_builder()
2686 .unwrap()
2687 .with_row_filter(RowFilter::new(vec![Box::new(filter)]))
2688 .build()
2689 .unwrap();
2690
2691 let reader1 = expect_data(decoder.try_next_reader());
2696 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2697 let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2698 assert_eq!(batch1, TEST_BATCH.slice(251, 149));
2699 expect_finished(decoder.try_next_reader());
2700 }
2701
2702 #[test]
2706 fn test_into_builder_rejected_mid_row_group() {
2707 let mut decoder = prefetched_decoder(50);
2708
2709 expect_data(decoder.try_decode());
2712 assert!(!decoder.is_at_row_group_boundary());
2713
2714 let err = decoder.into_builder().unwrap_err();
2715 let err_msg = format!("{err}");
2716 assert!(
2717 err_msg.contains("is_at_row_group_boundary"),
2718 "unexpected error: {err_msg}"
2719 );
2720 }
2721
2722 #[test]
2724 fn test_into_builder_rejected_on_finished_decoder() {
2725 let mut decoder = prefetched_decoder(1024);
2726 expect_data(decoder.try_decode());
2727 expect_data(decoder.try_decode());
2728 expect_finished(decoder.try_decode());
2729 assert!(!decoder.is_at_row_group_boundary());
2730
2731 let err = decoder.into_builder().unwrap_err();
2732 assert!(
2733 format!("{err}").contains("finished"),
2734 "unexpected error: {err}"
2735 );
2736 }
2737
2738 #[test]
2744 fn test_into_builder_allowed_while_iterating_handed_off_reader() {
2745 let mut decoder = prefetched_decoder(1024);
2746
2747 let reader0 = expect_data(decoder.try_next_reader());
2748 assert!(decoder.is_at_row_group_boundary());
2751 let _builder = decoder.into_builder().unwrap();
2754 let batches: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2755 let batch0 = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
2756 assert_eq!(batch0, TEST_BATCH.slice(0, 200));
2757 }
2758
2759 #[test]
2763 fn test_into_builder_resumes_remaining_budget() {
2764 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2767 .unwrap()
2768 .with_batch_size(1024)
2769 .with_limit(250)
2770 .build()
2771 .unwrap();
2772 prefetch_test_file(&mut decoder);
2773
2774 let reader0 = expect_data(decoder.try_next_reader());
2776 let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2777 let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2778 assert_eq!(batch0, TEST_BATCH.slice(0, 200));
2779
2780 assert!(decoder.is_at_row_group_boundary());
2784 let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2785
2786 let reader1 = expect_data(decoder.try_next_reader());
2787 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2788 let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2789 assert_eq!(batch1, TEST_BATCH.slice(200, 50));
2791 expect_finished(decoder.try_next_reader());
2792 }
2793
2794 #[test]
2798 fn test_into_builder_preserves_buffered_bytes() {
2799 let mut decoder = prefetched_decoder(1024);
2800 assert_eq!(decoder.buffered_bytes(), test_file_len());
2801
2802 let reader0 = expect_data(decoder.try_next_reader());
2804 let _: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2805 let buffered = decoder.buffered_bytes();
2807 assert!(buffered > 0);
2808
2809 let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2812 assert!(decoder.buffered_bytes() < buffered);
2813 assert_eq!(decoder.buffered_bytes(), test_file_row_group_bytes(1));
2814
2815 let reader1 = expect_data(decoder.try_next_reader());
2818 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2819 let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2820 assert_eq!(batch1, TEST_BATCH.slice(200, 200));
2821 expect_finished(decoder.try_next_reader());
2822 }
2823
2824 #[test]
2827 fn test_into_builder_releases_bytes_of_skipped_row_groups() {
2828 let builder =
2829 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2830 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2831 let mut decoder = builder
2832 .with_row_groups(vec![1])
2833 .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"]))
2834 .build()
2835 .unwrap();
2836 prefetch_test_file(&mut decoder);
2837 assert!(decoder.is_at_row_group_boundary());
2838
2839 let mut decoder = decoder
2841 .into_builder()
2842 .unwrap()
2843 .with_row_groups(vec![0])
2844 .with_projection(ProjectionMask::all())
2845 .build()
2846 .unwrap();
2847 assert_eq!(decoder.buffered_bytes(), test_file_row_group_bytes(0));
2848
2849 let reader = expect_data(decoder.try_next_reader());
2850 let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
2851 let batch = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
2852 assert_eq!(batch, TEST_BATCH.slice(0, 200));
2853 expect_finished(decoder.try_next_reader());
2854 }
2855
2856 #[test]
2859 fn test_into_builder_keeps_bytes_of_predicate_columns() {
2860 let builder =
2861 ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2862 let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2863 let row_filter_a = ArrowPredicateFn::new(
2864 ProjectionMask::columns(&schema_descr, ["a"]),
2865 |batch: RecordBatch| {
2866 let scalar_250 = Int64Array::new_scalar(250);
2867 let column = batch.column(0).as_primitive::<Int64Type>();
2868 gt(column, &scalar_250)
2869 },
2870 );
2871 let mut decoder = builder
2872 .with_row_groups(vec![1])
2873 .with_projection(ProjectionMask::columns(&schema_descr, ["b"]))
2874 .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
2875 .build()
2876 .unwrap();
2877 prefetch_test_file(&mut decoder);
2878
2879 let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2882 let column_bytes = |column_idx: usize| {
2883 let range =
2884 column_chunk_range(test_file_parquet_metadata().row_group(1).column(column_idx));
2885 range.end - range.start
2886 };
2887 assert_eq!(decoder.buffered_bytes(), column_bytes(0) + column_bytes(1));
2888
2889 let reader = expect_data(decoder.try_next_reader());
2890 let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
2891 let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
2892 assert_eq!(batch, TEST_BATCH.slice(251, 149).project(&[1]).unwrap());
2893 expect_finished(decoder.try_next_reader());
2894 }
2895
2896 #[test]
2897 fn test_into_builder_preserves_remaining_row_group_local_selections() {
2898 let bitmap_selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(
2899 (0..45)
2900 .map(|row_idx| (25..45).contains(&row_idx))
2901 .collect::<Vec<_>>(),
2902 ));
2903 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2904 .unwrap()
2905 .with_row_group_selections(vec![
2906 RowGroupSelection::new(
2907 0,
2908 Some(RowSelection::from(vec![
2909 RowSelector::skip(190),
2910 RowSelector::select(10),
2911 ])),
2912 ),
2913 RowGroupSelection::new(1, Some(bitmap_selection)),
2914 ])
2915 .build()
2916 .unwrap();
2917 prefetch_test_file(&mut decoder);
2918
2919 let reader0 = expect_data(decoder.try_next_reader());
2920 let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2921 let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2922 assert_eq!(batch0, TEST_BATCH.slice(190, 10));
2923
2924 assert!(decoder.is_at_row_group_boundary());
2926 assert_eq!(decoder.row_groups_remaining(), 1);
2927 let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2928
2929 let reader1 = expect_data(decoder.try_next_reader());
2930 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2931 let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2932 assert_eq!(batch1, TEST_BATCH.slice(225, 20));
2933 expect_finished(decoder.try_next_reader());
2934 }
2935
2936 #[test]
2944 fn test_into_builder_expand_projection_requests_new_bytes() {
2945 let metadata = test_file_parquet_metadata();
2946 let schema_descr = metadata.file_metadata().schema_descr_ptr();
2947
2948 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata)
2949 .unwrap()
2950 .with_batch_size(1024)
2951 .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
2952 .build()
2953 .unwrap();
2954
2955 let ranges_rg0 = expect_needs_data(decoder.try_next_reader());
2958 assert_eq!(ranges_rg0, vec![4..1860]);
2959 push_ranges_to_decoder(&mut decoder, ranges_rg0);
2960
2961 let reader0 = expect_data(decoder.try_next_reader());
2962 let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2963 let batch0 = concat_batches(&batches0[0].schema(), &batches0).unwrap();
2964 assert_eq!(batch0, TEST_BATCH.slice(0, 200).project(&[0]).unwrap());
2965
2966 assert!(decoder.is_at_row_group_boundary());
2968 let mut decoder = decoder
2969 .into_builder()
2970 .unwrap()
2971 .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b", "c"]))
2972 .build()
2973 .unwrap();
2974
2975 let ranges_rg1 = expect_needs_data(decoder.try_next_reader());
2978 assert_eq!(ranges_rg1, vec![11062..12918, 12918..14774, 14774..22230]);
2979 push_ranges_to_decoder(&mut decoder, ranges_rg1);
2980
2981 let reader1 = expect_data(decoder.try_next_reader());
2982 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2983 let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2984 assert_eq!(batch1, TEST_BATCH.slice(200, 200));
2985 expect_finished(decoder.try_next_reader());
2986 }
2987
2988 #[test]
2994 fn test_into_builder_narrow_projection_requests_fewer_bytes() {
2995 let metadata = test_file_parquet_metadata();
2996 let schema_descr = metadata.file_metadata().schema_descr_ptr();
2997
2998 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata)
2999 .unwrap()
3000 .with_batch_size(1024)
3001 .build()
3002 .unwrap();
3003
3004 let ranges_rg0 = expect_needs_data(decoder.try_next_reader());
3006 assert_eq!(ranges_rg0, vec![4..1860, 1860..3716, 3716..11062]);
3007 push_ranges_to_decoder(&mut decoder, ranges_rg0);
3008
3009 let reader0 = expect_data(decoder.try_next_reader());
3010 let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
3011 let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
3012 assert_eq!(batch0, TEST_BATCH.slice(0, 200));
3013
3014 assert!(decoder.is_at_row_group_boundary());
3016 let mut decoder = decoder
3017 .into_builder()
3018 .unwrap()
3019 .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
3020 .build()
3021 .unwrap();
3022
3023 let ranges_rg1 = expect_needs_data(decoder.try_next_reader());
3025 assert_eq!(ranges_rg1, vec![11062..12918]);
3026 push_ranges_to_decoder(&mut decoder, ranges_rg1);
3027
3028 let reader1 = expect_data(decoder.try_next_reader());
3029 let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
3030 let batch1 = concat_batches(&batches1[0].schema(), &batches1).unwrap();
3031 assert_eq!(batch1, TEST_BATCH.slice(200, 200).project(&[0]).unwrap());
3032 expect_finished(decoder.try_next_reader());
3033 }
3034
3035 static TEST_BATCH: LazyLock<RecordBatch> = LazyLock::new(|| {
3039 let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400));
3040 let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800));
3041 let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| {
3042 if i % 2 == 0 {
3043 format!("string_{i}")
3044 } else {
3045 format!("A string larger than 12 bytes and thus not inlined {i}")
3046 }
3047 })));
3048
3049 RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
3050 });
3051
3052 static TEST_FILE_DATA: LazyLock<Bytes> = LazyLock::new(|| {
3078 let input_batch = &TEST_BATCH;
3079 let mut output = Vec::new();
3080
3081 let writer_options = WriterProperties::builder()
3082 .set_max_row_group_row_count(Some(200))
3083 .set_data_page_row_count_limit(100)
3084 .build();
3085 let mut writer =
3086 ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap();
3087
3088 let mut row_remain = input_batch.num_rows();
3091 while row_remain > 0 {
3092 let chunk_size = row_remain.min(50);
3093 let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size);
3094 writer.write(&chunk).unwrap();
3095 row_remain -= chunk_size;
3096 }
3097 writer.close().unwrap();
3098 Bytes::from(output)
3099 });
3100
3101 fn test_file_len() -> u64 {
3103 TEST_FILE_DATA.len() as u64
3104 }
3105
3106 fn test_file_range() -> Range<u64> {
3108 0..test_file_len()
3109 }
3110
3111 pub fn test_file_slice(range: Range<u64>) -> Bytes {
3113 let start: usize = range.start.try_into().unwrap();
3114 let end: usize = range.end.try_into().unwrap();
3115 TEST_FILE_DATA.slice(start..end)
3116 }
3117
3118 pub fn test_file_parquet_metadata() -> Arc<ParquetMetaData> {
3120 let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap();
3121 push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]);
3122 let metadata = metadata_decoder.try_decode().unwrap();
3123 let DecodeResult::Data(metadata) = metadata else {
3124 panic!("Expected metadata to be decoded successfully");
3125 };
3126 Arc::new(metadata)
3127 }
3128
3129 #[test]
3132 fn test_decoder_offset_index_page_past_column_chunk() {
3133 let metadata = Arc::unwrap_or_clone(test_file_parquet_metadata_with_offset_index());
3134 let num_columns = metadata.row_group(0).num_columns();
3135 let mut page_index = PageIndexBuilder::new(metadata.num_row_groups(), num_columns);
3136 for row_group in 0..metadata.num_row_groups() {
3137 let row_group_index = metadata.page_index_for_row_group(row_group);
3138 for column in 0..num_columns {
3139 let mut offset_index = row_group_index.offset_index(column).unwrap().clone();
3140 if (row_group, column) == (0, 0) {
3141 let last = offset_index.page_locations.last_mut().unwrap();
3142 last.compressed_page_size += 1;
3143 }
3144 page_index.put_offset_index(offset_index, row_group, column);
3145 }
3146 }
3147 let metadata = metadata
3148 .into_builder()
3149 .set_page_index(Some(Arc::new(page_index.build())))
3150 .build();
3151
3152 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(Arc::new(metadata))
3153 .unwrap()
3154 .with_row_groups(vec![0])
3155 .build()
3156 .unwrap();
3157 let ranges = expect_needs_data(decoder.try_decode());
3158 push_ranges_to_decoder(&mut decoder, ranges);
3159 let err = loop {
3160 match decoder.try_decode() {
3161 Ok(DecodeResult::Data(_)) => {}
3162 Ok(other) => panic!("expected an error, got {other:?}"),
3163 Err(e) => break e.to_string(),
3164 }
3165 };
3166 assert!(err.contains("has 126 bytes, expected 127"), "{err}");
3167 }
3168
3169 fn test_file_parquet_metadata_with_offset_index() -> Arc<ParquetMetaData> {
3171 let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(test_file_len())
3172 .unwrap()
3173 .with_offset_index_policy(PageIndexPolicy::Required);
3174 push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]);
3175 let metadata = metadata_decoder.try_decode().unwrap();
3176 let DecodeResult::Data(metadata) = metadata else {
3177 panic!("Expected metadata to be decoded successfully");
3178 };
3179 Arc::new(metadata)
3180 }
3181
3182 fn push_ranges_to_metadata_decoder(
3184 metadata_decoder: &mut ParquetMetaDataPushDecoder,
3185 ranges: Vec<Range<u64>>,
3186 ) {
3187 let data = ranges
3188 .iter()
3189 .map(|range| test_file_slice(range.clone()))
3190 .collect::<Vec<_>>();
3191 metadata_decoder.push_ranges(ranges, data).unwrap();
3192 }
3193
3194 fn prefetch_test_file(decoder: &mut ParquetPushDecoder) {
3196 decoder
3197 .push_range(test_file_range(), TEST_FILE_DATA.clone())
3198 .unwrap();
3199 }
3200
3201 fn prefetched_decoder(batch_size: usize) -> ParquetPushDecoder {
3204 let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
3205 .unwrap()
3206 .with_batch_size(batch_size)
3207 .build()
3208 .unwrap();
3209 prefetch_test_file(&mut decoder);
3210 decoder
3211 }
3212
3213 fn test_file_row_group_bytes(row_group_idx: usize) -> u64 {
3216 test_file_parquet_metadata()
3217 .row_group(row_group_idx)
3218 .columns()
3219 .iter()
3220 .map(|column| {
3221 let range = column_chunk_range(column);
3222 range.end - range.start
3223 })
3224 .sum()
3225 }
3226
3227 fn push_ranges_to_decoder(decoder: &mut ParquetPushDecoder, ranges: Vec<Range<u64>>) {
3228 let data = ranges
3229 .iter()
3230 .map(|range| test_file_slice(range.clone()))
3231 .collect::<Vec<_>>();
3232 decoder.push_ranges(ranges, data).unwrap();
3233 }
3234
3235 fn expect_data<T: Debug>(result: Result<DecodeResult<T>, ParquetError>) -> T {
3237 match result.expect("Expected Ok(DecodeResult::Data(T))") {
3238 DecodeResult::Data(data) => data,
3239 result => panic!("Expected DecodeResult::Data, got {result:?}"),
3240 }
3241 }
3242
3243 fn expect_needs_data<T: Debug>(
3245 result: Result<DecodeResult<T>, ParquetError>,
3246 ) -> Vec<Range<u64>> {
3247 match result.expect("Expected Ok(DecodeResult::NeedsData(ranges))") {
3248 DecodeResult::NeedsData(ranges) => ranges,
3249 result => panic!("Expected DecodeResult::NeedsData, got {result:?}"),
3250 }
3251 }
3252
3253 fn expect_finished<T: Debug>(result: Result<DecodeResult<T>, ParquetError>) {
3254 match result.expect("Expected Ok(DecodeResult::Finished)") {
3255 DecodeResult::Finished => {}
3256 result => panic!("Expected DecodeResult::Finished, got {result:?}"),
3257 }
3258 }
3259}