1use std::borrow::Cow;
24use std::collections::HashMap;
25use std::sync::Arc;
26use std::sync::atomic::AtomicU64;
27
28use arrow_schema::{DataType, Field};
29use futures::{StreamExt, TryStreamExt};
30use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
31use parquet::arrow::{
32 PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder, ProjectionMask, RowNumber,
33};
34use parquet::encryption::decrypt::FileDecryptionProperties;
35
36use super::row_lineage::synthesize_row_id_column;
37use super::{
38 ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema,
39 apply_name_mapping_to_arrow_schema, find_leaf_by_field_id,
40};
41use crate::arrow::build_partition_constant;
42use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader;
43use crate::arrow::int96::coerce_int96_timestamps;
44use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder;
45use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics, ScanResult};
46use crate::encryption::StandardKeyMetadata;
47use crate::error::{Result, invalid_data};
48use crate::expr::BoundPredicate;
49use crate::expr::visitors::bloom_filter_evaluator::{
50 BloomFilterEvaluator, ColumnBloomFilter, collect_bloom_filter_field_ids,
51};
52use crate::io::{FileIO, FileMetadata, FileRead};
53use crate::metadata_columns::{
54 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER, RESERVED_COL_NAME_POS,
55 RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE,
56 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, RESERVED_FIELD_ID_PARTITION,
57 RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID, RESERVED_FIELD_ID_SPEC_ID, is_metadata_field,
58};
59use crate::runtime::Runtime;
60use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream};
61use crate::spec::{Datum, PartitionSpec, Struct};
62use crate::{Error, ErrorKind};
63
64impl ArrowReader {
65 pub fn read(self, tasks: FileScanTaskStream) -> Result<ScanResult> {
68 let concurrency_limit_data_files = self.concurrency_limit_data_files;
69 let scan_metrics = ScanMetrics::new();
70
71 let task_reader = FileScanTaskReader {
72 batch_size: self.batch_size,
73 file_io: self.file_io,
74 delete_file_loader: self
75 .delete_file_loader
76 .with_scan_metrics(scan_metrics.clone()),
77 row_group_filtering_enabled: self.row_group_filtering_enabled,
78 row_selection_enabled: self.row_selection_enabled,
79 bloom_filter_enabled: self.bloom_filter_enabled,
80 parquet_read_options: self.parquet_read_options,
81 scan_metrics: scan_metrics.clone(),
82 runtime: self.runtime,
83 };
84
85 let stream: ArrowRecordBatchStream = if concurrency_limit_data_files == 1 {
87 Box::pin(
88 tasks
89 .and_then(move |task| task_reader.clone().process(task))
90 .map_err(|err| {
91 Error::new(ErrorKind::Unexpected, "file scan task generate failed")
92 .with_source(err)
93 })
94 .try_flatten(),
95 )
96 } else {
97 Box::pin(
98 tasks
99 .map_ok(move |task| task_reader.clone().process(task))
100 .map_err(|err| {
101 Error::new(ErrorKind::Unexpected, "file scan task generate failed")
102 .with_source(err)
103 })
104 .try_buffer_unordered(concurrency_limit_data_files)
105 .try_flatten_unordered(concurrency_limit_data_files),
106 )
107 };
108
109 Ok(ScanResult::new(stream, scan_metrics))
110 }
111}
112
113const PRUNABLE_METADATA_FIELDS: &[i32] = &[
118 RESERVED_FIELD_ID_FILE,
119 RESERVED_FIELD_ID_SPEC_ID,
120 RESERVED_FIELD_ID_PARTITION,
121 RESERVED_FIELD_ID_POS,
122 RESERVED_FIELD_ID_ROW_ID,
123 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
124];
125
126#[derive(Clone)]
129struct FileScanTaskReader {
130 batch_size: Option<usize>,
131 file_io: FileIO,
132 delete_file_loader: CachingDeleteFileLoader,
133 row_group_filtering_enabled: bool,
134 row_selection_enabled: bool,
135 bloom_filter_enabled: bool,
136 parquet_read_options: ParquetReadOptions,
137 scan_metrics: ScanMetrics,
138 runtime: Runtime,
139}
140
141fn candidate_row_groups(
143 selected_row_groups: Option<&[usize]>,
144 num_row_groups: usize,
145) -> Cow<'_, [usize]> {
146 match selected_row_groups {
147 Some(indices) => Cow::Borrowed(indices),
148 None => Cow::Owned((0..num_row_groups).collect()),
149 }
150}
151
152impl FileScanTaskReader {
153 async fn process(self, task: FileScanTask) -> Result<ArrowRecordBatchStream> {
154 let should_load_page_index = (self.row_selection_enabled && task.predicate().is_some())
155 || !task.deletes().is_empty();
156 let mut parquet_read_options = self.parquet_read_options;
157 parquet_read_options.preload_page_index = should_load_page_index;
158
159 let delete_filter_rx = self
160 .delete_file_loader
161 .load_deletes(task.deletes(), task.schema_ref());
162
163 let path = task.data_file_path().to_string();
168 let file_io = self.file_io.clone();
169 let file_size_in_bytes = task.file_size_in_bytes();
170 let bytes_read = Arc::clone(self.scan_metrics.bytes_read_counter());
171 let key_metadata = task.key_metadata().map(|km| km.to_vec());
172 let opening = self
173 .runtime
174 .io()
175 .spawn(async move {
176 ArrowReader::open_parquet_file(
177 &path,
178 &file_io,
179 file_size_in_bytes,
180 parquet_read_options,
181 &bytes_read,
182 key_metadata.as_deref(),
183 )
184 .await
185 })
186 .abort_on_drop();
187 let (parquet_file_reader, arrow_metadata) = opening.await??;
188
189 let missing_field_ids = arrow_metadata
193 .schema()
194 .fields()
195 .iter()
196 .next()
197 .is_some_and(|f| f.metadata().get(PARQUET_FIELD_ID_META_KEY).is_none());
198
199 let use_position_fallback = missing_field_ids && task.name_mapping().is_none();
204
205 let project_pos = task.project_field_ids().contains(&RESERVED_FIELD_ID_POS);
206 let project_row_id = task.project_field_ids().contains(&RESERVED_FIELD_ID_ROW_ID);
207
208 let need_row_number = project_pos || (project_row_id && task.first_row_id().is_some());
213
214 let field_ids = task.project_field_ids();
215 let no_data_columns_projected = field_ids
217 .iter()
218 .all(|id| PRUNABLE_METADATA_FIELDS.contains(id));
219
220 let install_row_number = need_row_number || no_data_columns_projected;
221
222 let arrow_metadata = Self::configure_arrow_reader_metadata(
223 arrow_metadata,
224 &task,
225 missing_field_ids,
226 install_row_number,
227 )?;
228
229 let mut record_batch_stream_builder =
231 ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_file_reader, arrow_metadata);
232
233 let project_last_updated_seq = task
237 .project_field_ids()
238 .contains(&RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER);
239
240 let phys_last_updated_seq_leaf = if project_last_updated_seq {
246 find_leaf_by_field_id(
247 record_batch_stream_builder.parquet_schema(),
248 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
249 )
250 } else {
251 None
252 };
253
254 let last_updated_seq_present_by_name_only = project_last_updated_seq
257 && phys_last_updated_seq_leaf.is_none()
258 && record_batch_stream_builder
259 .schema()
260 .fields()
261 .iter()
262 .any(|f| f.name() == RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER);
263
264 let coalesce_last_updated_seq_leaf = phys_last_updated_seq_leaf
269 .filter(|_| task.first_row_id().is_some() && task.data_sequence_number().is_some());
270
271 let phys_row_id_leaf = if project_row_id {
272 find_leaf_by_field_id(
273 record_batch_stream_builder.parquet_schema(),
274 RESERVED_FIELD_ID_ROW_ID,
275 )
276 } else {
277 None
278 };
279
280 let row_id_present_by_name_only = project_row_id
286 && !use_position_fallback
287 && phys_row_id_leaf.is_none()
288 && record_batch_stream_builder
289 .schema()
290 .fields()
291 .iter()
292 .any(|f| f.name() == RESERVED_COL_NAME_ROW_ID);
293
294 let coalesce_row_id_leaf = phys_row_id_leaf.filter(|_| task.first_row_id().is_some());
297
298 let project_field_ids_without_metadata: Vec<i32> = task
300 .project_field_ids()
301 .iter()
302 .filter(|&&id| !is_metadata_field(id))
303 .copied()
304 .collect();
305
306 let mut projection_mask = ArrowReader::get_arrow_projection_mask(
312 &project_field_ids_without_metadata,
313 task.schema(),
314 record_batch_stream_builder.parquet_schema(),
315 record_batch_stream_builder.schema(),
316 use_position_fallback, )?;
318
319 if project_field_ids_without_metadata.is_empty() && install_row_number {
328 projection_mask =
329 ProjectionMask::none(record_batch_stream_builder.parquet_schema().num_columns());
330 }
331
332 for leaf in [coalesce_last_updated_seq_leaf, coalesce_row_id_leaf]
337 .into_iter()
338 .flatten()
339 {
340 let phys_mask =
341 ProjectionMask::leaves(record_batch_stream_builder.parquet_schema(), vec![leaf]);
342 projection_mask.union(&phys_mask);
343 }
344
345 record_batch_stream_builder =
346 record_batch_stream_builder.with_projection(projection_mask.clone());
347
348 let mut record_batch_transformer_builder =
352 RecordBatchTransformerBuilder::new(task.schema_ref(), task.project_field_ids());
353
354 if task.project_field_ids().contains(&RESERVED_FIELD_ID_FILE) {
356 let file_datum = Datum::string(task.data_file_path().to_string());
357 record_batch_transformer_builder =
358 record_batch_transformer_builder.with_constant(RESERVED_FIELD_ID_FILE, file_datum);
359 }
360
361 if task
362 .project_field_ids()
363 .contains(&RESERVED_FIELD_ID_SPEC_ID)
364 {
365 let partition_spec = task
366 .partition_spec()
367 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Partition spec is missing"))?;
368
369 let spec_id_datum = Datum::int(partition_spec.spec_id());
370 record_batch_transformer_builder = record_batch_transformer_builder
371 .with_constant(RESERVED_FIELD_ID_SPEC_ID, spec_id_datum);
372 }
373
374 if project_last_updated_seq {
375 record_batch_transformer_builder =
380 match (task.first_row_id(), task.data_sequence_number()) {
381 (Some(_), Some(seq)) => {
382 let datum = Datum::long(seq);
383 if coalesce_last_updated_seq_leaf.is_some() {
384 record_batch_transformer_builder.with_coalesced_last_updated_seq_column(
387 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
388 datum,
389 )
390 } else if last_updated_seq_present_by_name_only {
391 return Err(Error::new(
399 ErrorKind::FeatureUnsupported,
400 "Reading a physically-stored _last_updated_sequence_number column \
401 without an embedded field id is not supported",
402 ));
403 } else {
404 record_batch_transformer_builder.with_constant(
406 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
407 datum,
408 )
409 }
410 }
411 (None, _) => record_batch_transformer_builder.with_null_metadata_column(
413 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
414 )?,
415 (Some(_), None) => {
419 return Err(invalid_data!(
420 "Data file {} has a first_row_id but no data sequence number",
421 task.data_file_path()
422 ));
423 }
424 };
425 }
426
427 if project_row_id {
428 if task.first_row_id().is_some() && row_id_present_by_name_only {
434 return Err(Error::new(
435 ErrorKind::FeatureUnsupported,
436 "Reading a physically-stored _row_id column without an embedded field id \
437 is not supported",
438 ));
439 }
440
441 record_batch_transformer_builder =
445 record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_ROW_ID);
446 }
447
448 if let (Some(partition_spec), Some(partition_data)) =
449 (task.partition_spec(), task.partition())
450 {
451 record_batch_transformer_builder = record_batch_transformer_builder
452 .with_partition(Arc::clone(partition_spec), partition_data.clone())?;
453 }
454
455 if project_pos {
456 record_batch_transformer_builder =
457 record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_POS);
458 }
459
460 if task
463 .project_field_ids()
464 .contains(&RESERVED_FIELD_ID_PARTITION)
465 && let Some(unified_type) = task.unified_partition_type()
466 {
467 let (spec, partition_data) = match (task.partition_spec(), task.partition()) {
468 (Some(spec), Some(data)) => (Arc::clone(spec), data.clone()),
469 _ if unified_type.fields().is_empty() => {
474 (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty())
475 }
476 _ => {
477 return Err(Error::new(
478 ErrorKind::Unexpected,
479 "cannot build _partition column: unified partition type has fields \
480 but the scan task is missing its partition spec or data",
481 ));
482 }
483 };
484 let constant = build_partition_constant(unified_type, &spec, &partition_data)?;
485 record_batch_transformer_builder =
486 record_batch_transformer_builder.with_partition_constant(constant);
487 }
488
489 let mut record_batch_transformer = record_batch_transformer_builder.build();
490
491 if let Some(batch_size) = self.batch_size {
492 record_batch_stream_builder = record_batch_stream_builder.with_batch_size(batch_size);
493 }
494
495 let delete_filter = delete_filter_rx.await.unwrap()?;
496 let delete_predicate = delete_filter.build_equality_delete_predicate(&task).await?;
497
498 let final_predicate = match (task.predicate(), delete_predicate) {
503 (None, None) => None,
504 (Some(predicate), None) => Some(predicate.clone()),
505 (None, Some(ref predicate)) => Some(predicate.clone()),
506 (Some(filter_predicate), Some(delete_predicate)) => {
507 Some(filter_predicate.clone().and(delete_predicate))
508 }
509 };
510
511 let mut selected_row_group_indices = None;
527 let mut row_selection = None;
528
529 if task.start() != 0 || task.length() != 0 {
532 let byte_range_filtered_row_groups = ArrowReader::filter_row_groups_by_byte_range(
533 record_batch_stream_builder.metadata(),
534 task.start(),
535 task.length(),
536 )?;
537 selected_row_group_indices = Some(byte_range_filtered_row_groups);
538 }
539
540 if let Some(predicate) = final_predicate {
541 let predicate = predicate.rewrite_not();
548
549 let (iceberg_field_ids, field_id_map) = ArrowReader::build_field_id_set_and_map(
550 record_batch_stream_builder.parquet_schema(),
551 record_batch_stream_builder.schema(),
552 &predicate,
553 use_position_fallback,
554 )?;
555
556 let row_filter = ArrowReader::get_row_filter(
557 &predicate,
558 record_batch_stream_builder.parquet_schema(),
559 &iceberg_field_ids,
560 &field_id_map,
561 )?;
562 record_batch_stream_builder = record_batch_stream_builder.with_row_filter(row_filter);
563
564 if self.row_group_filtering_enabled {
565 let candidates = candidate_row_groups(
566 selected_row_group_indices.as_deref(),
567 record_batch_stream_builder.metadata().num_row_groups(),
568 );
569
570 selected_row_group_indices = Some(ArrowReader::get_selected_row_group_indices(
571 &predicate,
572 record_batch_stream_builder.metadata(),
573 &field_id_map,
574 task.schema(),
575 candidates.as_ref(),
576 )?);
577 }
578
579 if self.bloom_filter_enabled {
580 let candidates = candidate_row_groups(
581 selected_row_group_indices.as_deref(),
582 record_batch_stream_builder.metadata().num_row_groups(),
583 );
584
585 let bloom_predicate = predicate.clone();
587 let bloom_candidates = candidates.as_ref().to_vec();
588 let bloom_field_ids = field_id_map.clone();
589 let mut builder = record_batch_stream_builder;
590 let (builder, bloom_filtered) = self
591 .runtime
592 .io()
593 .spawn(async move {
594 let filtered = Self::filter_row_groups_by_bloom_filter(
595 &bloom_predicate,
596 &mut builder,
597 &bloom_candidates,
598 &bloom_field_ids,
599 )
600 .await;
601 (builder, filtered)
602 })
603 .abort_on_drop()
604 .await?;
605 record_batch_stream_builder = builder;
606 let bloom_filtered = bloom_filtered?;
607
608 if bloom_filtered.len() < candidates.len() {
609 selected_row_group_indices = Some(bloom_filtered);
610 }
611 }
612
613 if self.row_selection_enabled {
614 row_selection = ArrowReader::get_row_selection_for_filter_predicate(
615 &predicate,
616 record_batch_stream_builder.metadata(),
617 &selected_row_group_indices,
618 &field_id_map,
619 task.schema(),
620 )?;
621 }
622 }
623
624 let positional_delete_indexes = delete_filter.get_delete_vector(&task);
625
626 if let Some(positional_delete_indexes) = positional_delete_indexes {
627 let delete_row_selection = {
628 let positional_delete_indexes = positional_delete_indexes.lock().unwrap();
629
630 ArrowReader::build_deletes_row_selection(
631 record_batch_stream_builder.metadata().row_groups(),
632 &selected_row_group_indices,
633 &positional_delete_indexes,
634 )
635 }?;
636
637 row_selection = match row_selection {
640 None => Some(delete_row_selection),
641 Some(filter_row_selection) => {
642 Some(filter_row_selection.intersection(&delete_row_selection))
643 }
644 };
645 }
646
647 if let Some(row_selection) = row_selection {
648 record_batch_stream_builder =
649 record_batch_stream_builder.with_row_selection(row_selection);
650 }
651
652 if let Some(selected_row_group_indices) = selected_row_group_indices {
653 record_batch_stream_builder =
654 record_batch_stream_builder.with_row_groups(selected_row_group_indices);
655 }
656
657 let first_row_id = task.first_row_id();
662 let record_batch_stream = record_batch_stream_builder.build()?.map(move |batch| {
663 let mut batch = batch.map_err(|err| -> Error { err.into() })?;
664 if project_row_id {
665 batch = synthesize_row_id_column(batch, first_row_id)?;
666 }
667 record_batch_transformer.process_record_batch(batch)
669 });
670
671 Ok(Box::pin(record_batch_stream) as ArrowRecordBatchStream)
672 }
673
674 fn configure_arrow_reader_metadata(
677 arrow_metadata: ArrowReaderMetadata,
678 task: &FileScanTask,
679 missing_field_ids: bool,
680 install_row_number: bool,
681 ) -> Result<ArrowReaderMetadata> {
682 let arrow_schema = if missing_field_ids {
691 let schema = if let Some(name_mapping) = task.name_mapping() {
692 apply_name_mapping_to_arrow_schema(
693 Arc::clone(arrow_metadata.schema()),
694 name_mapping,
695 )?
696 } else {
697 add_fallback_field_ids_to_arrow_schema(arrow_metadata.schema())
698 };
699 coerce_int96_timestamps(&schema, task.schema()).unwrap_or(schema)
703 } else if let Some(coerced) =
704 coerce_int96_timestamps(arrow_metadata.schema(), task.schema())
705 {
706 coerced
707 } else if install_row_number {
708 Arc::clone(arrow_metadata.schema())
709 } else {
710 return Ok(arrow_metadata);
711 };
712
713 let mut options = ArrowReaderOptions::new().with_schema(Arc::clone(&arrow_schema));
714 if install_row_number {
715 let row_number_field = Arc::new(
716 Field::new(RESERVED_COL_NAME_POS, DataType::Int64, false)
717 .with_metadata(HashMap::from([(
718 PARQUET_FIELD_ID_META_KEY.to_string(),
719 RESERVED_FIELD_ID_POS.to_string(),
720 )]))
721 .with_extension_type(RowNumber),
722 );
723 options = options.with_virtual_columns(vec![row_number_field])?;
724 }
725
726 ArrowReaderMetadata::try_new(Arc::clone(arrow_metadata.metadata()), options).map_err(|e| {
727 Error::new(
728 ErrorKind::Unexpected,
729 format!(
730 "Failed to create ArrowReaderMetadata with the configured reader options \
731 (missing_field_ids: {missing_field_ids}, \
732 install_row_number: {install_row_number}, schema: {arrow_schema})"
733 ),
734 )
735 .with_source(e)
736 })
737 }
738
739 async fn filter_row_groups_by_bloom_filter(
742 predicate: &BoundPredicate,
743 builder: &mut ParquetRecordBatchStreamBuilder<ArrowFileReader>,
744 candidate_row_groups: &[usize],
745 field_id_map: &HashMap<i32, usize>,
746 ) -> Result<Vec<usize>> {
747 let bloom_filter_field_ids: Vec<i32> = collect_bloom_filter_field_ids(predicate)?
750 .into_iter()
751 .filter(|id| field_id_map.contains_key(id))
752 .collect();
753
754 if bloom_filter_field_ids.is_empty() {
755 return Ok(candidate_row_groups.to_vec());
756 }
757
758 let mut result = Vec::with_capacity(candidate_row_groups.len());
759
760 for &rg_idx in candidate_row_groups {
761 let mut bloom_filters: HashMap<i32, ColumnBloomFilter> = HashMap::new();
762
763 for &field_id in &bloom_filter_field_ids {
764 let col_idx = field_id_map[&field_id];
765 let col_meta = builder.metadata().row_group(rg_idx).column(col_idx);
766
767 if col_meta.bloom_filter_offset().is_none() {
769 continue;
770 }
771
772 let physical_type = col_meta.column_type();
773 let type_length = col_meta.column_descr().type_length();
774
775 match builder
776 .get_row_group_column_bloom_filter(rg_idx, col_idx)
777 .await
778 {
779 Ok(Some(sbbf)) => {
780 bloom_filters.insert(
781 field_id,
782 ColumnBloomFilter::new(sbbf, physical_type, type_length),
783 );
784 }
785 Ok(None) => {}
786 Err(e) => {
787 tracing::debug!(
790 "Bloom filter for field {field_id} in row group {rg_idx} could not be read: {e}"
791 );
792 }
793 }
794 }
795
796 match BloomFilterEvaluator::eval(predicate, &bloom_filters) {
797 Ok(true) => result.push(rg_idx),
798 Ok(false) => { }
799 Err(e) => {
800 tracing::debug!(
801 "Bloom filter evaluation failed for row group {rg_idx}, including it: {e}"
802 );
803 result.push(rg_idx);
804 }
805 }
806 }
807
808 Ok(result)
809 }
810}
811
812impl ArrowReader {
813 pub(crate) async fn open_parquet_file(
816 data_file_path: &str,
817 file_io: &FileIO,
818 file_size_in_bytes: u64,
819 parquet_read_options: ParquetReadOptions,
820 bytes_read: &Arc<AtomicU64>,
821 key_metadata: Option<&[u8]>,
822 ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> {
823 let parquet_file = file_io.new_input(data_file_path)?;
824 let counting_reader =
825 CountingFileRead::new(parquet_file.reader().await?, Arc::clone(bytes_read));
826 Self::build_parquet_reader(
827 Box::new(counting_reader),
828 file_size_in_bytes,
829 parquet_read_options,
830 key_metadata,
831 )
832 .await
833 }
834
835 async fn build_parquet_reader(
836 parquet_reader: Box<dyn FileRead>,
837 file_size_in_bytes: u64,
838 parquet_read_options: ParquetReadOptions,
839 key_metadata: Option<&[u8]>,
840 ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> {
841 let mut reader = ArrowFileReader::new(
842 FileMetadata {
843 size: file_size_in_bytes,
844 },
845 parquet_reader,
846 )
847 .with_parquet_read_options(parquet_read_options);
848
849 let arrow_reader_options = Self::build_arrow_reader_options(key_metadata)?;
850
851 let arrow_metadata = ArrowReaderMetadata::load_async(&mut reader, arrow_reader_options)
852 .await
853 .map_err(|e| {
854 Error::new(ErrorKind::Unexpected, "Failed to load Parquet metadata").with_source(e)
855 })?;
856
857 Ok((reader, arrow_metadata))
858 }
859
860 fn build_arrow_reader_options(key_metadata: Option<&[u8]>) -> Result<ArrowReaderOptions> {
863 match key_metadata {
864 Some(km) => {
865 let standard_key_metadata = StandardKeyMetadata::decode(km)?;
866 let mut builder = FileDecryptionProperties::builder(
867 standard_key_metadata.encryption_key().as_bytes().to_vec(),
868 );
869 if let Some(aad) = standard_key_metadata.aad_prefix() {
870 builder = builder.with_aad_prefix(aad.to_vec());
871 }
872 let decryption_properties = builder.build().map_err(|e| {
873 Error::new(
874 ErrorKind::Unexpected,
875 "Failed to build Parquet file decryption properties",
876 )
877 .with_source(e)
878 })?;
879 Ok(
880 ArrowReaderOptions::new()
881 .with_file_decryption_properties(decryption_properties),
882 )
883 }
884 None => Ok(ArrowReaderOptions::default()),
885 }
886 }
887}
888
889#[cfg(test)]
890mod tests {
891 use std::collections::HashMap;
892 use std::fs::File;
893 use std::sync::Arc;
894
895 use arrow_array::cast::AsArray;
896 use arrow_array::{Array, ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray};
897 use arrow_cast::cast;
898 use arrow_schema::{DataType, Field, Schema as ArrowSchema};
899 use futures::TryStreamExt;
900 use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
901 use parquet::basic::Compression;
902 use parquet::file::properties::WriterProperties;
903 use tempfile::TempDir;
904
905 use crate::Runtime;
906 use crate::arrow::ArrowReaderBuilder;
907 use crate::arrow::test_utils::write_encrypted_parquet;
908 use crate::io::FileIO;
909 use crate::metadata_columns::{
910 RESERVED_COL_NAME_POS, RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE,
911 RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID,
912 };
913 use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream};
914 use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, SchemaRef, Type};
915
916 const UNIX_EPOCH_JULIAN: i64 = 2_440_588;
919 const MICROS_PER_DAY: i64 = 86_400_000_000;
920 const INT96_TEST_NANOS_WITHIN_DAY: u64 = 43_200_000_000_000;
922 const INT96_TEST_JULIAN_DAY: u32 = 2_953_529;
923
924 fn make_int96_test_value() -> (parquet::data_type::Int96, i64) {
925 let mut val = parquet::data_type::Int96::new();
926 val.set_data(
927 (INT96_TEST_NANOS_WITHIN_DAY & 0xFFFFFFFF) as u32,
928 (INT96_TEST_NANOS_WITHIN_DAY >> 32) as u32,
929 INT96_TEST_JULIAN_DAY,
930 );
931 let expected_micros = (INT96_TEST_JULIAN_DAY as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY
932 + (INT96_TEST_NANOS_WITHIN_DAY / 1_000) as i64;
933 (val, expected_micros)
934 }
935
936 async fn read_int96_batches(
937 file_path: &str,
938 schema: SchemaRef,
939 project_field_ids: Vec<i32>,
940 ) -> Vec<RecordBatch> {
941 let file_io = FileIO::new_with_fs();
942 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
943
944 let file_size = std::fs::metadata(file_path).unwrap().len();
945 let task = FileScanTask::builder()
946 .with_file_size_in_bytes(file_size)
947 .with_start(0)
948 .with_length(file_size)
949 .with_data_file_path(file_path.to_string())
950 .with_data_file_format(DataFileFormat::Parquet)
951 .with_schema(schema)
952 .with_project_field_ids(project_field_ids)
953 .with_case_sensitive(false)
954 .build()
955 .unwrap();
956
957 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
958 reader
959 .read(tasks)
960 .unwrap()
961 .stream()
962 .try_collect()
963 .await
964 .unwrap()
965 }
966
967 fn write_int96_parquet_file(
969 table_location: &str,
970 filename: &str,
971 with_field_ids: bool,
972 ) -> (String, Vec<i64>) {
973 use parquet::basic::{Repetition, Type as PhysicalType};
974 use parquet::data_type::{Int32Type, Int96, Int96Type};
975 use parquet::file::writer::SerializedFileWriter;
976 use parquet::schema::types::Type as SchemaType;
977
978 let file_path = format!("{table_location}/{filename}");
979
980 let mut ts_builder = SchemaType::primitive_type_builder("ts", PhysicalType::INT96)
981 .with_repetition(Repetition::OPTIONAL);
982 let mut id_builder = SchemaType::primitive_type_builder("id", PhysicalType::INT32)
983 .with_repetition(Repetition::REQUIRED);
984
985 if with_field_ids {
986 ts_builder = ts_builder.with_id(Some(1));
987 id_builder = id_builder.with_id(Some(2));
988 }
989
990 let schema = SchemaType::group_type_builder("schema")
991 .with_fields(vec![
992 Arc::new(ts_builder.build().unwrap()),
993 Arc::new(id_builder.build().unwrap()),
994 ])
995 .build()
996 .unwrap();
997
998 const NOON_NANOS: u64 = INT96_TEST_NANOS_WITHIN_DAY;
1000 const JULIAN_3333: u32 = INT96_TEST_JULIAN_DAY;
1001 const JULIAN_2100: u32 = 2_488_070;
1002
1003 let test_data: Vec<(u32, u32, u32, i64)> = vec![
1004 (
1006 0,
1007 0,
1008 JULIAN_3333,
1009 (JULIAN_3333 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY,
1010 ),
1011 (
1013 (NOON_NANOS & 0xFFFFFFFF) as u32,
1014 (NOON_NANOS >> 32) as u32,
1015 JULIAN_3333,
1016 (JULIAN_3333 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY
1017 + (NOON_NANOS / 1_000) as i64,
1018 ),
1019 (
1021 0,
1022 0,
1023 JULIAN_2100,
1024 (JULIAN_2100 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY,
1025 ),
1026 ];
1027
1028 let int96_values: Vec<Int96> = test_data
1029 .iter()
1030 .map(|(lo, hi, day, _)| {
1031 let mut v = Int96::new();
1032 v.set_data(*lo, *hi, *day);
1033 v
1034 })
1035 .collect();
1036
1037 let id_values: Vec<i32> = (0..test_data.len() as i32).collect();
1038 let expected_micros: Vec<i64> = test_data.iter().map(|(_, _, _, m)| *m).collect();
1039
1040 let file = File::create(&file_path).unwrap();
1041 let mut writer =
1042 SerializedFileWriter::new(file, Arc::new(schema), Default::default()).unwrap();
1043
1044 let mut row_group = writer.next_row_group().unwrap();
1045 {
1046 let mut col = row_group.next_column().unwrap().unwrap();
1048 col.typed::<Int96Type>()
1049 .write_batch(&int96_values, Some(&vec![1; test_data.len()]), None)
1050 .unwrap();
1051 col.close().unwrap();
1052 }
1053 {
1054 let mut col = row_group.next_column().unwrap().unwrap();
1055 col.typed::<Int32Type>()
1056 .write_batch(&id_values, None, None)
1057 .unwrap();
1058 col.close().unwrap();
1059 }
1060 row_group.close().unwrap();
1061 writer.close().unwrap();
1062
1063 (file_path, expected_micros)
1064 }
1065
1066 async fn assert_int96_read_matches(
1067 file_path: &str,
1068 schema: SchemaRef,
1069 project_field_ids: Vec<i32>,
1070 expected_micros: &[i64],
1071 ) {
1072 use arrow_array::TimestampMicrosecondArray;
1073
1074 let batches = read_int96_batches(file_path, schema, project_field_ids).await;
1075
1076 assert_eq!(batches.len(), 1);
1077 let ts_array = batches[0]
1078 .column(0)
1079 .as_any()
1080 .downcast_ref::<TimestampMicrosecondArray>()
1081 .expect("Expected TimestampMicrosecondArray");
1082
1083 for (i, expected) in expected_micros.iter().enumerate() {
1084 assert_eq!(
1085 ts_array.value(i),
1086 *expected,
1087 "Row {i}: got {}, expected {expected}",
1088 ts_array.value(i)
1089 );
1090 }
1091 }
1092
1093 async fn assert_encrypted_parquet_roundtrip(encryption_key: &[u8]) {
1097 let aad_prefix = b"aad_prefix";
1098
1099 let schema = Arc::new(
1100 Schema::builder()
1101 .with_schema_id(1)
1102 .with_fields(vec![
1103 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1104 ])
1105 .build()
1106 .unwrap(),
1107 );
1108
1109 let arrow_schema = Arc::new(ArrowSchema::new(vec![
1110 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1111 PARQUET_FIELD_ID_META_KEY.to_string(),
1112 "1".to_string(),
1113 )])),
1114 ]));
1115
1116 let tmp_dir = TempDir::new().unwrap();
1117 let table_location = tmp_dir.path().to_str().unwrap().to_string();
1118 let file_io = FileIO::new_with_fs();
1119
1120 let id_data = Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef;
1121 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
1122
1123 let file_path = format!("{table_location}/encrypted.parquet");
1124 write_encrypted_parquet(&file_path, &batch, encryption_key, Some(aad_prefix));
1125
1126 let key_metadata = crate::encryption::StandardKeyMetadata::try_new(encryption_key)
1127 .unwrap()
1128 .with_aad_prefix(aad_prefix)
1129 .encode()
1130 .unwrap();
1131
1132 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1133
1134 let task = FileScanTask::builder()
1135 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1136 .with_start(0)
1137 .with_length(0)
1138 .with_data_file_path(file_path)
1139 .with_data_file_format(DataFileFormat::Parquet)
1140 .with_schema(schema)
1141 .with_project_field_ids(vec![1])
1142 .with_case_sensitive(false)
1143 .with_key_metadata(Some(key_metadata))
1144 .build()
1145 .unwrap();
1146
1147 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1148 let batches: Vec<RecordBatch> = reader
1149 .read(tasks)
1150 .unwrap()
1151 .stream()
1152 .try_collect()
1153 .await
1154 .unwrap();
1155
1156 assert_eq!(batches.len(), 1);
1157 let ids = batches[0]
1158 .column(0)
1159 .as_any()
1160 .downcast_ref::<Int32Array>()
1161 .unwrap();
1162 assert_eq!(ids.values(), &[10, 20, 30]);
1163 }
1164
1165 #[tokio::test]
1166 async fn test_read_encrypted_parquet_aes_128() {
1167 assert_encrypted_parquet_roundtrip(b"0123456789abcdef").await;
1168 }
1169
1170 #[tokio::test]
1171 async fn test_read_encrypted_parquet_aes_256() {
1172 assert_encrypted_parquet_roundtrip(b"0123456789abcdef0123456789abcdef").await;
1173 }
1174
1175 #[tokio::test]
1176 async fn test_read_encrypted_parquet_without_key_metadata_fails() {
1177 let encryption_key = b"0123456789abcdef";
1178
1179 let schema = Arc::new(
1180 Schema::builder()
1181 .with_schema_id(1)
1182 .with_fields(vec![
1183 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1184 ])
1185 .build()
1186 .unwrap(),
1187 );
1188
1189 let arrow_schema = Arc::new(ArrowSchema::new(vec![
1190 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1191 PARQUET_FIELD_ID_META_KEY.to_string(),
1192 "1".to_string(),
1193 )])),
1194 ]));
1195
1196 let tmp_dir = TempDir::new().unwrap();
1197 let table_location = tmp_dir.path().to_str().unwrap().to_string();
1198 let file_io = FileIO::new_with_fs();
1199
1200 let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1201 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
1202
1203 let file_path = format!("{table_location}/encrypted_no_key.parquet");
1204 write_encrypted_parquet(&file_path, &batch, encryption_key, None);
1205
1206 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1207
1208 let task = FileScanTask::builder()
1209 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1210 .with_start(0)
1211 .with_length(0)
1212 .with_data_file_path(file_path)
1213 .with_data_file_format(DataFileFormat::Parquet)
1214 .with_schema(schema)
1215 .with_project_field_ids(vec![1])
1216 .with_case_sensitive(false)
1217 .build()
1218 .unwrap();
1219
1220 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1221 let result: Result<Vec<RecordBatch>, _> =
1222 reader.read(tasks).unwrap().stream().try_collect().await;
1223
1224 let err = result.unwrap_err();
1225 assert_eq!(err.kind(), crate::ErrorKind::Unexpected);
1226 let err_str = format!("{err}");
1227 assert!(
1228 err_str.contains("encrypted footer"),
1229 "Expected error about encrypted footer, got: {err_str}"
1230 );
1231 assert!(
1232 err_str.contains("decryption properties were not provided"),
1233 "Expected error about missing decryption properties, got: {err_str}"
1234 );
1235 }
1236
1237 fn write_plain_parquet(
1240 dir: &str,
1241 name: &str,
1242 extra_fields: Vec<Field>,
1243 extra_columns: Vec<ArrayRef>,
1244 ) -> String {
1245 let props = WriterProperties::builder()
1246 .set_compression(Compression::SNAPPY)
1247 .build();
1248 write_parquet_with(dir, name, extra_fields, extra_columns, props)
1249 }
1250
1251 fn write_parquet_with(
1252 dir: &str,
1253 name: &str,
1254 extra_fields: Vec<Field>,
1255 extra_columns: Vec<ArrayRef>,
1256 props: WriterProperties,
1257 ) -> String {
1258 let mut fields =
1259 vec![
1260 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1261 PARQUET_FIELD_ID_META_KEY.to_string(),
1262 "1".to_string(),
1263 )])),
1264 ];
1265 fields.extend(extra_fields);
1266 let arrow_schema = Arc::new(ArrowSchema::new(fields));
1267
1268 let mut columns: Vec<ArrayRef> = vec![Arc::new(Int32Array::from(vec![1, 2, 3]))];
1269 columns.extend(extra_columns);
1270 let batch = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap();
1271
1272 let file_path = format!("{dir}/{name}");
1273 let file = File::create(&file_path).unwrap();
1274 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
1275 writer.write(&batch).unwrap();
1276 writer.close().unwrap();
1277 file_path
1278 }
1279
1280 fn last_updated_seq_task(
1281 file_path: String,
1282 first_row_id: Option<i64>,
1283 data_sequence_number: Option<i64>,
1284 ) -> FileScanTask {
1285 use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1286
1287 let schema = Arc::new(
1288 Schema::builder()
1289 .with_schema_id(1)
1290 .with_fields(vec![
1291 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1292 ])
1293 .build()
1294 .unwrap(),
1295 );
1296
1297 FileScanTask::builder()
1298 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1299 .with_start(0)
1300 .with_length(0)
1301 .with_data_file_path(file_path)
1302 .with_data_file_format(DataFileFormat::Parquet)
1303 .with_schema(schema)
1304 .with_project_field_ids(vec![1, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER])
1305 .with_first_row_id(first_row_id)
1306 .with_data_sequence_number(data_sequence_number)
1307 .with_case_sensitive(false)
1308 .build()
1309 .unwrap()
1310 }
1311
1312 fn assert_last_updated_seq_column(batches: &[RecordBatch], expected: &[Option<i64>]) {
1315 use arrow_array::cast::AsArray;
1316 use arrow_cast::cast;
1317 use arrow_schema::DataType;
1318
1319 use crate::metadata_columns::RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER;
1320
1321 let mut actual = Vec::new();
1322 for batch in batches {
1323 let col = batch
1324 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
1325 .expect("_last_updated_sequence_number column should be present");
1326 let logical = cast(col, &DataType::Int64).unwrap();
1327 let values = logical.as_primitive::<arrow_array::types::Int64Type>();
1328 for i in 0..values.len() {
1329 actual.push((!values.is_null(i)).then(|| values.value(i)));
1330 }
1331 }
1332 assert_eq!(actual, expected);
1333 }
1334
1335 #[tokio::test]
1336 async fn test_last_updated_sequence_number_null_when_no_first_row_id() {
1337 let tmp_dir = TempDir::new().unwrap();
1338 let dir = tmp_dir.path().to_str().unwrap();
1339 let file_path = write_plain_parquet(dir, "no_first_row_id.parquet", vec![], vec![]);
1340
1341 let task = last_updated_seq_task(file_path, None, Some(9));
1345
1346 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1347 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1348 let batches: Vec<RecordBatch> = reader
1349 .read(tasks)
1350 .unwrap()
1351 .stream()
1352 .try_collect()
1353 .await
1354 .unwrap();
1355
1356 assert_last_updated_seq_column(&batches, &[None, None, None]);
1357 }
1358
1359 #[tokio::test]
1360 async fn test_last_updated_sequence_number_error_when_no_data_seq() {
1361 let tmp_dir = TempDir::new().unwrap();
1362 let dir = tmp_dir.path().to_str().unwrap();
1363 let file_path = write_plain_parquet(dir, "no_data_seq.parquet", vec![], vec![]);
1364
1365 let task = last_updated_seq_task(file_path, Some(42), None);
1369
1370 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1371 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1372 let result: Result<Vec<RecordBatch>, _> =
1373 reader.read(tasks).unwrap().stream().try_collect().await;
1374
1375 let err = result.unwrap_err();
1376 assert_eq!(err.kind(), crate::ErrorKind::DataInvalid);
1377 assert!(
1378 format!("{err}").contains("no data sequence number"),
1379 "unexpected error: {err}"
1380 );
1381 }
1382
1383 #[tokio::test]
1384 async fn test_last_updated_sequence_number_derived_from_data_seq() {
1385 let tmp_dir = TempDir::new().unwrap();
1386 let dir = tmp_dir.path().to_str().unwrap();
1387 let file_path = write_plain_parquet(dir, "with_first_row_id.parquet", vec![], vec![]);
1388
1389 let task = last_updated_seq_task(file_path, Some(42), Some(7));
1392
1393 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1394 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1395 let batches: Vec<RecordBatch> = reader
1396 .read(tasks)
1397 .unwrap()
1398 .stream()
1399 .try_collect()
1400 .await
1401 .unwrap();
1402
1403 assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]);
1404 }
1405
1406 #[tokio::test]
1407 async fn test_last_updated_sequence_number_mixed_files_share_schema() {
1408 use arrow_select::concat::concat_batches;
1409
1410 let tmp_dir = TempDir::new().unwrap();
1411 let dir = tmp_dir.path().to_str().unwrap();
1412
1413 let constant = last_updated_seq_task(
1419 write_plain_parquet(dir, "constant.parquet", vec![], vec![]),
1420 Some(42),
1421 Some(7),
1422 );
1423 let nulled = last_updated_seq_task(
1424 write_plain_parquet(dir, "nulled.parquet", vec![], vec![]),
1425 None,
1426 Some(7),
1427 );
1428 let coalesced = last_updated_seq_task(
1429 write_plain_parquet(
1430 dir,
1431 "coalesced.parquet",
1432 vec![physical_last_updated_seq_field()],
1433 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
1434 ),
1435 Some(50),
1436 Some(7),
1437 );
1438
1439 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1440 let tasks = Box::pin(futures::stream::iter(vec![
1441 Ok(constant),
1442 Ok(nulled),
1443 Ok(coalesced),
1444 ])) as FileScanTaskStream;
1445 let batches: Vec<RecordBatch> = reader
1446 .read(tasks)
1447 .unwrap()
1448 .stream()
1449 .try_collect()
1450 .await
1451 .unwrap();
1452
1453 assert_eq!(batches.len(), 3);
1454 let schema = batches[0].schema();
1456 concat_batches(&schema, &batches)
1457 .expect("constant, null and coalesce files must share one column type");
1458 }
1459
1460 fn physical_last_updated_seq_field() -> Field {
1462 use crate::metadata_columns::{
1463 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1464 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1465 };
1466 Field::new(
1467 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1468 DataType::Int64,
1469 true,
1470 )
1471 .with_metadata(HashMap::from([(
1472 PARQUET_FIELD_ID_META_KEY.to_string(),
1473 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER.to_string(),
1474 )]))
1475 }
1476
1477 #[tokio::test]
1478 async fn test_last_updated_sequence_number_physical_column_coalesced() {
1479 let tmp_dir = TempDir::new().unwrap();
1480 let dir = tmp_dir.path().to_str().unwrap();
1481 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1485 let file_path = write_plain_parquet(
1486 dir,
1487 "with_seq.parquet",
1488 vec![physical_last_updated_seq_field()],
1489 vec![seq_col],
1490 );
1491
1492 let task = last_updated_seq_task(file_path, Some(100), Some(9));
1493
1494 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1495 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1496 let batches: Vec<RecordBatch> = reader
1497 .read(tasks)
1498 .unwrap()
1499 .stream()
1500 .try_collect()
1501 .await
1502 .unwrap();
1503
1504 assert_last_updated_seq_column(&batches, &[Some(5), Some(9), Some(8)]);
1506 }
1507
1508 #[tokio::test]
1509 async fn test_last_updated_sequence_number_coalesced_with_pos_column() {
1510 use crate::metadata_columns::{
1511 RESERVED_COL_NAME_POS, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1512 RESERVED_FIELD_ID_POS,
1513 };
1514
1515 let tmp_dir = TempDir::new().unwrap();
1516 let dir = tmp_dir.path().to_str().unwrap();
1517 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1518 let file_path = write_plain_parquet(
1519 dir,
1520 "with_seq_and_pos.parquet",
1521 vec![physical_last_updated_seq_field()],
1522 vec![seq_col],
1523 );
1524
1525 let schema = Arc::new(
1530 Schema::builder()
1531 .with_schema_id(1)
1532 .with_fields(vec![
1533 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1534 ])
1535 .build()
1536 .unwrap(),
1537 );
1538 let task = FileScanTask::builder()
1539 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1540 .with_start(0)
1541 .with_length(0)
1542 .with_data_file_path(file_path)
1543 .with_data_file_format(DataFileFormat::Parquet)
1544 .with_schema(schema)
1545 .with_project_field_ids(vec![
1546 1,
1547 RESERVED_FIELD_ID_POS,
1548 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1549 ])
1550 .with_first_row_id(Some(100))
1551 .with_data_sequence_number(Some(9))
1552 .with_case_sensitive(false)
1553 .build()
1554 .unwrap();
1555
1556 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1557 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1558 let batches: Vec<RecordBatch> = reader
1559 .read(tasks)
1560 .unwrap()
1561 .stream()
1562 .try_collect()
1563 .await
1564 .unwrap();
1565
1566 assert_last_updated_seq_column(&batches, &[Some(5), Some(9), Some(8)]);
1568 let pos_col = batches[0]
1570 .column_by_name(RESERVED_COL_NAME_POS)
1571 .expect("_pos column should be present")
1572 .as_primitive::<arrow_array::types::Int64Type>();
1573 assert_eq!(pos_col.values(), &[0, 1, 2]);
1574 }
1575
1576 #[tokio::test]
1577 async fn test_last_updated_sequence_number_physical_column_nulled_without_first_row_id() {
1578 let tmp_dir = TempDir::new().unwrap();
1579 let dir = tmp_dir.path().to_str().unwrap();
1580 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1581 let file_path = write_plain_parquet(
1582 dir,
1583 "with_seq_no_first_row_id.parquet",
1584 vec![physical_last_updated_seq_field()],
1585 vec![seq_col],
1586 );
1587
1588 let task = last_updated_seq_task(file_path, None, Some(9));
1591
1592 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1593 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1594 let batches: Vec<RecordBatch> = reader
1595 .read(tasks)
1596 .unwrap()
1597 .stream()
1598 .try_collect()
1599 .await
1600 .unwrap();
1601
1602 assert_last_updated_seq_column(&batches, &[None, None, None]);
1603 }
1604
1605 #[tokio::test]
1606 async fn test_last_updated_sequence_number_present_by_name_without_id_unsupported() {
1607 let tmp_dir = TempDir::new().unwrap();
1608 let dir = tmp_dir.path().to_str().unwrap();
1609 let seq_field = Field::new(
1613 crate::metadata_columns::RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1614 DataType::Int64,
1615 true,
1616 );
1617 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1618 let file_path =
1619 write_plain_parquet(dir, "with_seq_by_name.parquet", vec![seq_field], vec![
1620 seq_col,
1621 ]);
1622
1623 let task = last_updated_seq_task(file_path, Some(100), Some(9));
1624
1625 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1626 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1627 let result: Result<Vec<RecordBatch>, _> =
1628 reader.read(tasks).unwrap().stream().try_collect().await;
1629
1630 let err = result.unwrap_err();
1631 assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported);
1632 assert!(
1633 format!("{err}").contains("without an embedded field id"),
1634 "unexpected error: {err}"
1635 );
1636 }
1637
1638 #[tokio::test]
1639 async fn test_last_updated_sequence_number_physical_column_first_row_id_without_data_seq() {
1640 let tmp_dir = TempDir::new().unwrap();
1641 let dir = tmp_dir.path().to_str().unwrap();
1642 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1643 let file_path = write_plain_parquet(
1644 dir,
1645 "with_seq_no_data_seq.parquet",
1646 vec![physical_last_updated_seq_field()],
1647 vec![seq_col],
1648 );
1649
1650 let task = last_updated_seq_task(file_path, Some(100), None);
1654
1655 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1656 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1657 let result: Result<Vec<RecordBatch>, _> =
1658 reader.read(tasks).unwrap().stream().try_collect().await;
1659
1660 let err = result.unwrap_err();
1661 assert_eq!(err.kind(), crate::ErrorKind::DataInvalid);
1662 assert!(
1663 format!("{err}").contains("no data sequence number"),
1664 "unexpected error: {err}"
1665 );
1666 }
1667
1668 fn row_id_task(file_path: String, first_row_id: Option<i64>) -> FileScanTask {
1670 row_id_task_with_options(file_path, first_row_id, 0, 0, vec![])
1671 }
1672
1673 fn row_id_task_with_options(
1674 file_path: String,
1675 first_row_id: Option<i64>,
1676 start: u64,
1677 length: u64,
1678 deletes: Vec<FileScanTaskDeleteFile>,
1679 ) -> FileScanTask {
1680 let schema = Arc::new(
1681 Schema::builder()
1682 .with_schema_id(1)
1683 .with_fields(vec![
1684 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1685 ])
1686 .build()
1687 .unwrap(),
1688 );
1689
1690 FileScanTask::builder()
1691 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1692 .with_start(start)
1693 .with_length(length)
1694 .with_data_file_path(file_path)
1695 .with_data_file_format(DataFileFormat::Parquet)
1696 .with_schema(schema)
1697 .with_project_field_ids(vec![1, RESERVED_FIELD_ID_ROW_ID])
1698 .with_first_row_id(first_row_id)
1699 .with_deletes(deletes)
1700 .with_case_sensitive(false)
1701 .build()
1702 .unwrap()
1703 }
1704
1705 fn assert_row_id_column(batches: &[RecordBatch], expected: &[Option<i64>]) {
1708 use arrow_array::cast::AsArray;
1709 use arrow_cast::cast;
1710 use arrow_schema::DataType;
1711
1712 let mut actual = Vec::new();
1713 for batch in batches {
1714 let col = batch
1715 .column_by_name(RESERVED_COL_NAME_ROW_ID)
1716 .expect("_row_id column should be present");
1717 let logical = cast(col, &DataType::Int64).unwrap();
1718 let values = logical.as_primitive::<arrow_array::types::Int64Type>();
1719 for i in 0..values.len() {
1720 actual.push((!values.is_null(i)).then(|| values.value(i)));
1721 }
1722 }
1723 assert_eq!(actual, expected);
1724 }
1725
1726 fn physical_row_id_field() -> Field {
1728 Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true).with_metadata(HashMap::from([
1729 (
1730 PARQUET_FIELD_ID_META_KEY.to_string(),
1731 RESERVED_FIELD_ID_ROW_ID.to_string(),
1732 ),
1733 ]))
1734 }
1735
1736 #[tokio::test]
1737 async fn test_row_id_synthesized_from_first_row_id_and_pos() {
1738 let tmp_dir = TempDir::new().unwrap();
1739 let dir = tmp_dir.path().to_str().unwrap();
1740 let file_path = write_plain_parquet(dir, "row_id_synth.parquet", vec![], vec![]);
1741
1742 let task = row_id_task(file_path, Some(100));
1744
1745 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1746 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1747 let batches: Vec<RecordBatch> = reader
1748 .read(tasks)
1749 .unwrap()
1750 .stream()
1751 .try_collect()
1752 .await
1753 .unwrap();
1754
1755 assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
1756 }
1757
1758 #[tokio::test]
1759 async fn test_row_id_physical_column_coalesced() {
1760 let tmp_dir = TempDir::new().unwrap();
1761 let dir = tmp_dir.path().to_str().unwrap();
1762 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1765 let file_path = write_plain_parquet(
1766 dir,
1767 "row_id_phys.parquet",
1768 vec![physical_row_id_field()],
1769 vec![id_col],
1770 );
1771
1772 let task = row_id_task(file_path, Some(100));
1773
1774 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1775 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1776 let batches: Vec<RecordBatch> = reader
1777 .read(tasks)
1778 .unwrap()
1779 .stream()
1780 .try_collect()
1781 .await
1782 .unwrap();
1783
1784 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
1786 }
1787
1788 #[tokio::test]
1789 async fn test_row_id_only_synthesis_reads_no_data_columns() {
1790 let tmp_dir = TempDir::new().unwrap();
1795 let dir = tmp_dir.path().to_str().unwrap();
1796
1797 let meta_only = metadata_projection_task_with_first_row_id(
1798 write_parquet_with_wide_column(dir, "row_id_only.parquet", vec![], vec![]),
1799 id_and_wide_schema(),
1800 vec![RESERVED_FIELD_ID_ROW_ID],
1801 Some(100),
1802 );
1803 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1804
1805 assert_eq!(batches[0].num_columns(), 1);
1806 assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
1807
1808 let with_data = metadata_projection_task_with_first_row_id(
1810 write_parquet_with_wide_column(dir, "row_id_only_ref.parquet", vec![], vec![]),
1811 id_and_wide_schema(),
1812 vec![2, RESERVED_FIELD_ID_ROW_ID],
1813 Some(100),
1814 );
1815 let (_, with_data_bytes) = scan_task(with_data).await;
1816
1817 assert!(
1818 meta_only_bytes < with_data_bytes,
1819 "_row_id-only synthesis should read fewer bytes than a scan of the wide column: \
1820 {meta_only_bytes} vs {with_data_bytes}"
1821 );
1822 }
1823
1824 #[tokio::test]
1825 async fn test_row_id_only_null_first_row_id_reads_no_data_columns() {
1826 let tmp_dir = TempDir::new().unwrap();
1831 let dir = tmp_dir.path().to_str().unwrap();
1832
1833 let meta_only = metadata_projection_task_with_first_row_id(
1834 write_parquet_with_wide_column(dir, "row_id_null.parquet", vec![], vec![]),
1835 id_and_wide_schema(),
1836 vec![RESERVED_FIELD_ID_ROW_ID],
1837 None,
1838 );
1839 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1840
1841 assert_eq!(batches[0].num_columns(), 1);
1842 assert_row_id_column(&batches, &[None, None, None]);
1843
1844 let with_data = metadata_projection_task_with_first_row_id(
1845 write_parquet_with_wide_column(dir, "row_id_null_ref.parquet", vec![], vec![]),
1846 id_and_wide_schema(),
1847 vec![2, RESERVED_FIELD_ID_ROW_ID],
1848 None,
1849 );
1850 let (_, with_data_bytes) = scan_task(with_data).await;
1851
1852 assert!(
1853 meta_only_bytes < with_data_bytes,
1854 "_row_id-only scan with a null first_row_id should read fewer bytes than a scan \
1855 of the wide column: {meta_only_bytes} vs {with_data_bytes}"
1856 );
1857 }
1858
1859 #[tokio::test]
1860 async fn test_last_updated_seq_only_reads_no_data_columns() {
1861 use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1862
1863 let tmp_dir = TempDir::new().unwrap();
1868 let dir = tmp_dir.path().to_str().unwrap();
1869
1870 let lusn_only_task = |file_path: String, project_field_ids: Vec<i32>| {
1871 FileScanTask::builder()
1872 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1873 .with_start(0)
1874 .with_length(0)
1875 .with_data_file_path(file_path)
1876 .with_data_file_format(DataFileFormat::Parquet)
1877 .with_schema(id_and_wide_schema())
1878 .with_project_field_ids(project_field_ids)
1879 .with_first_row_id(Some(42))
1880 .with_data_sequence_number(Some(7))
1881 .with_case_sensitive(false)
1882 .build()
1883 .unwrap()
1884 };
1885
1886 let meta_only = lusn_only_task(
1887 write_parquet_with_wide_column(dir, "lusn_only.parquet", vec![], vec![]),
1888 vec![RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER],
1889 );
1890 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1891
1892 assert_eq!(batches[0].num_columns(), 1);
1893 assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]);
1894
1895 let with_data = lusn_only_task(
1896 write_parquet_with_wide_column(dir, "lusn_only_ref.parquet", vec![], vec![]),
1897 vec![2, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER],
1898 );
1899 let (_, with_data_bytes) = scan_task(with_data).await;
1900
1901 assert!(
1902 meta_only_bytes < with_data_bytes,
1903 "_last_updated_sequence_number-only scan should read fewer bytes than a scan of \
1904 the wide column: {meta_only_bytes} vs {with_data_bytes}"
1905 );
1906 }
1907
1908 #[tokio::test]
1909 async fn test_partition_only_reads_no_data_columns() {
1910 use arrow_array::StructArray;
1911
1912 use crate::metadata_columns::{RESERVED_COL_NAME_PARTITION, RESERVED_FIELD_ID_PARTITION};
1913 use crate::spec::{Literal, PartitionSpec, Struct, Transform};
1914
1915 let tmp_dir = TempDir::new().unwrap();
1919 let dir = tmp_dir.path().to_str().unwrap();
1920 let schema = id_and_wide_schema();
1921 let spec = Arc::new(
1922 PartitionSpec::builder(schema.clone())
1923 .with_spec_id(7)
1924 .add_partition_field("id", "id", Transform::Identity)
1925 .unwrap()
1926 .build()
1927 .unwrap(),
1928 );
1929 let unified_type = Arc::new(spec.partition_type(&schema).unwrap());
1930 let partition_data = Struct::from_iter(vec![Some(Literal::int(42))]);
1931
1932 let partition_task = |file_path: String, project_field_ids: Vec<i32>| {
1933 FileScanTask::builder()
1934 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1935 .with_start(0)
1936 .with_length(0)
1937 .with_data_file_path(file_path)
1938 .with_data_file_format(DataFileFormat::Parquet)
1939 .with_schema(schema.clone())
1940 .with_project_field_ids(project_field_ids)
1941 .with_partition_spec(Some(spec.clone()))
1942 .with_partition(Some(partition_data.clone()))
1943 .with_unified_partition_type(Some(unified_type.clone()))
1944 .with_case_sensitive(false)
1945 .build()
1946 .unwrap()
1947 };
1948
1949 let meta_only = partition_task(
1950 write_parquet_with_wide_column(dir, "partition_only.parquet", vec![], vec![]),
1951 vec![RESERVED_FIELD_ID_PARTITION],
1952 );
1953 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1954
1955 assert_eq!(batches[0].num_columns(), 1);
1956 let partition_col = batches[0]
1957 .column_by_name(RESERVED_COL_NAME_PARTITION)
1958 .expect("_partition column should be present")
1959 .as_any()
1960 .downcast_ref::<StructArray>()
1961 .unwrap();
1962 assert_eq!(partition_col.len(), 3);
1963 let inner = partition_col
1964 .column(0)
1965 .as_any()
1966 .downcast_ref::<Int32Array>()
1967 .unwrap();
1968 assert_eq!(inner.values(), &[42, 42, 42]);
1969
1970 let with_data = partition_task(
1971 write_parquet_with_wide_column(dir, "partition_only_ref.parquet", vec![], vec![]),
1972 vec![2, RESERVED_FIELD_ID_PARTITION],
1973 );
1974 let (_, with_data_bytes) = scan_task(with_data).await;
1975
1976 assert!(
1977 meta_only_bytes < with_data_bytes,
1978 "_partition-only scan should read fewer bytes than a scan of the wide column: \
1979 {meta_only_bytes} vs {with_data_bytes}"
1980 );
1981 }
1982
1983 #[tokio::test]
1984 async fn test_row_id_resolves_alongside_id_less_leaf() {
1985 let tmp_dir = TempDir::new().unwrap();
1990 let dir = tmp_dir.path().to_str().unwrap();
1991 let idless_field = Field::new("variant_internal", DataType::Utf8, true);
1992 let idless_col = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
1993 let row_id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1994 let file_path = write_plain_parquet(
1995 dir,
1996 "row_id_with_idless_leaf.parquet",
1997 vec![idless_field, physical_row_id_field()],
1998 vec![idless_col, row_id_col],
1999 );
2000
2001 let task = row_id_task(file_path, Some(100));
2002
2003 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2004 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2005 let batches: Vec<RecordBatch> = reader
2006 .read(tasks)
2007 .unwrap()
2008 .stream()
2009 .try_collect()
2010 .await
2011 .unwrap();
2012
2013 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2015 }
2016
2017 #[tokio::test]
2018 async fn test_row_id_and_last_updated_seq_co_projected() {
2019 use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
2020
2021 let tmp_dir = TempDir::new().unwrap();
2025 let dir = tmp_dir.path().to_str().unwrap();
2026 let row_id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2027 let seq_col = Arc::new(Int64Array::from(vec![Some(50), None, Some(70)])) as ArrayRef;
2028 let file_path = write_plain_parquet(
2029 dir,
2030 "row_id_and_seq.parquet",
2031 vec![physical_row_id_field(), physical_last_updated_seq_field()],
2032 vec![row_id_col, seq_col],
2033 );
2034
2035 let schema = Arc::new(
2036 Schema::builder()
2037 .with_schema_id(1)
2038 .with_fields(vec![
2039 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2040 ])
2041 .build()
2042 .unwrap(),
2043 );
2044 let task = FileScanTask::builder()
2045 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2046 .with_start(0)
2047 .with_length(0)
2048 .with_data_file_path(file_path)
2049 .with_data_file_format(DataFileFormat::Parquet)
2050 .with_schema(schema)
2051 .with_project_field_ids(vec![
2052 1,
2053 RESERVED_FIELD_ID_ROW_ID,
2054 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
2055 ])
2056 .with_first_row_id(Some(100))
2057 .with_data_sequence_number(Some(9))
2058 .with_case_sensitive(false)
2059 .build()
2060 .unwrap();
2061
2062 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2063 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2064 let batches: Vec<RecordBatch> = reader
2065 .read(tasks)
2066 .unwrap()
2067 .stream()
2068 .try_collect()
2069 .await
2070 .unwrap();
2071
2072 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2074 assert_last_updated_seq_column(&batches, &[Some(50), Some(9), Some(70)]);
2076 }
2077
2078 #[tokio::test]
2079 async fn test_row_id_null_when_no_first_row_id() {
2080 let tmp_dir = TempDir::new().unwrap();
2081 let dir = tmp_dir.path().to_str().unwrap();
2082 let id_col = Arc::new(Int64Array::from(vec![Some(5), Some(6), Some(7)])) as ArrayRef;
2084 let file_path = write_plain_parquet(
2085 dir,
2086 "row_id_no_first.parquet",
2087 vec![physical_row_id_field()],
2088 vec![id_col],
2089 );
2090
2091 let task = row_id_task(file_path, None);
2093
2094 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2095 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2096 let batches: Vec<RecordBatch> = reader
2097 .read(tasks)
2098 .unwrap()
2099 .stream()
2100 .try_collect()
2101 .await
2102 .unwrap();
2103
2104 assert_row_id_column(&batches, &[None, None, None]);
2105 }
2106
2107 #[tokio::test]
2108 async fn test_row_id_with_pos_column() {
2109 use crate::metadata_columns::RESERVED_COL_NAME_POS;
2110
2111 let tmp_dir = TempDir::new().unwrap();
2112 let dir = tmp_dir.path().to_str().unwrap();
2113 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2114 let file_path = write_plain_parquet(
2115 dir,
2116 "row_id_and_pos.parquet",
2117 vec![physical_row_id_field()],
2118 vec![id_col],
2119 );
2120
2121 let schema = Arc::new(
2125 Schema::builder()
2126 .with_schema_id(1)
2127 .with_fields(vec![
2128 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2129 ])
2130 .build()
2131 .unwrap(),
2132 );
2133 let task = FileScanTask::builder()
2134 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2135 .with_start(0)
2136 .with_length(0)
2137 .with_data_file_path(file_path)
2138 .with_data_file_format(DataFileFormat::Parquet)
2139 .with_schema(schema)
2140 .with_project_field_ids(vec![1, RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID])
2141 .with_first_row_id(Some(100))
2142 .with_case_sensitive(false)
2143 .build()
2144 .unwrap();
2145
2146 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2147 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2148 let batches: Vec<RecordBatch> = reader
2149 .read(tasks)
2150 .unwrap()
2151 .stream()
2152 .try_collect()
2153 .await
2154 .unwrap();
2155
2156 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2158 let pos_col = batches[0]
2160 .column_by_name(RESERVED_COL_NAME_POS)
2161 .expect("_pos column should be present")
2162 .as_primitive::<arrow_array::types::Int64Type>();
2163 assert_eq!(pos_col.values(), &[0, 1, 2]);
2164 }
2165
2166 #[tokio::test]
2167 async fn test_row_id_mixed_files_share_schema() {
2168 use arrow_select::concat::concat_batches;
2169
2170 let tmp_dir = TempDir::new().unwrap();
2171 let dir = tmp_dir.path().to_str().unwrap();
2172
2173 let synth = row_id_task(
2179 write_plain_parquet(dir, "row_id_synth2.parquet", vec![], vec![]),
2180 Some(42),
2181 );
2182 let nulled = row_id_task(
2183 write_plain_parquet(dir, "row_id_null2.parquet", vec![], vec![]),
2184 None,
2185 );
2186 let coalesced = row_id_task(
2187 write_plain_parquet(
2188 dir,
2189 "row_id_coalesced2.parquet",
2190 vec![physical_row_id_field()],
2191 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
2192 ),
2193 Some(50),
2194 );
2195
2196 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2197 let tasks = Box::pin(futures::stream::iter(vec![
2198 Ok(synth),
2199 Ok(nulled),
2200 Ok(coalesced),
2201 ])) as FileScanTaskStream;
2202 let batches: Vec<RecordBatch> = reader
2203 .read(tasks)
2204 .unwrap()
2205 .stream()
2206 .try_collect()
2207 .await
2208 .unwrap();
2209
2210 assert_eq!(batches.len(), 3);
2211 let schema = batches[0].schema();
2212 concat_batches(&schema, &batches)
2213 .expect("synthesis, null and coalesce files must share one column type");
2214 }
2215
2216 #[tokio::test]
2217 async fn test_row_id_present_by_name_without_id_unsupported() {
2218 let tmp_dir = TempDir::new().unwrap();
2219 let dir = tmp_dir.path().to_str().unwrap();
2220 let id_field = Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true);
2223 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2224 let file_path =
2225 write_plain_parquet(dir, "row_id_by_name.parquet", vec![id_field], vec![id_col]);
2226
2227 let task = row_id_task(file_path, Some(100));
2228
2229 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2230 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2231 let result: Result<Vec<RecordBatch>, _> =
2232 reader.read(tasks).unwrap().stream().try_collect().await;
2233
2234 let err = result.unwrap_err();
2235 assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported);
2236 assert!(
2237 format!("{err}").contains("without an embedded field id"),
2238 "unexpected error: {err}"
2239 );
2240 }
2241
2242 #[tokio::test]
2243 async fn test_row_id_present_by_name_without_id_nulls_when_no_lineage() {
2244 let tmp_dir = TempDir::new().unwrap();
2245 let dir = tmp_dir.path().to_str().unwrap();
2246 let id_field = Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true);
2251 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2252 let file_path = write_plain_parquet(
2253 dir,
2254 "row_id_by_name_no_lineage.parquet",
2255 vec![id_field],
2256 vec![id_col],
2257 );
2258
2259 let task = row_id_task(file_path, None);
2260
2261 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2262 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2263 let batches: Vec<RecordBatch> = reader
2264 .read(tasks)
2265 .unwrap()
2266 .stream()
2267 .try_collect()
2268 .await
2269 .unwrap();
2270
2271 assert_row_id_column(&batches, &[None, None, None]);
2272 }
2273
2274 fn row_id_task_with_predicate(
2277 file_path: String,
2278 first_row_id: Option<i64>,
2279 extra_project_field_ids: Vec<i32>,
2280 predicate: crate::expr::Predicate,
2281 ) -> FileScanTask {
2282 use crate::expr::Bind;
2283
2284 let schema = Arc::new(
2285 Schema::builder()
2286 .with_schema_id(1)
2287 .with_fields(vec![
2288 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2289 ])
2290 .build()
2291 .unwrap(),
2292 );
2293 let bound = predicate.bind(Arc::clone(&schema), false).unwrap();
2294
2295 let mut project_field_ids = vec![1];
2296 project_field_ids.extend(extra_project_field_ids);
2297 project_field_ids.push(RESERVED_FIELD_ID_ROW_ID);
2298
2299 FileScanTask::builder()
2300 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2301 .with_start(0)
2302 .with_length(0)
2303 .with_data_file_path(file_path)
2304 .with_data_file_format(DataFileFormat::Parquet)
2305 .with_schema(schema)
2306 .with_project_field_ids(project_field_ids)
2307 .with_predicate(Some(bound))
2308 .with_first_row_id(first_row_id)
2309 .with_case_sensitive(false)
2310 .build()
2311 .unwrap()
2312 }
2313
2314 #[tokio::test]
2315 async fn test_row_id_stable_under_row_selection() {
2316 use crate::expr::Reference;
2317 use crate::spec::Datum;
2318
2319 let tmp_dir = TempDir::new().unwrap();
2320 let dir = tmp_dir.path().to_str().unwrap();
2321 let file_path = write_plain_parquet(dir, "row_id_selection.parquet", vec![], vec![]);
2323
2324 let task = row_id_task_with_predicate(
2325 file_path,
2326 Some(100),
2327 vec![],
2328 Reference::new("id").not_equal_to(Datum::int(2)),
2329 );
2330
2331 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
2333 .with_row_selection_enabled(true)
2334 .build();
2335 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2336 let batches: Vec<RecordBatch> = reader
2337 .read(tasks)
2338 .unwrap()
2339 .stream()
2340 .try_collect()
2341 .await
2342 .unwrap();
2343
2344 assert_row_id_column(&batches, &[Some(100), Some(102)]);
2347 }
2348
2349 #[tokio::test]
2350 async fn test_row_id_coalesce_stable_under_row_selection() {
2351 use crate::expr::Reference;
2352 use crate::spec::Datum;
2353
2354 let tmp_dir = TempDir::new().unwrap();
2355 let dir = tmp_dir.path().to_str().unwrap();
2356 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2359 let file_path = write_plain_parquet(
2360 dir,
2361 "row_id_coalesce_selection.parquet",
2362 vec![physical_row_id_field()],
2363 vec![id_col],
2364 );
2365
2366 let task = row_id_task_with_predicate(
2367 file_path,
2368 Some(100),
2369 vec![],
2370 Reference::new("id").not_equal_to(Datum::int(2)),
2371 );
2372
2373 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
2374 .with_row_selection_enabled(true)
2375 .build();
2376 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2377 let batches: Vec<RecordBatch> = reader
2378 .read(tasks)
2379 .unwrap()
2380 .stream()
2381 .try_collect()
2382 .await
2383 .unwrap();
2384
2385 assert_row_id_column(&batches, &[Some(5), Some(8)]);
2389 }
2390
2391 #[tokio::test]
2392 async fn test_row_id_global_across_row_groups() {
2393 let tmp_dir = TempDir::new().unwrap();
2394 let dir = tmp_dir.path().to_str().unwrap();
2395
2396 let file_path = format!("{dir}/row_id_multi_rg.parquet");
2400 let field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2401 PARQUET_FIELD_ID_META_KEY.to_string(),
2402 "1".to_string(),
2403 )]));
2404 let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
2405 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int32Array::from(
2406 vec![1, 2, 3, 4, 5],
2407 ))])
2408 .unwrap();
2409 let props = WriterProperties::builder()
2410 .set_compression(Compression::SNAPPY)
2411 .set_max_row_group_row_count(Some(2))
2412 .build();
2413 let file = File::create(&file_path).unwrap();
2414 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2415 writer.write(&batch).unwrap();
2416 writer.close().unwrap();
2417
2418 let task = row_id_task(file_path, Some(0));
2419 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2420 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2421 let batches: Vec<RecordBatch> = reader
2422 .read(tasks)
2423 .unwrap()
2424 .stream()
2425 .try_collect()
2426 .await
2427 .unwrap();
2428
2429 assert_row_id_column(&batches, &[Some(0), Some(1), Some(2), Some(3), Some(4)]);
2430 }
2431
2432 #[tokio::test]
2433 async fn test_row_id_global_when_first_row_group_pruned() {
2434 use parquet::file::reader::{FileReader, SerializedFileReader};
2435
2436 let tmp_dir = TempDir::new().unwrap();
2437 let dir = tmp_dir.path().to_str().unwrap();
2438
2439 let file_path = format!("{dir}/row_id_prune_rg0.parquet");
2444 let field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2445 PARQUET_FIELD_ID_META_KEY.to_string(),
2446 "1".to_string(),
2447 )]));
2448 let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
2449 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int32Array::from(
2450 vec![1, 2, 3, 4, 5, 6],
2451 ))])
2452 .unwrap();
2453 let props = WriterProperties::builder()
2454 .set_compression(Compression::SNAPPY)
2455 .set_max_row_group_row_count(Some(2))
2456 .build();
2457 let file = File::create(&file_path).unwrap();
2458 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2459 writer.write(&batch).unwrap();
2460 writer.close().unwrap();
2461
2462 let metadata = SerializedFileReader::new(File::open(&file_path).unwrap())
2465 .unwrap()
2466 .metadata()
2467 .clone();
2468 assert_eq!(metadata.num_row_groups(), 3);
2469 let start = 4 + metadata.row_group(0).compressed_size() as u64;
2470 let file_size = std::fs::metadata(&file_path).unwrap().len();
2471
2472 let task = row_id_task_with_options(file_path, Some(100), start, file_size - start, vec![]);
2473
2474 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2475 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2476 let batches: Vec<RecordBatch> = reader
2477 .read(tasks)
2478 .unwrap()
2479 .stream()
2480 .try_collect()
2481 .await
2482 .unwrap();
2483
2484 assert_row_id_column(&batches, &[Some(102), Some(103), Some(104), Some(105)]);
2487 }
2488
2489 fn write_positional_delete(
2492 dir: &str,
2493 name: &str,
2494 data_file_path: &str,
2495 positions: &[i64],
2496 ) -> String {
2497 use arrow_array::StringArray;
2498
2499 let file_path_field =
2500 Field::new("file_path", DataType::Utf8, false).with_metadata(HashMap::from([(
2501 PARQUET_FIELD_ID_META_KEY.to_string(),
2502 "2147483546".to_string(),
2503 )]));
2504 let pos_field = Field::new("pos", DataType::Int64, false).with_metadata(HashMap::from([(
2505 PARQUET_FIELD_ID_META_KEY.to_string(),
2506 "2147483545".to_string(),
2507 )]));
2508 let arrow_schema = Arc::new(ArrowSchema::new(vec![file_path_field, pos_field]));
2509 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![
2510 Arc::new(StringArray::from(vec![data_file_path; positions.len()])),
2511 Arc::new(Int64Array::from(positions.to_vec())),
2512 ])
2513 .unwrap();
2514
2515 let path = format!("{dir}/{name}");
2516 let file = File::create(&path).unwrap();
2517 let props = WriterProperties::builder()
2518 .set_compression(Compression::SNAPPY)
2519 .build();
2520 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2521 writer.write(&batch).unwrap();
2522 writer.close().unwrap();
2523 path
2524 }
2525
2526 #[tokio::test]
2527 async fn test_row_id_survives_positional_delete() {
2528 use crate::spec::DataContentType;
2529
2530 let tmp_dir = TempDir::new().unwrap();
2531 let dir = tmp_dir.path().to_str().unwrap();
2532 let data_path = write_plain_parquet(dir, "row_id_posdel_data.parquet", vec![], vec![]);
2534 let del_path = write_positional_delete(dir, "row_id_posdel.parquet", &data_path, &[1]);
2535
2536 let delete = FileScanTaskDeleteFile::builder()
2537 .with_file_path(del_path.clone())
2538 .with_file_size_in_bytes(std::fs::metadata(&del_path).unwrap().len())
2539 .with_file_type(DataContentType::PositionDeletes)
2540 .with_file_format(DataFileFormat::Parquet)
2541 .with_partition_spec_id(0)
2542 .build()
2543 .unwrap();
2544 let task = row_id_task_with_options(data_path, Some(100), 0, 0, vec![delete]);
2545
2546 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2547 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2548 let batches: Vec<RecordBatch> = reader
2549 .read(tasks)
2550 .unwrap()
2551 .stream()
2552 .try_collect()
2553 .await
2554 .unwrap();
2555
2556 assert_row_id_column(&batches, &[Some(100), Some(102)]);
2560 }
2561
2562 #[tokio::test]
2563 async fn test_row_id_name_collision_under_positional_fallback() {
2564 let tmp_dir = TempDir::new().unwrap();
2565 let dir = tmp_dir.path().to_str().unwrap();
2566
2567 let file_path = format!("{dir}/fallback_row_id_name.parquet");
2572 let arrow_schema = Arc::new(ArrowSchema::new(vec![
2573 Field::new("id", DataType::Int32, false),
2574 Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true),
2575 ]));
2576 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![
2577 Arc::new(Int32Array::from(vec![1, 2, 3])),
2578 Arc::new(Int64Array::from(vec![7i64, 8, 9])),
2579 ])
2580 .unwrap();
2581 let props = WriterProperties::builder()
2582 .set_compression(Compression::SNAPPY)
2583 .build();
2584 let file = File::create(&file_path).unwrap();
2585 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2586 writer.write(&batch).unwrap();
2587 writer.close().unwrap();
2588
2589 let task = row_id_task(file_path, Some(100));
2590 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2591 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2592 let batches: Vec<RecordBatch> = reader
2593 .read(tasks)
2594 .unwrap()
2595 .stream()
2596 .try_collect()
2597 .await
2598 .unwrap();
2599
2600 assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
2602 }
2603
2604 #[tokio::test]
2605 async fn test_read_encrypted_parquet_with_wrong_key_fails() {
2606 let encryption_key = b"0123456789abcdef";
2607 let wrong_key = b"fedcba9876543210";
2608
2609 let schema = Arc::new(
2610 Schema::builder()
2611 .with_schema_id(1)
2612 .with_fields(vec![
2613 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2614 ])
2615 .build()
2616 .unwrap(),
2617 );
2618
2619 let arrow_schema = Arc::new(ArrowSchema::new(vec![
2620 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2621 PARQUET_FIELD_ID_META_KEY.to_string(),
2622 "1".to_string(),
2623 )])),
2624 ]));
2625
2626 let tmp_dir = TempDir::new().unwrap();
2627 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2628 let file_io = FileIO::new_with_fs();
2629
2630 let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
2631 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
2632
2633 let file_path = format!("{table_location}/encrypted_wrong_key.parquet");
2634 write_encrypted_parquet(&file_path, &batch, encryption_key, None);
2635
2636 let wrong_key_metadata = crate::encryption::StandardKeyMetadata::try_new(wrong_key)
2637 .unwrap()
2638 .encode()
2639 .unwrap();
2640
2641 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
2642
2643 let task = FileScanTask::builder()
2644 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2645 .with_start(0)
2646 .with_length(0)
2647 .with_data_file_path(file_path)
2648 .with_data_file_format(DataFileFormat::Parquet)
2649 .with_schema(schema)
2650 .with_project_field_ids(vec![1])
2651 .with_case_sensitive(false)
2652 .with_key_metadata(Some(wrong_key_metadata))
2653 .build()
2654 .unwrap();
2655
2656 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2657 let result: Result<Vec<RecordBatch>, _> =
2658 reader.read(tasks).unwrap().stream().try_collect().await;
2659
2660 let err = result.unwrap_err();
2661 assert_eq!(err.kind(), crate::ErrorKind::Unexpected);
2662 let err_str = format!("{err}");
2663 assert!(
2664 err_str.contains("unable to decrypt parquet footer"),
2665 "Expected error about decryption failure, got: {err_str}"
2666 );
2667 }
2668
2669 #[tokio::test]
2672 async fn test_read_with_concurrency_one() {
2673 use arrow_array::Int32Array;
2674
2675 let schema = Arc::new(
2676 Schema::builder()
2677 .with_schema_id(1)
2678 .with_fields(vec![
2679 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2680 NestedField::required(2, "file_num", Type::Primitive(PrimitiveType::Int))
2681 .into(),
2682 ])
2683 .build()
2684 .unwrap(),
2685 );
2686
2687 let arrow_schema = Arc::new(ArrowSchema::new(vec![
2688 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2689 PARQUET_FIELD_ID_META_KEY.to_string(),
2690 "1".to_string(),
2691 )])),
2692 Field::new("file_num", DataType::Int32, false).with_metadata(HashMap::from([(
2693 PARQUET_FIELD_ID_META_KEY.to_string(),
2694 "2".to_string(),
2695 )])),
2696 ]));
2697
2698 let tmp_dir = TempDir::new().unwrap();
2699 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2700 let file_io = FileIO::new_with_fs();
2701
2702 let props = WriterProperties::builder()
2704 .set_compression(Compression::SNAPPY)
2705 .build();
2706
2707 for file_num in 0..3 {
2708 let id_data = Arc::new(Int32Array::from_iter_values(
2709 file_num * 10..(file_num + 1) * 10,
2710 )) as ArrayRef;
2711 let file_num_data = Arc::new(Int32Array::from(vec![file_num; 10])) as ArrayRef;
2712
2713 let to_write =
2714 RecordBatch::try_new(arrow_schema.clone(), vec![id_data, file_num_data]).unwrap();
2715
2716 let file = File::create(format!("{table_location}/file_{file_num}.parquet")).unwrap();
2717 let mut writer =
2718 ArrowWriter::try_new(file, to_write.schema(), Some(props.clone())).unwrap();
2719 writer.write(&to_write).expect("Writing batch");
2720 writer.close().unwrap();
2721 }
2722
2723 let reader = ArrowReaderBuilder::new(file_io, Runtime::current())
2725 .with_data_file_concurrency_limit(1)
2726 .build();
2727
2728 let tasks = vec![
2730 Ok(FileScanTask::builder()
2731 .with_file_size_in_bytes(
2732 std::fs::metadata(format!("{table_location}/file_0.parquet"))
2733 .unwrap()
2734 .len(),
2735 )
2736 .with_start(0)
2737 .with_length(0)
2738 .with_data_file_path(format!("{table_location}/file_0.parquet"))
2739 .with_data_file_format(DataFileFormat::Parquet)
2740 .with_schema(schema.clone())
2741 .with_project_field_ids(vec![1, 2])
2742 .with_case_sensitive(false)
2743 .build()
2744 .unwrap()),
2745 Ok(FileScanTask::builder()
2746 .with_file_size_in_bytes(
2747 std::fs::metadata(format!("{table_location}/file_1.parquet"))
2748 .unwrap()
2749 .len(),
2750 )
2751 .with_start(0)
2752 .with_length(0)
2753 .with_data_file_path(format!("{table_location}/file_1.parquet"))
2754 .with_data_file_format(DataFileFormat::Parquet)
2755 .with_schema(schema.clone())
2756 .with_project_field_ids(vec![1, 2])
2757 .with_case_sensitive(false)
2758 .build()
2759 .unwrap()),
2760 Ok(FileScanTask::builder()
2761 .with_file_size_in_bytes(
2762 std::fs::metadata(format!("{table_location}/file_2.parquet"))
2763 .unwrap()
2764 .len(),
2765 )
2766 .with_start(0)
2767 .with_length(0)
2768 .with_data_file_path(format!("{table_location}/file_2.parquet"))
2769 .with_data_file_format(DataFileFormat::Parquet)
2770 .with_schema(schema.clone())
2771 .with_project_field_ids(vec![1, 2])
2772 .with_case_sensitive(false)
2773 .build()
2774 .unwrap()),
2775 ];
2776
2777 let tasks_stream = Box::pin(futures::stream::iter(tasks)) as FileScanTaskStream;
2778
2779 let result = reader
2780 .read(tasks_stream)
2781 .unwrap()
2782 .stream()
2783 .try_collect::<Vec<RecordBatch>>()
2784 .await
2785 .unwrap();
2786
2787 let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
2789 assert_eq!(total_rows, 30, "Should have 30 total rows");
2790
2791 let mut all_ids = Vec::new();
2793 let mut all_file_nums = Vec::new();
2794
2795 for batch in &result {
2796 let id_col = batch
2797 .column(0)
2798 .as_primitive::<arrow_array::types::Int32Type>();
2799 let file_num_col = batch
2800 .column(1)
2801 .as_primitive::<arrow_array::types::Int32Type>();
2802
2803 for i in 0..batch.num_rows() {
2804 all_ids.push(id_col.value(i));
2805 all_file_nums.push(file_num_col.value(i));
2806 }
2807 }
2808
2809 assert_eq!(all_ids.len(), 30);
2810 assert_eq!(all_file_nums.len(), 30);
2811
2812 for i in 0..10 {
2817 assert_eq!(all_file_nums[i], 0, "First 10 rows should be from file_0");
2818 assert_eq!(all_ids[i], i as i32, "IDs should be 0-9");
2819 }
2820 for i in 10..20 {
2821 assert_eq!(all_file_nums[i], 1, "Next 10 rows should be from file_1");
2822 assert_eq!(all_ids[i], i as i32, "IDs should be 10-19");
2823 }
2824 for i in 20..30 {
2825 assert_eq!(all_file_nums[i], 2, "Last 10 rows should be from file_2");
2826 assert_eq!(all_ids[i], i as i32, "IDs should be 20-29");
2827 }
2828 }
2829
2830 #[tokio::test]
2831 async fn test_read_int96_timestamps_with_field_ids() {
2832 let schema = Arc::new(
2833 Schema::builder()
2834 .with_schema_id(1)
2835 .with_fields(vec![
2836 NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2837 .into(),
2838 NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2839 ])
2840 .build()
2841 .unwrap(),
2842 );
2843
2844 let tmp_dir = TempDir::new().unwrap();
2845 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2846 let (file_path, expected_micros) =
2847 write_int96_parquet_file(&table_location, "with_ids.parquet", true);
2848
2849 assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await;
2850 }
2851
2852 #[tokio::test]
2853 async fn test_read_int96_timestamps_without_field_ids() {
2854 let schema = Arc::new(
2855 Schema::builder()
2856 .with_schema_id(1)
2857 .with_fields(vec![
2858 NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2859 .into(),
2860 NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2861 ])
2862 .build()
2863 .unwrap(),
2864 );
2865
2866 let tmp_dir = TempDir::new().unwrap();
2867 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2868 let (file_path, expected_micros) =
2869 write_int96_parquet_file(&table_location, "no_ids.parquet", false);
2870
2871 assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await;
2872 }
2873
2874 #[tokio::test]
2875 async fn test_read_int96_timestamps_with_fallback_ids_and_pos() {
2876 use arrow_array::TimestampMicrosecondArray;
2877
2878 let schema = Arc::new(
2883 Schema::builder()
2884 .with_schema_id(1)
2885 .with_fields(vec![
2886 NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2887 .into(),
2888 NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2889 ])
2890 .build()
2891 .unwrap(),
2892 );
2893
2894 let tmp_dir = TempDir::new().unwrap();
2895 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2896 let (file_path, expected_micros) =
2897 write_int96_parquet_file(&table_location, "no_ids_with_pos.parquet", false);
2898
2899 let batches =
2900 read_int96_batches(&file_path, schema, vec![1, 2, RESERVED_FIELD_ID_POS]).await;
2901
2902 assert_eq!(batches.len(), 1);
2903 let ts_col = batches[0]
2905 .column_by_name("ts")
2906 .expect("ts column should be present")
2907 .as_any()
2908 .downcast_ref::<TimestampMicrosecondArray>()
2909 .expect("Expected TimestampMicrosecondArray");
2910 for (i, expected) in expected_micros.iter().enumerate() {
2911 assert_eq!(ts_col.value(i), *expected, "Row {i}");
2912 }
2913 let pos_col = batches[0]
2915 .column_by_name(RESERVED_COL_NAME_POS)
2916 .expect("_pos column should be present")
2917 .as_primitive::<arrow_array::types::Int64Type>();
2918 assert_eq!(pos_col.values(), &[0, 1, 2]);
2919 }
2920
2921 #[tokio::test]
2922 async fn test_read_int96_timestamps_with_field_ids_and_pos() {
2923 use arrow_array::TimestampMicrosecondArray;
2924
2925 let schema = Arc::new(
2926 Schema::builder()
2927 .with_schema_id(1)
2928 .with_fields(vec![
2929 NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2930 .into(),
2931 NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2932 ])
2933 .build()
2934 .unwrap(),
2935 );
2936
2937 let tmp_dir = TempDir::new().unwrap();
2938 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2939 let (file_path, expected_micros) =
2940 write_int96_parquet_file(&table_location, "with_ids_with_pos.parquet", true);
2941
2942 let batches =
2943 read_int96_batches(&file_path, schema, vec![1, 2, RESERVED_FIELD_ID_POS]).await;
2944
2945 assert_eq!(batches.len(), 1);
2946 let ts_col = batches[0]
2947 .column_by_name("ts")
2948 .expect("ts column should be present")
2949 .as_any()
2950 .downcast_ref::<TimestampMicrosecondArray>()
2951 .expect("Expected TimestampMicrosecondArray");
2952
2953 for (i, expected) in expected_micros.iter().enumerate() {
2954 assert_eq!(ts_col.value(i), *expected, "Row {i}");
2955 }
2956
2957 let pos_col = batches[0]
2958 .column_by_name(RESERVED_COL_NAME_POS)
2959 .expect("_pos column should be present")
2960 .as_primitive::<arrow_array::types::Int64Type>();
2961
2962 assert_eq!(pos_col.values(), &[0, 1, 2]);
2963 }
2964
2965 #[tokio::test]
2966 async fn test_read_int96_timestamps_in_struct() {
2967 use arrow_array::{StructArray, TimestampMicrosecondArray};
2968 use parquet::basic::{Repetition, Type as PhysicalType};
2969 use parquet::data_type::Int96Type;
2970 use parquet::file::writer::SerializedFileWriter;
2971 use parquet::schema::types::Type as SchemaType;
2972
2973 let tmp_dir = TempDir::new().unwrap();
2974 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2975 let file_path = format!("{table_location}/struct_int96.parquet");
2976
2977 let ts_type = SchemaType::primitive_type_builder("ts", PhysicalType::INT96)
2978 .with_repetition(Repetition::OPTIONAL)
2979 .with_id(Some(2))
2980 .build()
2981 .unwrap();
2982
2983 let struct_type = SchemaType::group_type_builder("data")
2984 .with_repetition(Repetition::REQUIRED)
2985 .with_id(Some(1))
2986 .with_fields(vec![Arc::new(ts_type)])
2987 .build()
2988 .unwrap();
2989
2990 let parquet_schema = SchemaType::group_type_builder("schema")
2991 .with_fields(vec![Arc::new(struct_type)])
2992 .build()
2993 .unwrap();
2994
2995 let (int96_val, expected_micros) = make_int96_test_value();
2996
2997 let file = File::create(&file_path).unwrap();
2998 let mut writer =
2999 SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3000
3001 let mut row_group = writer.next_row_group().unwrap();
3004 {
3005 let mut col = row_group.next_column().unwrap().unwrap();
3006 col.typed::<Int96Type>()
3007 .write_batch(&[int96_val], Some(&[1]), None)
3008 .unwrap();
3009 col.close().unwrap();
3010 }
3011 row_group.close().unwrap();
3012 writer.close().unwrap();
3013
3014 let iceberg_schema = Arc::new(
3015 Schema::builder()
3016 .with_schema_id(1)
3017 .with_fields(vec![
3018 NestedField::required(
3019 1,
3020 "data",
3021 Type::Struct(crate::spec::StructType::new(vec![
3022 NestedField::optional(
3023 2,
3024 "ts",
3025 Type::Primitive(PrimitiveType::Timestamp),
3026 )
3027 .into(),
3028 ])),
3029 )
3030 .into(),
3031 ])
3032 .build()
3033 .unwrap(),
3034 );
3035
3036 let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3037
3038 assert_eq!(batches.len(), 1);
3039 let struct_array = batches[0]
3040 .column(0)
3041 .as_any()
3042 .downcast_ref::<StructArray>()
3043 .expect("Expected StructArray");
3044 let ts_array = struct_array
3045 .column(0)
3046 .as_any()
3047 .downcast_ref::<TimestampMicrosecondArray>()
3048 .expect("Expected TimestampMicrosecondArray inside struct");
3049
3050 assert_eq!(
3051 ts_array.value(0),
3052 expected_micros,
3053 "INT96 in struct: got {}, expected {expected_micros}",
3054 ts_array.value(0)
3055 );
3056 }
3057
3058 #[tokio::test]
3059 async fn test_read_int96_timestamps_in_list() {
3060 use arrow_array::{ListArray, TimestampMicrosecondArray};
3061 use parquet::basic::{Repetition, Type as PhysicalType};
3062 use parquet::data_type::Int96Type;
3063 use parquet::file::writer::SerializedFileWriter;
3064 use parquet::schema::types::Type as SchemaType;
3065
3066 let tmp_dir = TempDir::new().unwrap();
3067 let table_location = tmp_dir.path().to_str().unwrap().to_string();
3068 let file_path = format!("{table_location}/list_int96.parquet");
3069
3070 let element_type = SchemaType::primitive_type_builder("element", PhysicalType::INT96)
3077 .with_repetition(Repetition::OPTIONAL)
3078 .with_id(Some(2))
3079 .build()
3080 .unwrap();
3081
3082 let list_group = SchemaType::group_type_builder("list")
3083 .with_repetition(Repetition::REPEATED)
3084 .with_fields(vec![Arc::new(element_type)])
3085 .build()
3086 .unwrap();
3087
3088 let list_type = SchemaType::group_type_builder("timestamps")
3089 .with_repetition(Repetition::OPTIONAL)
3090 .with_id(Some(1))
3091 .with_logical_type(Some(parquet::basic::LogicalType::List))
3092 .with_fields(vec![Arc::new(list_group)])
3093 .build()
3094 .unwrap();
3095
3096 let parquet_schema = SchemaType::group_type_builder("schema")
3097 .with_fields(vec![Arc::new(list_type)])
3098 .build()
3099 .unwrap();
3100
3101 let (int96_val, expected_micros) = make_int96_test_value();
3102
3103 let file = File::create(&file_path).unwrap();
3104 let mut writer =
3105 SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3106
3107 let mut row_group = writer.next_row_group().unwrap();
3111 {
3112 let mut col = row_group.next_column().unwrap().unwrap();
3113 col.typed::<Int96Type>()
3114 .write_batch(&[int96_val], Some(&[3]), Some(&[0]))
3115 .unwrap();
3116 col.close().unwrap();
3117 }
3118 row_group.close().unwrap();
3119 writer.close().unwrap();
3120
3121 let iceberg_schema = Arc::new(
3122 Schema::builder()
3123 .with_schema_id(1)
3124 .with_fields(vec![
3125 NestedField::optional(
3126 1,
3127 "timestamps",
3128 Type::List(crate::spec::ListType {
3129 element_field: NestedField::optional(
3130 2,
3131 "element",
3132 Type::Primitive(PrimitiveType::Timestamp),
3133 )
3134 .into(),
3135 }),
3136 )
3137 .into(),
3138 ])
3139 .build()
3140 .unwrap(),
3141 );
3142
3143 let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3144
3145 assert_eq!(batches.len(), 1);
3146 let list_array = batches[0]
3147 .column(0)
3148 .as_any()
3149 .downcast_ref::<ListArray>()
3150 .expect("Expected ListArray");
3151 let ts_array = list_array
3152 .values()
3153 .as_any()
3154 .downcast_ref::<TimestampMicrosecondArray>()
3155 .expect("Expected TimestampMicrosecondArray inside list");
3156
3157 assert_eq!(
3158 ts_array.value(0),
3159 expected_micros,
3160 "INT96 in list: got {}, expected {expected_micros}",
3161 ts_array.value(0)
3162 );
3163 }
3164
3165 #[tokio::test]
3166 async fn test_read_int96_timestamps_in_map() {
3167 use arrow_array::{MapArray, TimestampMicrosecondArray};
3168 use parquet::basic::{Repetition, Type as PhysicalType};
3169 use parquet::data_type::{ByteArrayType, Int96Type};
3170 use parquet::file::writer::SerializedFileWriter;
3171 use parquet::schema::types::Type as SchemaType;
3172
3173 let tmp_dir = TempDir::new().unwrap();
3174 let table_location = tmp_dir.path().to_str().unwrap().to_string();
3175 let file_path = format!("{table_location}/map_int96.parquet");
3176
3177 let key_type = SchemaType::primitive_type_builder("key", PhysicalType::BYTE_ARRAY)
3185 .with_repetition(Repetition::REQUIRED)
3186 .with_logical_type(Some(parquet::basic::LogicalType::String))
3187 .with_id(Some(2))
3188 .build()
3189 .unwrap();
3190
3191 let value_type = SchemaType::primitive_type_builder("value", PhysicalType::INT96)
3192 .with_repetition(Repetition::OPTIONAL)
3193 .with_id(Some(3))
3194 .build()
3195 .unwrap();
3196
3197 let key_value_group = SchemaType::group_type_builder("key_value")
3198 .with_repetition(Repetition::REPEATED)
3199 .with_fields(vec![Arc::new(key_type), Arc::new(value_type)])
3200 .build()
3201 .unwrap();
3202
3203 let map_type = SchemaType::group_type_builder("ts_map")
3204 .with_repetition(Repetition::OPTIONAL)
3205 .with_id(Some(1))
3206 .with_logical_type(Some(parquet::basic::LogicalType::Map))
3207 .with_fields(vec![Arc::new(key_value_group)])
3208 .build()
3209 .unwrap();
3210
3211 let parquet_schema = SchemaType::group_type_builder("schema")
3212 .with_fields(vec![Arc::new(map_type)])
3213 .build()
3214 .unwrap();
3215
3216 let (int96_val, expected_micros) = make_int96_test_value();
3217
3218 let file = File::create(&file_path).unwrap();
3219 let mut writer =
3220 SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3221
3222 let mut row_group = writer.next_row_group().unwrap();
3227 {
3228 let mut col = row_group.next_column().unwrap().unwrap();
3229 col.typed::<ByteArrayType>()
3230 .write_batch(
3231 &[parquet::data_type::ByteArray::from("event_time")],
3232 Some(&[2]),
3233 Some(&[0]),
3234 )
3235 .unwrap();
3236 col.close().unwrap();
3237 }
3238 {
3239 let mut col = row_group.next_column().unwrap().unwrap();
3240 col.typed::<Int96Type>()
3241 .write_batch(&[int96_val], Some(&[3]), Some(&[0]))
3242 .unwrap();
3243 col.close().unwrap();
3244 }
3245 row_group.close().unwrap();
3246 writer.close().unwrap();
3247
3248 let iceberg_schema = Arc::new(
3249 Schema::builder()
3250 .with_schema_id(1)
3251 .with_fields(vec![
3252 NestedField::optional(
3253 1,
3254 "ts_map",
3255 Type::Map(crate::spec::MapType {
3256 key_field: NestedField::required(
3257 2,
3258 "key",
3259 Type::Primitive(PrimitiveType::String),
3260 )
3261 .into(),
3262 value_field: NestedField::optional(
3263 3,
3264 "value",
3265 Type::Primitive(PrimitiveType::Timestamp),
3266 )
3267 .into(),
3268 }),
3269 )
3270 .into(),
3271 ])
3272 .build()
3273 .unwrap(),
3274 );
3275
3276 let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3277
3278 assert_eq!(batches.len(), 1);
3279 let map_array = batches[0]
3280 .column(0)
3281 .as_any()
3282 .downcast_ref::<MapArray>()
3283 .expect("Expected MapArray");
3284 let ts_array = map_array
3285 .values()
3286 .as_any()
3287 .downcast_ref::<TimestampMicrosecondArray>()
3288 .expect("Expected TimestampMicrosecondArray as map values");
3289
3290 assert_eq!(
3291 ts_array.value(0),
3292 expected_micros,
3293 "INT96 in map: got {}, expected {expected_micros}",
3294 ts_array.value(0)
3295 );
3296 }
3297
3298 fn write_parquet_with_wide_column(
3304 dir: &str,
3305 name: &str,
3306 extra_fields: Vec<Field>,
3307 extra_columns: Vec<ArrayRef>,
3308 ) -> String {
3309 let wide_field =
3310 Field::new("wide", DataType::Utf8, false).with_metadata(HashMap::from([(
3311 PARQUET_FIELD_ID_META_KEY.to_string(),
3312 "2".to_string(),
3313 )]));
3314 let wide_values: Vec<String> = (0..3)
3317 .map(|i| {
3318 (0..2048)
3319 .map(|j| ((i * 2048 + j) % 251) as u8 as char)
3320 .collect()
3321 })
3322 .collect();
3323
3324 let mut fields = vec![wide_field];
3325 fields.extend(extra_fields);
3326 let mut columns: Vec<ArrayRef> = vec![Arc::new(StringArray::from(wide_values))];
3327 columns.extend(extra_columns);
3328 write_plain_parquet(dir, name, fields, columns)
3329 }
3330
3331 fn id_and_wide_schema() -> SchemaRef {
3334 Arc::new(
3335 Schema::builder()
3336 .with_schema_id(1)
3337 .with_fields(vec![
3338 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3339 NestedField::required(2, "wide", Type::Primitive(PrimitiveType::String)).into(),
3340 ])
3341 .build()
3342 .unwrap(),
3343 )
3344 }
3345
3346 fn metadata_projection_task(
3348 file_path: String,
3349 schema: SchemaRef,
3350 project_field_ids: Vec<i32>,
3351 ) -> FileScanTask {
3352 metadata_projection_task_with_first_row_id(file_path, schema, project_field_ids, None)
3353 }
3354
3355 fn metadata_projection_task_with_first_row_id(
3356 file_path: String,
3357 schema: SchemaRef,
3358 project_field_ids: Vec<i32>,
3359 first_row_id: Option<i64>,
3360 ) -> FileScanTask {
3361 FileScanTask::builder()
3362 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3363 .with_start(0)
3364 .with_length(0)
3365 .with_data_file_path(file_path)
3366 .with_data_file_format(DataFileFormat::Parquet)
3367 .with_schema(schema)
3368 .with_project_field_ids(project_field_ids)
3369 .with_first_row_id(first_row_id)
3370 .with_case_sensitive(false)
3371 .build()
3372 .unwrap()
3373 }
3374
3375 async fn scan_task(task: FileScanTask) -> (Vec<RecordBatch>, u64) {
3377 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
3378 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
3379 let scan = reader.read(tasks).unwrap();
3380 let metrics = scan.metrics().clone();
3381 let batches = scan.stream().try_collect().await.unwrap();
3382 (batches, metrics.bytes_read())
3383 }
3384
3385 #[tokio::test]
3386 async fn test_pos_only_projection_reads_no_data_columns() {
3387 let tmp_dir = TempDir::new().unwrap();
3388 let dir = tmp_dir.path().to_str().unwrap();
3389
3390 let pos_only = metadata_projection_task(
3391 write_parquet_with_wide_column(dir, "pos_only.parquet", vec![], vec![]),
3392 id_and_wide_schema(),
3393 vec![RESERVED_FIELD_ID_POS],
3394 );
3395 let (batches, pos_only_bytes) = scan_task(pos_only).await;
3396
3397 assert_eq!(batches[0].num_columns(), 1);
3399 let pos_col = batches[0]
3400 .column_by_name(RESERVED_COL_NAME_POS)
3401 .expect("_pos column should be present")
3402 .as_primitive::<arrow_array::types::Int64Type>();
3403 assert_eq!(pos_col.values(), &[0, 1, 2]);
3404
3405 let with_data = metadata_projection_task(
3408 write_parquet_with_wide_column(dir, "pos_only_ref.parquet", vec![], vec![]),
3409 id_and_wide_schema(),
3410 vec![2, RESERVED_FIELD_ID_POS],
3411 );
3412 let (_, with_data_bytes) = scan_task(with_data).await;
3413
3414 assert!(
3415 pos_only_bytes < with_data_bytes,
3416 "_pos-only scan should read fewer bytes than a scan of the wide column: \
3417 {pos_only_bytes} vs {with_data_bytes}"
3418 );
3419 }
3420
3421 #[tokio::test]
3422 async fn test_empty_projection_reads_no_data_columns() {
3423 let tmp_dir = TempDir::new().unwrap();
3424 let dir = tmp_dir.path().to_str().unwrap();
3425
3426 let empty = metadata_projection_task(
3427 write_parquet_with_wide_column(dir, "empty.parquet", vec![], vec![]),
3428 id_and_wide_schema(),
3429 vec![],
3430 );
3431 let (batches, empty_bytes) = scan_task(empty).await;
3432
3433 assert!(batches.iter().all(|b| b.num_columns() == 0));
3435 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
3436
3437 let with_data = metadata_projection_task(
3438 write_parquet_with_wide_column(dir, "empty_ref.parquet", vec![], vec![]),
3439 id_and_wide_schema(),
3440 vec![2],
3441 );
3442 let (_, with_data_bytes) = scan_task(with_data).await;
3443
3444 assert!(
3445 empty_bytes < with_data_bytes,
3446 "empty projection should read fewer bytes than the wide column: \
3447 {empty_bytes} vs {with_data_bytes}"
3448 );
3449 }
3450
3451 #[tokio::test]
3452 async fn test_empty_projection_counts_rows_after_positional_delete() {
3453 use crate::spec::DataContentType;
3454
3455 let tmp_dir = TempDir::new().unwrap();
3456 let dir = tmp_dir.path().to_str().unwrap();
3457 let data_path = write_plain_parquet(dir, "empty_posdel_data.parquet", vec![], vec![]);
3458 let del_path = write_positional_delete(dir, "empty_posdel.parquet", &data_path, &[1]);
3459 let delete = FileScanTaskDeleteFile::builder()
3460 .with_file_path(del_path.clone())
3461 .with_file_size_in_bytes(std::fs::metadata(&del_path).unwrap().len())
3462 .with_file_type(DataContentType::PositionDeletes)
3463 .with_file_format(DataFileFormat::Parquet)
3464 .with_partition_spec_id(0)
3465 .build()
3466 .unwrap();
3467 let task = FileScanTask::builder()
3468 .with_file_size_in_bytes(std::fs::metadata(&data_path).unwrap().len())
3469 .with_start(0)
3470 .with_length(0)
3471 .with_data_file_path(data_path)
3472 .with_data_file_format(DataFileFormat::Parquet)
3473 .with_schema(id_and_wide_schema())
3474 .with_project_field_ids(vec![])
3475 .with_deletes(vec![delete])
3476 .with_case_sensitive(false)
3477 .build()
3478 .unwrap();
3479
3480 let (batches, _) = scan_task(task).await;
3481
3482 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
3483 }
3484
3485 #[tokio::test]
3486 async fn test_empty_projection_counts_rows_after_predicate() {
3487 use crate::expr::{Bind, Reference};
3488 use crate::spec::Datum;
3489
3490 let tmp_dir = TempDir::new().unwrap();
3491 let dir = tmp_dir.path().to_str().unwrap();
3492 let file_path = write_plain_parquet(dir, "empty_predicate.parquet", vec![], vec![]);
3493 let schema = id_and_wide_schema();
3494 let bound = Reference::new("id")
3496 .not_equal_to(Datum::int(2))
3497 .bind(Arc::clone(&schema), false)
3498 .unwrap();
3499 let task = FileScanTask::builder()
3500 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3501 .with_start(0)
3502 .with_length(0)
3503 .with_data_file_path(file_path)
3504 .with_data_file_format(DataFileFormat::Parquet)
3505 .with_schema(schema)
3506 .with_project_field_ids(vec![])
3507 .with_predicate(Some(bound))
3508 .with_case_sensitive(false)
3509 .build()
3510 .unwrap();
3511
3512 let (batches, _) = scan_task(task).await;
3513
3514 assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
3515 }
3516
3517 #[tokio::test]
3518 async fn test_pos_only_projection_keeps_absolute_pos_under_predicate() {
3519 use crate::expr::{Bind, Reference};
3520 use crate::spec::Datum;
3521
3522 let tmp_dir = TempDir::new().unwrap();
3523 let dir = tmp_dir.path().to_str().unwrap();
3524 let file_path = write_plain_parquet(dir, "pos_only_predicate.parquet", vec![], vec![]);
3526
3527 let schema = Arc::new(
3528 Schema::builder()
3529 .with_schema_id(1)
3530 .with_fields(vec![
3531 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3532 ])
3533 .build()
3534 .unwrap(),
3535 );
3536 let bound = Reference::new("id")
3537 .not_equal_to(Datum::int(2))
3538 .bind(Arc::clone(&schema), false)
3539 .unwrap();
3540 let task = FileScanTask::builder()
3541 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3542 .with_start(0)
3543 .with_length(0)
3544 .with_data_file_path(file_path)
3545 .with_data_file_format(DataFileFormat::Parquet)
3546 .with_schema(schema)
3547 .with_project_field_ids(vec![RESERVED_FIELD_ID_POS])
3548 .with_predicate(Some(bound))
3549 .with_case_sensitive(false)
3550 .build()
3551 .unwrap();
3552
3553 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
3557 .with_row_selection_enabled(true)
3558 .build();
3559 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
3560 let batches: Vec<RecordBatch> = reader
3561 .read(tasks)
3562 .unwrap()
3563 .stream()
3564 .try_collect()
3565 .await
3566 .unwrap();
3567
3568 let pos: Vec<i64> = batches
3569 .iter()
3570 .flat_map(|b| {
3571 b.column_by_name(RESERVED_COL_NAME_POS)
3572 .expect("_pos column should be present")
3573 .as_primitive::<arrow_array::types::Int64Type>()
3574 .values()
3575 .to_vec()
3576 })
3577 .collect();
3578 assert_eq!(pos, vec![0, 2]);
3579 }
3580
3581 #[tokio::test]
3582 async fn test_pos_and_file_projection() {
3583 use crate::metadata_columns::RESERVED_COL_NAME_FILE;
3584
3585 let tmp_dir = TempDir::new().unwrap();
3586 let dir = tmp_dir.path().to_str().unwrap();
3587 let file_path = write_parquet_with_wide_column(dir, "pos_and_file.parquet", vec![], vec![]);
3590 let task = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![
3591 RESERVED_FIELD_ID_POS,
3592 RESERVED_FIELD_ID_FILE,
3593 ]);
3594 let (batches, _) = scan_task(task).await;
3595
3596 assert_eq!(batches[0].num_columns(), 2);
3598 let pos_col = batches[0]
3599 .column_by_name(RESERVED_COL_NAME_POS)
3600 .expect("_pos column should be present")
3601 .as_primitive::<arrow_array::types::Int64Type>();
3602 assert_eq!(pos_col.values(), &[0, 1, 2]);
3603 let file_col = batches[0]
3604 .column_by_name(RESERVED_COL_NAME_FILE)
3605 .expect("_file column should be present");
3606 let file_col = cast(file_col, &DataType::Utf8).unwrap();
3607 let file_col = file_col.as_any().downcast_ref::<StringArray>().unwrap();
3608 assert_eq!(file_col.value(0), file_path);
3609 }
3610
3611 #[tokio::test]
3612 async fn test_pos_and_physical_seq_projection_reads_only_the_leaf() {
3613 use crate::metadata_columns::{
3614 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3615 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3616 };
3617
3618 let tmp_dir = TempDir::new().unwrap();
3622 let dir = tmp_dir.path().to_str().unwrap();
3623
3624 let write = |name: &str| {
3626 write_parquet_with_wide_column(
3627 dir,
3628 name,
3629 vec![physical_last_updated_seq_field()],
3630 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
3631 )
3632 };
3633
3634 let seq_task = |path: String, ids: Vec<i32>| {
3635 FileScanTask::builder()
3636 .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len())
3637 .with_start(0)
3638 .with_length(0)
3639 .with_data_file_path(path)
3640 .with_data_file_format(DataFileFormat::Parquet)
3641 .with_schema(id_and_wide_schema())
3642 .with_project_field_ids(ids)
3643 .with_first_row_id(Some(100))
3644 .with_data_sequence_number(Some(9))
3645 .with_case_sensitive(false)
3646 .build()
3647 .unwrap()
3648 };
3649
3650 let meta_only = seq_task(write("pos_seq.parquet"), vec![
3651 RESERVED_FIELD_ID_POS,
3652 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3653 ]);
3654 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3655
3656 let pos_col = batches[0]
3658 .column_by_name(RESERVED_COL_NAME_POS)
3659 .expect("_pos column should be present")
3660 .as_primitive::<arrow_array::types::Int64Type>();
3661 assert_eq!(pos_col.values(), &[0, 1, 2]);
3662 let seq_col = batches[0]
3663 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3664 .expect("_last_updated_sequence_number column should be present");
3665 let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3666 let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3667 assert_eq!(seq_col.value(0), 5);
3669 assert_eq!(seq_col.value(1), 9);
3670 assert_eq!(seq_col.value(2), 8);
3671 assert!(batches[0].column_by_name("wide").is_none());
3672
3673 let with_data = seq_task(write("pos_seq_ref.parquet"), vec![
3676 2,
3677 RESERVED_FIELD_ID_POS,
3678 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3679 ]);
3680 let (_, with_data_bytes) = scan_task(with_data).await;
3681
3682 assert!(
3683 meta_only_bytes < with_data_bytes,
3684 "_pos + physical sequence scan should read fewer bytes than one that also \
3685 reads the wide column: {meta_only_bytes} vs {with_data_bytes}"
3686 );
3687 }
3688
3689 #[tokio::test]
3690 async fn test_seq_only_projection_reads_only_the_leaf() {
3691 use crate::metadata_columns::{
3692 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3693 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3694 };
3695
3696 let tmp_dir = TempDir::new().unwrap();
3700 let dir = tmp_dir.path().to_str().unwrap();
3701 let write = |name: &str| {
3702 write_parquet_with_wide_column(
3703 dir,
3704 name,
3705 vec![physical_last_updated_seq_field()],
3706 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
3707 )
3708 };
3709 let seq_task = |path: String, ids: Vec<i32>| {
3710 FileScanTask::builder()
3711 .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len())
3712 .with_start(0)
3713 .with_length(0)
3714 .with_data_file_path(path)
3715 .with_data_file_format(DataFileFormat::Parquet)
3716 .with_schema(id_and_wide_schema())
3717 .with_project_field_ids(ids)
3718 .with_first_row_id(Some(100))
3719 .with_data_sequence_number(Some(9))
3720 .with_case_sensitive(false)
3721 .build()
3722 .unwrap()
3723 };
3724
3725 let meta_only = seq_task(write("seq_only.parquet"), vec![
3726 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3727 ]);
3728 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3729
3730 let seq_col = batches[0]
3731 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3732 .expect("_last_updated_sequence_number column should be present");
3733 let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3734 let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3735 assert_eq!(seq_col.value(0), 5);
3736 assert_eq!(seq_col.value(1), 9);
3737 assert_eq!(seq_col.value(2), 8);
3738 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3739 assert_eq!(total_rows, 3);
3740 assert!(batches[0].column_by_name("wide").is_none());
3741
3742 let with_data = seq_task(write("seq_only_ref.parquet"), vec![
3743 2,
3744 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3745 ]);
3746 let (_, with_data_bytes) = scan_task(with_data).await;
3747
3748 assert!(
3749 meta_only_bytes < with_data_bytes,
3750 "seq-only scan should read fewer bytes than one that also reads the wide \
3751 column: {meta_only_bytes} vs {with_data_bytes}"
3752 );
3753 }
3754
3755 #[tokio::test]
3756 async fn test_seq_only_null_first_row_id_reads_no_data_columns() {
3757 use crate::metadata_columns::{
3758 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3759 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3760 };
3761
3762 let tmp_dir = TempDir::new().unwrap();
3767 let dir = tmp_dir.path().to_str().unwrap();
3768 let seq_task = |name: &str, ids: Vec<i32>| {
3769 let file_path = write_parquet_with_wide_column(
3770 dir,
3771 name,
3772 vec![physical_last_updated_seq_field()],
3773 vec![Arc::new(Int64Array::from(vec![Some(5), Some(6), Some(7)])) as ArrayRef],
3774 );
3775 FileScanTask::builder()
3776 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3777 .with_start(0)
3778 .with_length(0)
3779 .with_data_file_path(file_path)
3780 .with_data_file_format(DataFileFormat::Parquet)
3781 .with_schema(id_and_wide_schema())
3782 .with_project_field_ids(ids)
3783 .with_first_row_id(None)
3784 .with_data_sequence_number(Some(9))
3785 .with_case_sensitive(false)
3786 .build()
3787 .unwrap()
3788 };
3789
3790 let meta_only = seq_task("seq_only_null_first.parquet", vec![
3791 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3792 ]);
3793 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3794
3795 assert_eq!(batches[0].num_columns(), 1);
3796 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3797 assert_eq!(total_rows, 3);
3798 let seq_col = batches[0]
3799 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3800 .expect("_last_updated_sequence_number column should be present");
3801 let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3802 let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3803 assert!((0..3).all(|i| seq_col.is_null(i)));
3804
3805 let with_data = seq_task("seq_only_null_first_ref.parquet", vec![
3806 2,
3807 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3808 ]);
3809 let (_, with_data_bytes) = scan_task(with_data).await;
3810
3811 assert!(
3812 meta_only_bytes < with_data_bytes,
3813 "seq-only scan with a null first_row_id should read fewer bytes than a scan of \
3814 the wide column: {meta_only_bytes} vs {with_data_bytes}"
3815 );
3816 }
3817
3818 #[tokio::test]
3819 async fn test_file_only_reads_no_data_columns() {
3820 use crate::metadata_columns::RESERVED_COL_NAME_FILE;
3821
3822 let tmp_dir = TempDir::new().unwrap();
3826 let dir = tmp_dir.path().to_str().unwrap();
3827
3828 let file_path = write_parquet_with_wide_column(dir, "file_only.parquet", vec![], vec![]);
3829 let meta_only = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![
3830 RESERVED_FIELD_ID_FILE,
3831 ]);
3832 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3833
3834 assert_eq!(batches[0].num_columns(), 1);
3835 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3836 assert_eq!(total_rows, 3);
3837 let file_col = batches[0]
3838 .column_by_name(RESERVED_COL_NAME_FILE)
3839 .expect("_file column should be present");
3840 let file_col = cast(file_col, &DataType::Utf8).unwrap();
3841 let file_col = file_col.as_any().downcast_ref::<StringArray>().unwrap();
3842 assert_eq!(file_col.value(0), file_path);
3843
3844 let with_data = metadata_projection_task(
3845 write_parquet_with_wide_column(dir, "file_only_ref.parquet", vec![], vec![]),
3846 id_and_wide_schema(),
3847 vec![2, RESERVED_FIELD_ID_FILE],
3848 );
3849 let (_, with_data_bytes) = scan_task(with_data).await;
3850
3851 assert!(
3852 meta_only_bytes < with_data_bytes,
3853 "_file-only scan should read fewer bytes than a scan of the wide column: \
3854 {meta_only_bytes} vs {with_data_bytes}"
3855 );
3856 }
3857
3858 #[tokio::test]
3859 async fn test_spec_id_only_reads_no_data_columns() {
3860 use crate::metadata_columns::{RESERVED_COL_NAME_SPEC_ID, RESERVED_FIELD_ID_SPEC_ID};
3861 use crate::spec::{Literal, PartitionSpec, Struct, Transform};
3862
3863 let tmp_dir = TempDir::new().unwrap();
3867 let dir = tmp_dir.path().to_str().unwrap();
3868 let schema = id_and_wide_schema();
3869 let spec = Arc::new(
3870 PartitionSpec::builder(schema.clone())
3871 .with_spec_id(7)
3872 .add_partition_field("id", "id", Transform::Identity)
3873 .unwrap()
3874 .build()
3875 .unwrap(),
3876 );
3877
3878 let spec_id_task = |file_path: String, project_field_ids: Vec<i32>| {
3879 FileScanTask::builder()
3880 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3881 .with_start(0)
3882 .with_length(0)
3883 .with_data_file_path(file_path)
3884 .with_data_file_format(DataFileFormat::Parquet)
3885 .with_schema(schema.clone())
3886 .with_project_field_ids(project_field_ids)
3887 .with_partition(Some(Struct::from_iter([Some(Literal::int(42))])))
3888 .with_partition_spec(Some(spec.clone()))
3889 .with_case_sensitive(false)
3890 .build()
3891 .unwrap()
3892 };
3893
3894 let meta_only = spec_id_task(
3895 write_parquet_with_wide_column(dir, "spec_id_only.parquet", vec![], vec![]),
3896 vec![RESERVED_FIELD_ID_SPEC_ID],
3897 );
3898 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3899
3900 assert_eq!(batches[0].num_columns(), 1);
3901 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3902 assert_eq!(total_rows, 3);
3903 let spec_id_col = batches[0]
3904 .column_by_name(RESERVED_COL_NAME_SPEC_ID)
3905 .expect("_spec_id column should be present");
3906 let spec_id_col = cast(spec_id_col, &DataType::Int32).unwrap();
3907 let spec_id_col = spec_id_col.as_any().downcast_ref::<Int32Array>().unwrap();
3908 assert_eq!(spec_id_col.values(), &[7, 7, 7]);
3909
3910 let with_data = spec_id_task(
3911 write_parquet_with_wide_column(dir, "spec_id_only_ref.parquet", vec![], vec![]),
3912 vec![2, RESERVED_FIELD_ID_SPEC_ID],
3913 );
3914 let (_, with_data_bytes) = scan_task(with_data).await;
3915
3916 assert!(
3917 meta_only_bytes < with_data_bytes,
3918 "_spec_id-only scan should read fewer bytes than a scan of the wide column: \
3919 {meta_only_bytes} vs {with_data_bytes}"
3920 );
3921 }
3922
3923 #[tokio::test]
3924 async fn test_empty_projection_preserves_row_count() {
3925 let tmp_dir = TempDir::new().unwrap();
3926 let dir = tmp_dir.path().to_str().unwrap();
3927 let file_path = write_plain_parquet(dir, "empty_projection.parquet", vec![], vec![]);
3928 let schema = Arc::new(
3929 Schema::builder()
3930 .with_schema_id(1)
3931 .with_fields(vec![
3932 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3933 ])
3934 .build()
3935 .unwrap(),
3936 );
3937 let task = metadata_projection_task(file_path, schema, vec![]);
3938 let (batches, _) = scan_task(task).await;
3939
3940 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3942 assert_eq!(total_rows, 3);
3943 }
3944
3945 mod open_progress {
3948 use std::ops::Range;
3949 use std::sync::Mutex;
3950
3951 use async_trait::async_trait;
3952 use bytes::Bytes;
3953 use futures::StreamExt;
3954 use futures::stream::BoxStream;
3955 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
3956 use serde::{Deserialize, Serialize};
3957
3958 use super::*;
3959 use crate::Result;
3960 use crate::arrow::scan_metrics::ScanMetrics;
3961 use crate::expr::{Bind, Predicate, Reference};
3962 use crate::io::{
3963 FileIOBuilder, FileMetadata, FileRead, FileWrite, InputFile, MemoryStorage, OutputFile,
3964 Storage, StorageConfig, StorageFactory,
3965 };
3966 use crate::scan::ArrowRecordBatchStream;
3967 use crate::spec::Datum;
3968
3969 const CONCURRENCY: usize = 2;
3970 const METADATA_SIZE_HINT: usize = 64;
3972 const OPEN_BYTES: u64 = 547;
3976 const YIELDS: usize = 100;
3979
3980 #[derive(Debug, Default)]
3981 struct Reads {
3982 last_start: Option<u64>,
3983 started: usize,
3984 finished: usize,
3985 }
3986
3987 #[derive(Debug, Clone, Default)]
3991 struct ReadLog(Arc<Mutex<Reads>>);
3992
3993 impl ReadLog {
3994 fn last_start(&self) -> Option<u64> {
3996 self.0.lock().unwrap().last_start
3997 }
3998
3999 fn in_flight(&self) -> usize {
4000 let reads = self.0.lock().unwrap();
4001 reads.started - reads.finished
4002 }
4003 }
4004
4005 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
4007 struct DeferredStorage {
4008 #[serde(skip)]
4009 files: MemoryStorage,
4010 #[serde(skip)]
4011 log: ReadLog,
4012 }
4013
4014 struct DeferredRead {
4015 contents: Bytes,
4016 log: ReadLog,
4017 }
4018
4019 #[async_trait]
4020 impl FileRead for DeferredRead {
4021 async fn read(&self, range: Range<u64>) -> Result<Bytes> {
4022 {
4023 let mut reads = self.log.0.lock().unwrap();
4024 reads.last_start = Some(range.start);
4025 reads.started += 1;
4026 }
4027 tokio::task::yield_now().await;
4030 self.log.0.lock().unwrap().finished += 1;
4031 Ok(self
4032 .contents
4033 .slice(range.start as usize..range.end as usize))
4034 }
4035 }
4036
4037 #[async_trait]
4038 #[typetag::serde]
4039 impl Storage for DeferredStorage {
4040 async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
4041 Ok(Box::new(DeferredRead {
4042 contents: self.files.read(path).await?,
4043 log: self.log.clone(),
4044 }))
4045 }
4046
4047 fn new_input(&self, path: &str) -> Result<InputFile> {
4048 Ok(InputFile::new(Arc::new(self.clone()), path.to_string()))
4049 }
4050
4051 async fn exists(&self, path: &str) -> Result<bool> {
4052 self.files.exists(path).await
4053 }
4054
4055 async fn metadata(&self, path: &str) -> Result<FileMetadata> {
4056 self.files.metadata(path).await
4057 }
4058
4059 async fn read(&self, path: &str) -> Result<Bytes> {
4060 self.files.read(path).await
4061 }
4062
4063 async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
4064 self.files.write(path, bs).await
4065 }
4066
4067 async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
4068 self.files.writer(path).await
4069 }
4070
4071 async fn delete(&self, path: &str) -> Result<()> {
4072 self.files.delete(path).await
4073 }
4074
4075 async fn delete_prefix(&self, path: &str) -> Result<()> {
4076 self.files.delete_prefix(path).await
4077 }
4078
4079 async fn delete_stream(&self, paths: BoxStream<'static, String>) -> Result<()> {
4080 self.files.delete_stream(paths).await
4081 }
4082
4083 fn new_output(&self, path: &str) -> Result<OutputFile> {
4084 Ok(OutputFile::new(Arc::new(self.clone()), path.to_string()))
4085 }
4086 }
4087
4088 #[derive(Debug, Default, Serialize, Deserialize)]
4089 struct DeferredStorageFactory {
4090 #[serde(skip)]
4091 storage: DeferredStorage,
4092 }
4093
4094 #[typetag::serde]
4095 impl StorageFactory for DeferredStorageFactory {
4096 fn build(&self, _config: &StorageConfig) -> Result<Arc<dyn Storage>> {
4097 Ok(Arc::new(self.storage.clone()))
4098 }
4099 }
4100
4101 async fn deferred_scan_of(
4104 contents: Bytes,
4105 files: usize,
4106 predicate: Option<Predicate>,
4107 ) -> (ScanMetrics, ReadLog, ArrowRecordBatchStream) {
4108 let schema = Arc::new(
4109 Schema::builder()
4110 .with_schema_id(1)
4111 .with_fields(vec![
4112 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
4113 ])
4114 .build()
4115 .unwrap(),
4116 );
4117 let bound = predicate.map(|p| p.bind(Arc::clone(&schema), false).unwrap());
4118
4119 let storage = DeferredStorage::default();
4120 let log = storage.log.clone();
4121 let mut tasks = Vec::with_capacity(files);
4122 for file_num in 0..files {
4123 let path = format!("memory:/data/file_{file_num}.parquet");
4124 storage.write(&path, contents.clone()).await.unwrap();
4125 tasks.push(
4126 FileScanTask::builder()
4127 .with_data_file_path(path)
4128 .with_file_size_in_bytes(contents.len() as u64)
4129 .with_start(0)
4130 .with_length(0)
4131 .with_data_file_format(DataFileFormat::Parquet)
4132 .with_schema(Arc::clone(&schema))
4133 .with_project_field_ids(vec![1])
4134 .with_predicate(bound.clone())
4135 .with_case_sensitive(false)
4136 .build()
4137 .unwrap(),
4138 );
4139 }
4140
4141 let file_io = FileIOBuilder::new(Arc::new(DeferredStorageFactory { storage })).build();
4142
4143 let reader = ArrowReaderBuilder::new(file_io, Runtime::current())
4144 .with_data_file_concurrency_limit(CONCURRENCY)
4145 .with_metadata_size_hint(METADATA_SIZE_HINT)
4146 .with_bloom_filter_enabled(bound.is_some())
4147 .build();
4148 let tasks =
4149 Box::pin(futures::stream::iter(tasks.into_iter().map(Ok))) as FileScanTaskStream;
4150 let result = reader.read(tasks).unwrap();
4151 (result.metrics().clone(), log, result.stream())
4152 }
4153
4154 async fn deferred_scan() -> (ScanMetrics, ArrowRecordBatchStream) {
4156 let tmp_dir = TempDir::new().unwrap();
4157 let on_disk = write_plain_parquet(
4158 tmp_dir.path().to_str().unwrap(),
4159 "data.parquet",
4160 vec![],
4161 vec![],
4162 );
4163 let contents = Bytes::from(std::fs::read(on_disk).unwrap());
4164 let (metrics, _, scan) = deferred_scan_of(contents, CONCURRENCY, None).await;
4165 (metrics, scan)
4166 }
4167
4168 async fn spin(done: impl Fn() -> bool) {
4170 for _ in 0..YIELDS {
4171 if done() {
4172 return;
4173 }
4174 tokio::task::yield_now().await;
4175 }
4176 }
4177
4178 #[tokio::test]
4180 async fn test_open_proceeds_without_the_scan_being_polled() {
4181 let (metrics, scan) = deferred_scan().await;
4182 let mut scan = Box::pin(scan);
4183 let every_open = OPEN_BYTES * CONCURRENCY as u64;
4184
4185 let _ = futures::poll!(scan.next());
4188 spin(|| metrics.bytes_read() >= every_open).await;
4189 assert_eq!(
4190 metrics.bytes_read(),
4191 every_open,
4192 "the opens stopped where the scan stopped polling"
4193 );
4194 }
4195
4196 #[tokio::test]
4198 async fn test_dropped_scan_stops_opening_files() {
4199 let (metrics, scan) = deferred_scan().await;
4200 let mut scan = Box::pin(scan);
4201 let every_open = OPEN_BYTES * CONCURRENCY as u64;
4202
4203 let _ = futures::poll!(scan.next());
4205 spin(|| metrics.bytes_read() > 0).await;
4206 let in_flight = metrics.bytes_read();
4207 assert!(in_flight > 0, "no open reached the storage");
4208 assert!(in_flight < every_open, "the opens finished before the drop");
4209
4210 drop(scan);
4211 spin(|| false).await;
4212 assert_eq!(
4213 metrics.bytes_read(),
4214 in_flight,
4215 "a dropped scan read {} more bytes",
4216 metrics.bytes_read() - in_flight
4217 );
4218 }
4219
4220 #[tokio::test]
4223 async fn test_bloom_filter_fetch_proceeds_without_the_scan_being_polled() {
4224 let tmp_dir = TempDir::new().unwrap();
4225 let props = WriterProperties::builder()
4226 .set_compression(Compression::SNAPPY)
4227 .set_bloom_filter_enabled(true)
4228 .build();
4229 let on_disk = write_parquet_with(
4230 tmp_dir.path().to_str().unwrap(),
4231 "bloom.parquet",
4232 vec![],
4233 vec![],
4234 props,
4235 );
4236 let bloom_offset =
4237 ParquetRecordBatchReaderBuilder::try_new(File::open(&on_disk).unwrap())
4238 .unwrap()
4239 .metadata()
4240 .row_group(0)
4241 .column(0)
4242 .bloom_filter_offset()
4243 .expect("the writer put a bloom filter on `id`") as u64;
4244
4245 let contents = Bytes::from(std::fs::read(&on_disk).unwrap());
4246 let predicate = Reference::new("id").equal_to(Datum::int(2));
4247 let (_, log, scan) = deferred_scan_of(contents, 1, Some(predicate)).await;
4248 let mut scan = Box::pin(scan);
4249
4250 let at_the_bloom_filter = || log.last_start() == Some(bloom_offset);
4253 for _ in 0..YIELDS {
4254 if at_the_bloom_filter() {
4255 break;
4256 }
4257 let _ = futures::poll!(scan.next());
4258 tokio::task::yield_now().await;
4259 }
4260 assert!(at_the_bloom_filter(), "the bloom filter was never read");
4261
4262 spin(|| log.in_flight() == 0).await;
4264 assert_eq!(
4265 log.in_flight(),
4266 0,
4267 "a read was left in flight where the scan stopped polling"
4268 );
4269 }
4270 }
4271}