1use std::collections::HashMap;
24use std::sync::Arc;
25use std::sync::atomic::AtomicU64;
26
27use arrow_schema::{DataType, Field};
28use futures::{StreamExt, TryStreamExt};
29use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
30use parquet::arrow::{
31 PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder, ProjectionMask, RowNumber,
32};
33use parquet::encryption::decrypt::FileDecryptionProperties;
34
35use super::row_lineage::synthesize_row_id_column;
36use super::{
37 ArrowFileReader, ArrowReader, ParquetReadOptions, add_fallback_field_ids_to_arrow_schema,
38 apply_name_mapping_to_arrow_schema, find_leaf_by_field_id,
39};
40use crate::arrow::build_partition_constant;
41use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader;
42use crate::arrow::int96::coerce_int96_timestamps;
43use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder;
44use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics, ScanResult};
45use crate::encryption::StandardKeyMetadata;
46use crate::error::Result;
47use crate::io::{FileIO, FileMetadata, FileRead};
48use crate::metadata_columns::{
49 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER, RESERVED_COL_NAME_POS,
50 RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE,
51 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, RESERVED_FIELD_ID_PARTITION,
52 RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID, RESERVED_FIELD_ID_SPEC_ID, is_metadata_field,
53};
54use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream};
55use crate::spec::{Datum, PartitionSpec, Struct};
56use crate::{Error, ErrorKind};
57
58impl ArrowReader {
59 pub fn read(self, tasks: FileScanTaskStream) -> Result<ScanResult> {
62 let concurrency_limit_data_files = self.concurrency_limit_data_files;
63 let scan_metrics = ScanMetrics::new();
64
65 let task_reader = FileScanTaskReader {
66 batch_size: self.batch_size,
67 file_io: self.file_io,
68 delete_file_loader: self
69 .delete_file_loader
70 .with_scan_metrics(scan_metrics.clone()),
71 row_group_filtering_enabled: self.row_group_filtering_enabled,
72 row_selection_enabled: self.row_selection_enabled,
73 parquet_read_options: self.parquet_read_options,
74 scan_metrics: scan_metrics.clone(),
75 };
76
77 let stream: ArrowRecordBatchStream = if concurrency_limit_data_files == 1 {
79 Box::pin(
80 tasks
81 .and_then(move |task| task_reader.clone().process(task))
82 .map_err(|err| {
83 Error::new(ErrorKind::Unexpected, "file scan task generate failed")
84 .with_source(err)
85 })
86 .try_flatten(),
87 )
88 } else {
89 Box::pin(
90 tasks
91 .map_ok(move |task| task_reader.clone().process(task))
92 .map_err(|err| {
93 Error::new(ErrorKind::Unexpected, "file scan task generate failed")
94 .with_source(err)
95 })
96 .try_buffer_unordered(concurrency_limit_data_files)
97 .try_flatten_unordered(concurrency_limit_data_files),
98 )
99 };
100
101 Ok(ScanResult::new(stream, scan_metrics))
102 }
103}
104
105const PRUNABLE_METADATA_FIELDS: &[i32] = &[
109 RESERVED_FIELD_ID_FILE,
110 RESERVED_FIELD_ID_SPEC_ID,
111 RESERVED_FIELD_ID_PARTITION,
112 RESERVED_FIELD_ID_POS,
113 RESERVED_FIELD_ID_ROW_ID,
114 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
115];
116
117#[derive(Clone)]
120struct FileScanTaskReader {
121 batch_size: Option<usize>,
122 file_io: FileIO,
123 delete_file_loader: CachingDeleteFileLoader,
124 row_group_filtering_enabled: bool,
125 row_selection_enabled: bool,
126 parquet_read_options: ParquetReadOptions,
127 scan_metrics: ScanMetrics,
128}
129
130impl FileScanTaskReader {
131 async fn process(self, task: FileScanTask) -> Result<ArrowRecordBatchStream> {
132 let should_load_page_index = (self.row_selection_enabled && task.predicate().is_some())
133 || !task.deletes().is_empty();
134 let mut parquet_read_options = self.parquet_read_options;
135 parquet_read_options.preload_page_index = should_load_page_index;
136
137 let delete_filter_rx = self
138 .delete_file_loader
139 .load_deletes(task.deletes(), task.schema_ref());
140
141 let (parquet_file_reader, arrow_metadata) = ArrowReader::open_parquet_file(
143 task.data_file_path(),
144 &self.file_io,
145 task.file_size_in_bytes(),
146 parquet_read_options,
147 self.scan_metrics.bytes_read_counter(),
148 task.key_metadata(),
149 )
150 .await?;
151
152 let missing_field_ids = arrow_metadata
156 .schema()
157 .fields()
158 .iter()
159 .next()
160 .is_some_and(|f| f.metadata().get(PARQUET_FIELD_ID_META_KEY).is_none());
161
162 let use_position_fallback = missing_field_ids && task.name_mapping().is_none();
167
168 let arrow_metadata = if missing_field_ids {
184 let arrow_schema = if let Some(name_mapping) = task.name_mapping() {
186 apply_name_mapping_to_arrow_schema(
191 Arc::clone(arrow_metadata.schema()),
192 name_mapping,
193 )?
194 } else {
195 add_fallback_field_ids_to_arrow_schema(arrow_metadata.schema())
198 };
199
200 let options = ArrowReaderOptions::new().with_schema(arrow_schema);
201 ArrowReaderMetadata::try_new(Arc::clone(arrow_metadata.metadata()), options).map_err(
202 |e| {
203 Error::new(
204 ErrorKind::Unexpected,
205 "Failed to create ArrowReaderMetadata with field ID schema",
206 )
207 .with_source(e)
208 },
209 )?
210 } else {
211 arrow_metadata
213 };
214
215 let arrow_metadata = if let Some(coerced_schema) =
218 coerce_int96_timestamps(arrow_metadata.schema(), task.schema())
219 {
220 let options = ArrowReaderOptions::new().with_schema(Arc::clone(&coerced_schema));
221 ArrowReaderMetadata::try_new(Arc::clone(arrow_metadata.metadata()), options).map_err(
222 |e| {
223 Error::new(
224 ErrorKind::Unexpected,
225 format!(
226 "Failed to create ArrowReaderMetadata with INT96-coerced schema: {coerced_schema}"
227 ),
228 )
229 .with_source(e)
230 },
231 )?
232 } else {
233 arrow_metadata
234 };
235
236 let project_pos = task.project_field_ids().contains(&RESERVED_FIELD_ID_POS);
237 let project_row_id = task.project_field_ids().contains(&RESERVED_FIELD_ID_ROW_ID);
238
239 let need_row_number = project_pos || (project_row_id && task.first_row_id().is_some());
244
245 let field_ids = task.project_field_ids();
246 let metadata_only_projection = !field_ids.is_empty()
247 && field_ids
248 .iter()
249 .all(|id| PRUNABLE_METADATA_FIELDS.contains(id));
250
251 let install_row_number = need_row_number || metadata_only_projection;
252
253 let arrow_metadata = if install_row_number {
254 let row_number_field = Arc::new(
255 Field::new(RESERVED_COL_NAME_POS, DataType::Int64, false)
256 .with_metadata(HashMap::from([(
257 PARQUET_FIELD_ID_META_KEY.to_string(),
258 RESERVED_FIELD_ID_POS.to_string(),
259 )]))
260 .with_extension_type(RowNumber),
261 );
262
263 let options = ArrowReaderOptions::new()
264 .with_schema(Arc::clone(arrow_metadata.schema()))
265 .with_virtual_columns(vec![row_number_field])?;
266
267 ArrowReaderMetadata::try_new(Arc::clone(arrow_metadata.metadata()), options).map_err(
268 |e| {
269 Error::new(
270 ErrorKind::Unexpected,
271 "Failed to create ArrowReaderMetadata with the 'row_number' virtual_column",
272 )
273 .with_source(e)
274 },
275 )?
276 } else {
277 arrow_metadata
278 };
279
280 let mut record_batch_stream_builder =
282 ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_file_reader, arrow_metadata);
283
284 let project_last_updated_seq = task
288 .project_field_ids()
289 .contains(&RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER);
290
291 let phys_last_updated_seq_leaf = if project_last_updated_seq {
297 find_leaf_by_field_id(
298 record_batch_stream_builder.parquet_schema(),
299 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
300 )
301 } else {
302 None
303 };
304
305 let last_updated_seq_present_by_name_only = project_last_updated_seq
308 && phys_last_updated_seq_leaf.is_none()
309 && record_batch_stream_builder
310 .schema()
311 .fields()
312 .iter()
313 .any(|f| f.name() == RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER);
314
315 let coalesce_last_updated_seq_leaf = phys_last_updated_seq_leaf
320 .filter(|_| task.first_row_id().is_some() && task.data_sequence_number().is_some());
321
322 let phys_row_id_leaf = if project_row_id {
323 find_leaf_by_field_id(
324 record_batch_stream_builder.parquet_schema(),
325 RESERVED_FIELD_ID_ROW_ID,
326 )
327 } else {
328 None
329 };
330
331 let row_id_present_by_name_only = project_row_id
337 && !use_position_fallback
338 && phys_row_id_leaf.is_none()
339 && record_batch_stream_builder
340 .schema()
341 .fields()
342 .iter()
343 .any(|f| f.name() == RESERVED_COL_NAME_ROW_ID);
344
345 let coalesce_row_id_leaf = phys_row_id_leaf.filter(|_| task.first_row_id().is_some());
348
349 let project_field_ids_without_metadata: Vec<i32> = task
351 .project_field_ids()
352 .iter()
353 .filter(|&&id| !is_metadata_field(id))
354 .copied()
355 .collect();
356
357 let mut projection_mask = ArrowReader::get_arrow_projection_mask(
363 &project_field_ids_without_metadata,
364 task.schema(),
365 record_batch_stream_builder.parquet_schema(),
366 record_batch_stream_builder.schema(),
367 use_position_fallback, )?;
369
370 if project_field_ids_without_metadata.is_empty() && install_row_number {
381 projection_mask =
382 ProjectionMask::none(record_batch_stream_builder.parquet_schema().num_columns());
383 }
384
385 for leaf in [coalesce_last_updated_seq_leaf, coalesce_row_id_leaf]
390 .into_iter()
391 .flatten()
392 {
393 let phys_mask =
394 ProjectionMask::leaves(record_batch_stream_builder.parquet_schema(), vec![leaf]);
395 projection_mask.union(&phys_mask);
396 }
397
398 record_batch_stream_builder =
399 record_batch_stream_builder.with_projection(projection_mask.clone());
400
401 let mut record_batch_transformer_builder =
405 RecordBatchTransformerBuilder::new(task.schema_ref(), task.project_field_ids());
406
407 if task.project_field_ids().contains(&RESERVED_FIELD_ID_FILE) {
409 let file_datum = Datum::string(task.data_file_path().to_string());
410 record_batch_transformer_builder =
411 record_batch_transformer_builder.with_constant(RESERVED_FIELD_ID_FILE, file_datum);
412 }
413
414 if task
415 .project_field_ids()
416 .contains(&RESERVED_FIELD_ID_SPEC_ID)
417 {
418 let partition_spec = task
419 .partition_spec()
420 .ok_or_else(|| Error::new(ErrorKind::Unexpected, "Partition spec is missing"))?;
421
422 let spec_id_datum = Datum::int(partition_spec.spec_id());
423 record_batch_transformer_builder = record_batch_transformer_builder
424 .with_constant(RESERVED_FIELD_ID_SPEC_ID, spec_id_datum);
425 }
426
427 if project_last_updated_seq {
428 record_batch_transformer_builder =
433 match (task.first_row_id(), task.data_sequence_number()) {
434 (Some(_), Some(seq)) => {
435 let datum = Datum::long(seq);
436 if coalesce_last_updated_seq_leaf.is_some() {
437 record_batch_transformer_builder.with_coalesced_last_updated_seq_column(
440 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
441 datum,
442 )
443 } else if last_updated_seq_present_by_name_only {
444 return Err(Error::new(
452 ErrorKind::FeatureUnsupported,
453 "Reading a physically-stored _last_updated_sequence_number column \
454 without an embedded field id is not supported",
455 ));
456 } else {
457 record_batch_transformer_builder.with_constant(
459 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
460 datum,
461 )
462 }
463 }
464 (None, _) => record_batch_transformer_builder.with_null_metadata_column(
466 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
467 )?,
468 (Some(_), None) => {
472 return Err(Error::new(
473 ErrorKind::DataInvalid,
474 format!(
475 "Data file {} has a first_row_id but no data sequence number",
476 task.data_file_path()
477 ),
478 ));
479 }
480 };
481 }
482
483 if project_row_id {
484 if task.first_row_id().is_some() && row_id_present_by_name_only {
490 return Err(Error::new(
491 ErrorKind::FeatureUnsupported,
492 "Reading a physically-stored _row_id column without an embedded field id \
493 is not supported",
494 ));
495 }
496
497 record_batch_transformer_builder =
501 record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_ROW_ID);
502 }
503
504 if let (Some(partition_spec), Some(partition_data)) =
505 (task.partition_spec(), task.partition())
506 {
507 record_batch_transformer_builder = record_batch_transformer_builder
508 .with_partition(Arc::clone(partition_spec), partition_data.clone())?;
509 }
510
511 if project_pos {
512 record_batch_transformer_builder =
513 record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_POS);
514 }
515
516 if task
519 .project_field_ids()
520 .contains(&RESERVED_FIELD_ID_PARTITION)
521 && let Some(unified_type) = task.unified_partition_type()
522 {
523 let (spec, partition_data) = match (task.partition_spec(), task.partition()) {
524 (Some(spec), Some(data)) => (Arc::clone(spec), data.clone()),
525 _ if unified_type.fields().is_empty() => {
530 (Arc::new(PartitionSpec::unpartition_spec()), Struct::empty())
531 }
532 _ => {
533 return Err(Error::new(
534 ErrorKind::Unexpected,
535 "cannot build _partition column: unified partition type has fields \
536 but the scan task is missing its partition spec or data",
537 ));
538 }
539 };
540 let constant = build_partition_constant(unified_type, &spec, &partition_data)?;
541 record_batch_transformer_builder =
542 record_batch_transformer_builder.with_partition_constant(constant);
543 }
544
545 let mut record_batch_transformer = record_batch_transformer_builder.build();
546
547 if let Some(batch_size) = self.batch_size {
548 record_batch_stream_builder = record_batch_stream_builder.with_batch_size(batch_size);
549 }
550
551 let delete_filter = delete_filter_rx.await.unwrap()?;
552 let delete_predicate = delete_filter.build_equality_delete_predicate(&task).await?;
553
554 let final_predicate = match (task.predicate(), delete_predicate) {
559 (None, None) => None,
560 (Some(predicate), None) => Some(predicate.clone()),
561 (None, Some(ref predicate)) => Some(predicate.clone()),
562 (Some(filter_predicate), Some(delete_predicate)) => {
563 Some(filter_predicate.clone().and(delete_predicate))
564 }
565 };
566
567 let mut selected_row_group_indices = None;
583 let mut row_selection = None;
584
585 if task.start() != 0 || task.length() != 0 {
588 let byte_range_filtered_row_groups = ArrowReader::filter_row_groups_by_byte_range(
589 record_batch_stream_builder.metadata(),
590 task.start(),
591 task.length(),
592 )?;
593 selected_row_group_indices = Some(byte_range_filtered_row_groups);
594 }
595
596 if let Some(predicate) = final_predicate {
597 let (iceberg_field_ids, field_id_map) = ArrowReader::build_field_id_set_and_map(
598 record_batch_stream_builder.parquet_schema(),
599 record_batch_stream_builder.schema(),
600 &predicate,
601 use_position_fallback,
602 )?;
603
604 let row_filter = ArrowReader::get_row_filter(
605 &predicate,
606 record_batch_stream_builder.parquet_schema(),
607 &iceberg_field_ids,
608 &field_id_map,
609 )?;
610 record_batch_stream_builder = record_batch_stream_builder.with_row_filter(row_filter);
611
612 if self.row_group_filtering_enabled {
613 let predicate_filtered_row_groups = ArrowReader::get_selected_row_group_indices(
614 &predicate,
615 record_batch_stream_builder.metadata(),
616 &field_id_map,
617 task.schema(),
618 )?;
619
620 selected_row_group_indices = match selected_row_group_indices {
623 Some(byte_range_filtered) => {
624 let intersection: Vec<usize> = byte_range_filtered
626 .into_iter()
627 .filter(|idx| predicate_filtered_row_groups.contains(idx))
628 .collect();
629 Some(intersection)
630 }
631 None => Some(predicate_filtered_row_groups),
632 };
633 }
634
635 if self.row_selection_enabled {
636 row_selection = ArrowReader::get_row_selection_for_filter_predicate(
637 &predicate,
638 record_batch_stream_builder.metadata(),
639 &selected_row_group_indices,
640 &field_id_map,
641 task.schema(),
642 )?;
643 }
644 }
645
646 let positional_delete_indexes = delete_filter.get_delete_vector(&task);
647
648 if let Some(positional_delete_indexes) = positional_delete_indexes {
649 let delete_row_selection = {
650 let positional_delete_indexes = positional_delete_indexes.lock().unwrap();
651
652 ArrowReader::build_deletes_row_selection(
653 record_batch_stream_builder.metadata().row_groups(),
654 &selected_row_group_indices,
655 &positional_delete_indexes,
656 )
657 }?;
658
659 row_selection = match row_selection {
662 None => Some(delete_row_selection),
663 Some(filter_row_selection) => {
664 Some(filter_row_selection.intersection(&delete_row_selection))
665 }
666 };
667 }
668
669 if let Some(row_selection) = row_selection {
670 record_batch_stream_builder =
671 record_batch_stream_builder.with_row_selection(row_selection);
672 }
673
674 if let Some(selected_row_group_indices) = selected_row_group_indices {
675 record_batch_stream_builder =
676 record_batch_stream_builder.with_row_groups(selected_row_group_indices);
677 }
678
679 let first_row_id = task.first_row_id();
684 let record_batch_stream = record_batch_stream_builder.build()?.map(move |batch| {
685 let mut batch = batch.map_err(|err| -> Error { err.into() })?;
686 if project_row_id {
687 batch = synthesize_row_id_column(batch, first_row_id)?;
688 }
689 record_batch_transformer.process_record_batch(batch)
691 });
692
693 Ok(Box::pin(record_batch_stream) as ArrowRecordBatchStream)
694 }
695}
696
697impl ArrowReader {
698 pub(crate) async fn open_parquet_file(
701 data_file_path: &str,
702 file_io: &FileIO,
703 file_size_in_bytes: u64,
704 parquet_read_options: ParquetReadOptions,
705 bytes_read: &Arc<AtomicU64>,
706 key_metadata: Option<&[u8]>,
707 ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> {
708 let parquet_file = file_io.new_input(data_file_path)?;
709 let counting_reader =
710 CountingFileRead::new(parquet_file.reader().await?, Arc::clone(bytes_read));
711 Self::build_parquet_reader(
712 Box::new(counting_reader),
713 file_size_in_bytes,
714 parquet_read_options,
715 key_metadata,
716 )
717 .await
718 }
719
720 async fn build_parquet_reader(
721 parquet_reader: Box<dyn FileRead>,
722 file_size_in_bytes: u64,
723 parquet_read_options: ParquetReadOptions,
724 key_metadata: Option<&[u8]>,
725 ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> {
726 let mut reader = ArrowFileReader::new(
727 FileMetadata {
728 size: file_size_in_bytes,
729 },
730 parquet_reader,
731 )
732 .with_parquet_read_options(parquet_read_options);
733
734 let arrow_reader_options = Self::build_arrow_reader_options(key_metadata)?;
735
736 let arrow_metadata = ArrowReaderMetadata::load_async(&mut reader, arrow_reader_options)
737 .await
738 .map_err(|e| {
739 Error::new(ErrorKind::Unexpected, "Failed to load Parquet metadata").with_source(e)
740 })?;
741
742 Ok((reader, arrow_metadata))
743 }
744
745 fn build_arrow_reader_options(key_metadata: Option<&[u8]>) -> Result<ArrowReaderOptions> {
748 match key_metadata {
749 Some(km) => {
750 let standard_key_metadata = StandardKeyMetadata::decode(km)?;
751 let mut builder = FileDecryptionProperties::builder(
752 standard_key_metadata.encryption_key().as_bytes().to_vec(),
753 );
754 if let Some(aad) = standard_key_metadata.aad_prefix() {
755 builder = builder.with_aad_prefix(aad.to_vec());
756 }
757 let decryption_properties = builder.build().map_err(|e| {
758 Error::new(
759 ErrorKind::Unexpected,
760 "Failed to build Parquet file decryption properties",
761 )
762 .with_source(e)
763 })?;
764 Ok(
765 ArrowReaderOptions::new()
766 .with_file_decryption_properties(decryption_properties),
767 )
768 }
769 None => Ok(ArrowReaderOptions::default()),
770 }
771 }
772}
773
774#[cfg(test)]
775mod tests {
776 use std::collections::HashMap;
777 use std::fs::File;
778 use std::sync::Arc;
779
780 use arrow_array::cast::AsArray;
781 use arrow_array::{Array, ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray};
782 use arrow_cast::cast;
783 use arrow_schema::{DataType, Field, Schema as ArrowSchema};
784 use futures::TryStreamExt;
785 use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
786 use parquet::basic::Compression;
787 use parquet::file::properties::WriterProperties;
788 use tempfile::TempDir;
789
790 use crate::Runtime;
791 use crate::arrow::ArrowReaderBuilder;
792 use crate::arrow::test_utils::write_encrypted_parquet;
793 use crate::io::FileIO;
794 use crate::metadata_columns::{
795 RESERVED_COL_NAME_POS, RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE,
796 RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID,
797 };
798 use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream};
799 use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, SchemaRef, Type};
800
801 const UNIX_EPOCH_JULIAN: i64 = 2_440_588;
804 const MICROS_PER_DAY: i64 = 86_400_000_000;
805 const INT96_TEST_NANOS_WITHIN_DAY: u64 = 43_200_000_000_000;
807 const INT96_TEST_JULIAN_DAY: u32 = 2_953_529;
808
809 fn make_int96_test_value() -> (parquet::data_type::Int96, i64) {
810 let mut val = parquet::data_type::Int96::new();
811 val.set_data(
812 (INT96_TEST_NANOS_WITHIN_DAY & 0xFFFFFFFF) as u32,
813 (INT96_TEST_NANOS_WITHIN_DAY >> 32) as u32,
814 INT96_TEST_JULIAN_DAY,
815 );
816 let expected_micros = (INT96_TEST_JULIAN_DAY as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY
817 + (INT96_TEST_NANOS_WITHIN_DAY / 1_000) as i64;
818 (val, expected_micros)
819 }
820
821 async fn read_int96_batches(
822 file_path: &str,
823 schema: SchemaRef,
824 project_field_ids: Vec<i32>,
825 ) -> Vec<RecordBatch> {
826 let file_io = FileIO::new_with_fs();
827 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
828
829 let file_size = std::fs::metadata(file_path).unwrap().len();
830 let task = FileScanTask::builder()
831 .with_file_size_in_bytes(file_size)
832 .with_start(0)
833 .with_length(file_size)
834 .with_data_file_path(file_path.to_string())
835 .with_data_file_format(DataFileFormat::Parquet)
836 .with_schema(schema)
837 .with_project_field_ids(project_field_ids)
838 .with_case_sensitive(false)
839 .build()
840 .unwrap();
841
842 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
843 reader
844 .read(tasks)
845 .unwrap()
846 .stream()
847 .try_collect()
848 .await
849 .unwrap()
850 }
851
852 fn write_int96_parquet_file(
854 table_location: &str,
855 filename: &str,
856 with_field_ids: bool,
857 ) -> (String, Vec<i64>) {
858 use parquet::basic::{Repetition, Type as PhysicalType};
859 use parquet::data_type::{Int32Type, Int96, Int96Type};
860 use parquet::file::writer::SerializedFileWriter;
861 use parquet::schema::types::Type as SchemaType;
862
863 let file_path = format!("{table_location}/{filename}");
864
865 let mut ts_builder = SchemaType::primitive_type_builder("ts", PhysicalType::INT96)
866 .with_repetition(Repetition::OPTIONAL);
867 let mut id_builder = SchemaType::primitive_type_builder("id", PhysicalType::INT32)
868 .with_repetition(Repetition::REQUIRED);
869
870 if with_field_ids {
871 ts_builder = ts_builder.with_id(Some(1));
872 id_builder = id_builder.with_id(Some(2));
873 }
874
875 let schema = SchemaType::group_type_builder("schema")
876 .with_fields(vec![
877 Arc::new(ts_builder.build().unwrap()),
878 Arc::new(id_builder.build().unwrap()),
879 ])
880 .build()
881 .unwrap();
882
883 const NOON_NANOS: u64 = INT96_TEST_NANOS_WITHIN_DAY;
885 const JULIAN_3333: u32 = INT96_TEST_JULIAN_DAY;
886 const JULIAN_2100: u32 = 2_488_070;
887
888 let test_data: Vec<(u32, u32, u32, i64)> = vec![
889 (
891 0,
892 0,
893 JULIAN_3333,
894 (JULIAN_3333 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY,
895 ),
896 (
898 (NOON_NANOS & 0xFFFFFFFF) as u32,
899 (NOON_NANOS >> 32) as u32,
900 JULIAN_3333,
901 (JULIAN_3333 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY
902 + (NOON_NANOS / 1_000) as i64,
903 ),
904 (
906 0,
907 0,
908 JULIAN_2100,
909 (JULIAN_2100 as i64 - UNIX_EPOCH_JULIAN) * MICROS_PER_DAY,
910 ),
911 ];
912
913 let int96_values: Vec<Int96> = test_data
914 .iter()
915 .map(|(lo, hi, day, _)| {
916 let mut v = Int96::new();
917 v.set_data(*lo, *hi, *day);
918 v
919 })
920 .collect();
921
922 let id_values: Vec<i32> = (0..test_data.len() as i32).collect();
923 let expected_micros: Vec<i64> = test_data.iter().map(|(_, _, _, m)| *m).collect();
924
925 let file = File::create(&file_path).unwrap();
926 let mut writer =
927 SerializedFileWriter::new(file, Arc::new(schema), Default::default()).unwrap();
928
929 let mut row_group = writer.next_row_group().unwrap();
930 {
931 let mut col = row_group.next_column().unwrap().unwrap();
933 col.typed::<Int96Type>()
934 .write_batch(&int96_values, Some(&vec![1; test_data.len()]), None)
935 .unwrap();
936 col.close().unwrap();
937 }
938 {
939 let mut col = row_group.next_column().unwrap().unwrap();
940 col.typed::<Int32Type>()
941 .write_batch(&id_values, None, None)
942 .unwrap();
943 col.close().unwrap();
944 }
945 row_group.close().unwrap();
946 writer.close().unwrap();
947
948 (file_path, expected_micros)
949 }
950
951 async fn assert_int96_read_matches(
952 file_path: &str,
953 schema: SchemaRef,
954 project_field_ids: Vec<i32>,
955 expected_micros: &[i64],
956 ) {
957 use arrow_array::TimestampMicrosecondArray;
958
959 let batches = read_int96_batches(file_path, schema, project_field_ids).await;
960
961 assert_eq!(batches.len(), 1);
962 let ts_array = batches[0]
963 .column(0)
964 .as_any()
965 .downcast_ref::<TimestampMicrosecondArray>()
966 .expect("Expected TimestampMicrosecondArray");
967
968 for (i, expected) in expected_micros.iter().enumerate() {
969 assert_eq!(
970 ts_array.value(i),
971 *expected,
972 "Row {i}: got {}, expected {expected}",
973 ts_array.value(i)
974 );
975 }
976 }
977
978 async fn assert_encrypted_parquet_roundtrip(encryption_key: &[u8]) {
982 let aad_prefix = b"aad_prefix";
983
984 let schema = Arc::new(
985 Schema::builder()
986 .with_schema_id(1)
987 .with_fields(vec![
988 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
989 ])
990 .build()
991 .unwrap(),
992 );
993
994 let arrow_schema = Arc::new(ArrowSchema::new(vec![
995 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
996 PARQUET_FIELD_ID_META_KEY.to_string(),
997 "1".to_string(),
998 )])),
999 ]));
1000
1001 let tmp_dir = TempDir::new().unwrap();
1002 let table_location = tmp_dir.path().to_str().unwrap().to_string();
1003 let file_io = FileIO::new_with_fs();
1004
1005 let id_data = Arc::new(Int32Array::from(vec![10, 20, 30])) as ArrayRef;
1006 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
1007
1008 let file_path = format!("{table_location}/encrypted.parquet");
1009 write_encrypted_parquet(&file_path, &batch, encryption_key, Some(aad_prefix));
1010
1011 let key_metadata = crate::encryption::StandardKeyMetadata::try_new(encryption_key)
1012 .unwrap()
1013 .with_aad_prefix(aad_prefix)
1014 .encode()
1015 .unwrap();
1016
1017 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1018
1019 let task = FileScanTask::builder()
1020 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1021 .with_start(0)
1022 .with_length(0)
1023 .with_data_file_path(file_path)
1024 .with_data_file_format(DataFileFormat::Parquet)
1025 .with_schema(schema)
1026 .with_project_field_ids(vec![1])
1027 .with_case_sensitive(false)
1028 .with_key_metadata(Some(key_metadata))
1029 .build()
1030 .unwrap();
1031
1032 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1033 let batches: Vec<RecordBatch> = reader
1034 .read(tasks)
1035 .unwrap()
1036 .stream()
1037 .try_collect()
1038 .await
1039 .unwrap();
1040
1041 assert_eq!(batches.len(), 1);
1042 let ids = batches[0]
1043 .column(0)
1044 .as_any()
1045 .downcast_ref::<Int32Array>()
1046 .unwrap();
1047 assert_eq!(ids.values(), &[10, 20, 30]);
1048 }
1049
1050 #[tokio::test]
1051 async fn test_read_encrypted_parquet_aes_128() {
1052 assert_encrypted_parquet_roundtrip(b"0123456789abcdef").await;
1053 }
1054
1055 #[tokio::test]
1056 async fn test_read_encrypted_parquet_aes_256() {
1057 assert_encrypted_parquet_roundtrip(b"0123456789abcdef0123456789abcdef").await;
1058 }
1059
1060 #[tokio::test]
1061 async fn test_read_encrypted_parquet_without_key_metadata_fails() {
1062 let encryption_key = b"0123456789abcdef";
1063
1064 let schema = Arc::new(
1065 Schema::builder()
1066 .with_schema_id(1)
1067 .with_fields(vec![
1068 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1069 ])
1070 .build()
1071 .unwrap(),
1072 );
1073
1074 let arrow_schema = Arc::new(ArrowSchema::new(vec![
1075 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1076 PARQUET_FIELD_ID_META_KEY.to_string(),
1077 "1".to_string(),
1078 )])),
1079 ]));
1080
1081 let tmp_dir = TempDir::new().unwrap();
1082 let table_location = tmp_dir.path().to_str().unwrap().to_string();
1083 let file_io = FileIO::new_with_fs();
1084
1085 let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
1086 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
1087
1088 let file_path = format!("{table_location}/encrypted_no_key.parquet");
1089 write_encrypted_parquet(&file_path, &batch, encryption_key, None);
1090
1091 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1092
1093 let task = FileScanTask::builder()
1094 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1095 .with_start(0)
1096 .with_length(0)
1097 .with_data_file_path(file_path)
1098 .with_data_file_format(DataFileFormat::Parquet)
1099 .with_schema(schema)
1100 .with_project_field_ids(vec![1])
1101 .with_case_sensitive(false)
1102 .build()
1103 .unwrap();
1104
1105 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1106 let result: Result<Vec<RecordBatch>, _> =
1107 reader.read(tasks).unwrap().stream().try_collect().await;
1108
1109 let err = result.unwrap_err();
1110 assert_eq!(err.kind(), crate::ErrorKind::Unexpected);
1111 let err_str = format!("{err}");
1112 assert!(
1113 err_str.contains("encrypted footer"),
1114 "Expected error about encrypted footer, got: {err_str}"
1115 );
1116 assert!(
1117 err_str.contains("decryption properties were not provided"),
1118 "Expected error about missing decryption properties, got: {err_str}"
1119 );
1120 }
1121
1122 fn write_plain_parquet(
1125 dir: &str,
1126 name: &str,
1127 extra_fields: Vec<Field>,
1128 extra_columns: Vec<ArrayRef>,
1129 ) -> String {
1130 let mut fields =
1131 vec![
1132 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1133 PARQUET_FIELD_ID_META_KEY.to_string(),
1134 "1".to_string(),
1135 )])),
1136 ];
1137 fields.extend(extra_fields);
1138 let arrow_schema = Arc::new(ArrowSchema::new(fields));
1139
1140 let mut columns: Vec<ArrayRef> = vec![Arc::new(Int32Array::from(vec![1, 2, 3]))];
1141 columns.extend(extra_columns);
1142 let batch = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap();
1143
1144 let file_path = format!("{dir}/{name}");
1145 let file = File::create(&file_path).unwrap();
1146 let props = WriterProperties::builder()
1147 .set_compression(Compression::SNAPPY)
1148 .build();
1149 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
1150 writer.write(&batch).unwrap();
1151 writer.close().unwrap();
1152 file_path
1153 }
1154
1155 fn last_updated_seq_task(
1156 file_path: String,
1157 first_row_id: Option<i64>,
1158 data_sequence_number: Option<i64>,
1159 ) -> FileScanTask {
1160 use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1161
1162 let schema = Arc::new(
1163 Schema::builder()
1164 .with_schema_id(1)
1165 .with_fields(vec![
1166 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1167 ])
1168 .build()
1169 .unwrap(),
1170 );
1171
1172 FileScanTask::builder()
1173 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1174 .with_start(0)
1175 .with_length(0)
1176 .with_data_file_path(file_path)
1177 .with_data_file_format(DataFileFormat::Parquet)
1178 .with_schema(schema)
1179 .with_project_field_ids(vec![1, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER])
1180 .with_first_row_id(first_row_id)
1181 .with_data_sequence_number(data_sequence_number)
1182 .with_case_sensitive(false)
1183 .build()
1184 .unwrap()
1185 }
1186
1187 fn assert_last_updated_seq_column(batches: &[RecordBatch], expected: &[Option<i64>]) {
1190 use arrow_array::cast::AsArray;
1191 use arrow_cast::cast;
1192 use arrow_schema::DataType;
1193
1194 use crate::metadata_columns::RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER;
1195
1196 let mut actual = Vec::new();
1197 for batch in batches {
1198 let col = batch
1199 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
1200 .expect("_last_updated_sequence_number column should be present");
1201 let logical = cast(col, &DataType::Int64).unwrap();
1202 let values = logical.as_primitive::<arrow_array::types::Int64Type>();
1203 for i in 0..values.len() {
1204 actual.push((!values.is_null(i)).then(|| values.value(i)));
1205 }
1206 }
1207 assert_eq!(actual, expected);
1208 }
1209
1210 #[tokio::test]
1211 async fn test_last_updated_sequence_number_null_when_no_first_row_id() {
1212 let tmp_dir = TempDir::new().unwrap();
1213 let dir = tmp_dir.path().to_str().unwrap();
1214 let file_path = write_plain_parquet(dir, "no_first_row_id.parquet", vec![], vec![]);
1215
1216 let task = last_updated_seq_task(file_path, None, Some(9));
1220
1221 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1222 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1223 let batches: Vec<RecordBatch> = reader
1224 .read(tasks)
1225 .unwrap()
1226 .stream()
1227 .try_collect()
1228 .await
1229 .unwrap();
1230
1231 assert_last_updated_seq_column(&batches, &[None, None, None]);
1232 }
1233
1234 #[tokio::test]
1235 async fn test_last_updated_sequence_number_error_when_no_data_seq() {
1236 let tmp_dir = TempDir::new().unwrap();
1237 let dir = tmp_dir.path().to_str().unwrap();
1238 let file_path = write_plain_parquet(dir, "no_data_seq.parquet", vec![], vec![]);
1239
1240 let task = last_updated_seq_task(file_path, Some(42), None);
1244
1245 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1246 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1247 let result: Result<Vec<RecordBatch>, _> =
1248 reader.read(tasks).unwrap().stream().try_collect().await;
1249
1250 let err = result.unwrap_err();
1251 assert_eq!(err.kind(), crate::ErrorKind::DataInvalid);
1252 assert!(
1253 format!("{err}").contains("no data sequence number"),
1254 "unexpected error: {err}"
1255 );
1256 }
1257
1258 #[tokio::test]
1259 async fn test_last_updated_sequence_number_derived_from_data_seq() {
1260 let tmp_dir = TempDir::new().unwrap();
1261 let dir = tmp_dir.path().to_str().unwrap();
1262 let file_path = write_plain_parquet(dir, "with_first_row_id.parquet", vec![], vec![]);
1263
1264 let task = last_updated_seq_task(file_path, Some(42), Some(7));
1267
1268 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1269 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1270 let batches: Vec<RecordBatch> = reader
1271 .read(tasks)
1272 .unwrap()
1273 .stream()
1274 .try_collect()
1275 .await
1276 .unwrap();
1277
1278 assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]);
1279 }
1280
1281 #[tokio::test]
1282 async fn test_last_updated_sequence_number_mixed_files_share_schema() {
1283 use arrow_select::concat::concat_batches;
1284
1285 let tmp_dir = TempDir::new().unwrap();
1286 let dir = tmp_dir.path().to_str().unwrap();
1287
1288 let constant = last_updated_seq_task(
1294 write_plain_parquet(dir, "constant.parquet", vec![], vec![]),
1295 Some(42),
1296 Some(7),
1297 );
1298 let nulled = last_updated_seq_task(
1299 write_plain_parquet(dir, "nulled.parquet", vec![], vec![]),
1300 None,
1301 Some(7),
1302 );
1303 let coalesced = last_updated_seq_task(
1304 write_plain_parquet(
1305 dir,
1306 "coalesced.parquet",
1307 vec![physical_last_updated_seq_field()],
1308 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
1309 ),
1310 Some(50),
1311 Some(7),
1312 );
1313
1314 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1315 let tasks = Box::pin(futures::stream::iter(vec![
1316 Ok(constant),
1317 Ok(nulled),
1318 Ok(coalesced),
1319 ])) as FileScanTaskStream;
1320 let batches: Vec<RecordBatch> = reader
1321 .read(tasks)
1322 .unwrap()
1323 .stream()
1324 .try_collect()
1325 .await
1326 .unwrap();
1327
1328 assert_eq!(batches.len(), 3);
1329 let schema = batches[0].schema();
1331 concat_batches(&schema, &batches)
1332 .expect("constant, null and coalesce files must share one column type");
1333 }
1334
1335 fn physical_last_updated_seq_field() -> Field {
1337 use crate::metadata_columns::{
1338 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1339 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1340 };
1341 Field::new(
1342 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1343 DataType::Int64,
1344 true,
1345 )
1346 .with_metadata(HashMap::from([(
1347 PARQUET_FIELD_ID_META_KEY.to_string(),
1348 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER.to_string(),
1349 )]))
1350 }
1351
1352 #[tokio::test]
1353 async fn test_last_updated_sequence_number_physical_column_coalesced() {
1354 let tmp_dir = TempDir::new().unwrap();
1355 let dir = tmp_dir.path().to_str().unwrap();
1356 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1360 let file_path = write_plain_parquet(
1361 dir,
1362 "with_seq.parquet",
1363 vec![physical_last_updated_seq_field()],
1364 vec![seq_col],
1365 );
1366
1367 let task = last_updated_seq_task(file_path, Some(100), Some(9));
1368
1369 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1370 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1371 let batches: Vec<RecordBatch> = reader
1372 .read(tasks)
1373 .unwrap()
1374 .stream()
1375 .try_collect()
1376 .await
1377 .unwrap();
1378
1379 assert_last_updated_seq_column(&batches, &[Some(5), Some(9), Some(8)]);
1381 }
1382
1383 #[tokio::test]
1384 async fn test_last_updated_sequence_number_coalesced_with_pos_column() {
1385 use crate::metadata_columns::{
1386 RESERVED_COL_NAME_POS, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1387 RESERVED_FIELD_ID_POS,
1388 };
1389
1390 let tmp_dir = TempDir::new().unwrap();
1391 let dir = tmp_dir.path().to_str().unwrap();
1392 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1393 let file_path = write_plain_parquet(
1394 dir,
1395 "with_seq_and_pos.parquet",
1396 vec![physical_last_updated_seq_field()],
1397 vec![seq_col],
1398 );
1399
1400 let schema = Arc::new(
1405 Schema::builder()
1406 .with_schema_id(1)
1407 .with_fields(vec![
1408 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1409 ])
1410 .build()
1411 .unwrap(),
1412 );
1413 let task = FileScanTask::builder()
1414 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1415 .with_start(0)
1416 .with_length(0)
1417 .with_data_file_path(file_path)
1418 .with_data_file_format(DataFileFormat::Parquet)
1419 .with_schema(schema)
1420 .with_project_field_ids(vec![
1421 1,
1422 RESERVED_FIELD_ID_POS,
1423 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1424 ])
1425 .with_first_row_id(Some(100))
1426 .with_data_sequence_number(Some(9))
1427 .with_case_sensitive(false)
1428 .build()
1429 .unwrap();
1430
1431 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1432 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1433 let batches: Vec<RecordBatch> = reader
1434 .read(tasks)
1435 .unwrap()
1436 .stream()
1437 .try_collect()
1438 .await
1439 .unwrap();
1440
1441 assert_last_updated_seq_column(&batches, &[Some(5), Some(9), Some(8)]);
1443 let pos_col = batches[0]
1445 .column_by_name(RESERVED_COL_NAME_POS)
1446 .expect("_pos column should be present")
1447 .as_primitive::<arrow_array::types::Int64Type>();
1448 assert_eq!(pos_col.values(), &[0, 1, 2]);
1449 }
1450
1451 #[tokio::test]
1452 async fn test_last_updated_sequence_number_physical_column_nulled_without_first_row_id() {
1453 let tmp_dir = TempDir::new().unwrap();
1454 let dir = tmp_dir.path().to_str().unwrap();
1455 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1456 let file_path = write_plain_parquet(
1457 dir,
1458 "with_seq_no_first_row_id.parquet",
1459 vec![physical_last_updated_seq_field()],
1460 vec![seq_col],
1461 );
1462
1463 let task = last_updated_seq_task(file_path, None, Some(9));
1466
1467 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1468 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1469 let batches: Vec<RecordBatch> = reader
1470 .read(tasks)
1471 .unwrap()
1472 .stream()
1473 .try_collect()
1474 .await
1475 .unwrap();
1476
1477 assert_last_updated_seq_column(&batches, &[None, None, None]);
1478 }
1479
1480 #[tokio::test]
1481 async fn test_last_updated_sequence_number_present_by_name_without_id_unsupported() {
1482 let tmp_dir = TempDir::new().unwrap();
1483 let dir = tmp_dir.path().to_str().unwrap();
1484 let seq_field = Field::new(
1488 crate::metadata_columns::RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
1489 DataType::Int64,
1490 true,
1491 );
1492 let seq_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1493 let file_path =
1494 write_plain_parquet(dir, "with_seq_by_name.parquet", vec![seq_field], vec![
1495 seq_col,
1496 ]);
1497
1498 let task = last_updated_seq_task(file_path, Some(100), Some(9));
1499
1500 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1501 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1502 let result: Result<Vec<RecordBatch>, _> =
1503 reader.read(tasks).unwrap().stream().try_collect().await;
1504
1505 let err = result.unwrap_err();
1506 assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported);
1507 assert!(
1508 format!("{err}").contains("without an embedded field id"),
1509 "unexpected error: {err}"
1510 );
1511 }
1512
1513 #[tokio::test]
1514 async fn test_last_updated_sequence_number_physical_column_first_row_id_without_data_seq() {
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_no_data_seq.parquet",
1521 vec![physical_last_updated_seq_field()],
1522 vec![seq_col],
1523 );
1524
1525 let task = last_updated_seq_task(file_path, Some(100), None);
1529
1530 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1531 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1532 let result: Result<Vec<RecordBatch>, _> =
1533 reader.read(tasks).unwrap().stream().try_collect().await;
1534
1535 let err = result.unwrap_err();
1536 assert_eq!(err.kind(), crate::ErrorKind::DataInvalid);
1537 assert!(
1538 format!("{err}").contains("no data sequence number"),
1539 "unexpected error: {err}"
1540 );
1541 }
1542
1543 fn row_id_task(file_path: String, first_row_id: Option<i64>) -> FileScanTask {
1545 row_id_task_with_options(file_path, first_row_id, 0, 0, vec![])
1546 }
1547
1548 fn row_id_task_with_options(
1549 file_path: String,
1550 first_row_id: Option<i64>,
1551 start: u64,
1552 length: u64,
1553 deletes: Vec<FileScanTaskDeleteFile>,
1554 ) -> FileScanTask {
1555 let schema = Arc::new(
1556 Schema::builder()
1557 .with_schema_id(1)
1558 .with_fields(vec![
1559 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1560 ])
1561 .build()
1562 .unwrap(),
1563 );
1564
1565 FileScanTask::builder()
1566 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1567 .with_start(start)
1568 .with_length(length)
1569 .with_data_file_path(file_path)
1570 .with_data_file_format(DataFileFormat::Parquet)
1571 .with_schema(schema)
1572 .with_project_field_ids(vec![1, RESERVED_FIELD_ID_ROW_ID])
1573 .with_first_row_id(first_row_id)
1574 .with_deletes(deletes)
1575 .with_case_sensitive(false)
1576 .build()
1577 .unwrap()
1578 }
1579
1580 fn assert_row_id_column(batches: &[RecordBatch], expected: &[Option<i64>]) {
1583 use arrow_array::cast::AsArray;
1584 use arrow_cast::cast;
1585 use arrow_schema::DataType;
1586
1587 let mut actual = Vec::new();
1588 for batch in batches {
1589 let col = batch
1590 .column_by_name(RESERVED_COL_NAME_ROW_ID)
1591 .expect("_row_id column should be present");
1592 let logical = cast(col, &DataType::Int64).unwrap();
1593 let values = logical.as_primitive::<arrow_array::types::Int64Type>();
1594 for i in 0..values.len() {
1595 actual.push((!values.is_null(i)).then(|| values.value(i)));
1596 }
1597 }
1598 assert_eq!(actual, expected);
1599 }
1600
1601 fn physical_row_id_field() -> Field {
1603 Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true).with_metadata(HashMap::from([
1604 (
1605 PARQUET_FIELD_ID_META_KEY.to_string(),
1606 RESERVED_FIELD_ID_ROW_ID.to_string(),
1607 ),
1608 ]))
1609 }
1610
1611 #[tokio::test]
1612 async fn test_row_id_synthesized_from_first_row_id_and_pos() {
1613 let tmp_dir = TempDir::new().unwrap();
1614 let dir = tmp_dir.path().to_str().unwrap();
1615 let file_path = write_plain_parquet(dir, "row_id_synth.parquet", vec![], vec![]);
1616
1617 let task = row_id_task(file_path, Some(100));
1619
1620 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1621 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1622 let batches: Vec<RecordBatch> = reader
1623 .read(tasks)
1624 .unwrap()
1625 .stream()
1626 .try_collect()
1627 .await
1628 .unwrap();
1629
1630 assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
1631 }
1632
1633 #[tokio::test]
1634 async fn test_row_id_physical_column_coalesced() {
1635 let tmp_dir = TempDir::new().unwrap();
1636 let dir = tmp_dir.path().to_str().unwrap();
1637 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1640 let file_path = write_plain_parquet(
1641 dir,
1642 "row_id_phys.parquet",
1643 vec![physical_row_id_field()],
1644 vec![id_col],
1645 );
1646
1647 let task = row_id_task(file_path, Some(100));
1648
1649 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1650 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1651 let batches: Vec<RecordBatch> = reader
1652 .read(tasks)
1653 .unwrap()
1654 .stream()
1655 .try_collect()
1656 .await
1657 .unwrap();
1658
1659 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
1661 }
1662
1663 #[tokio::test]
1664 async fn test_row_id_only_synthesis_reads_no_data_columns() {
1665 let tmp_dir = TempDir::new().unwrap();
1670 let dir = tmp_dir.path().to_str().unwrap();
1671
1672 let meta_only = metadata_projection_task_with_first_row_id(
1673 write_parquet_with_wide_column(dir, "row_id_only.parquet", vec![], vec![]),
1674 id_and_wide_schema(),
1675 vec![RESERVED_FIELD_ID_ROW_ID],
1676 Some(100),
1677 );
1678 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1679
1680 assert_eq!(batches[0].num_columns(), 1);
1681 assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
1682
1683 let with_data = metadata_projection_task_with_first_row_id(
1685 write_parquet_with_wide_column(dir, "row_id_only_ref.parquet", vec![], vec![]),
1686 id_and_wide_schema(),
1687 vec![2, RESERVED_FIELD_ID_ROW_ID],
1688 Some(100),
1689 );
1690 let (_, with_data_bytes) = scan_task(with_data).await;
1691
1692 assert!(
1693 meta_only_bytes < with_data_bytes,
1694 "_row_id-only synthesis should read fewer bytes than a scan of the wide column: \
1695 {meta_only_bytes} vs {with_data_bytes}"
1696 );
1697 }
1698
1699 #[tokio::test]
1700 async fn test_row_id_only_null_first_row_id_reads_no_data_columns() {
1701 let tmp_dir = TempDir::new().unwrap();
1706 let dir = tmp_dir.path().to_str().unwrap();
1707
1708 let meta_only = metadata_projection_task_with_first_row_id(
1709 write_parquet_with_wide_column(dir, "row_id_null.parquet", vec![], vec![]),
1710 id_and_wide_schema(),
1711 vec![RESERVED_FIELD_ID_ROW_ID],
1712 None,
1713 );
1714 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1715
1716 assert_eq!(batches[0].num_columns(), 1);
1717 assert_row_id_column(&batches, &[None, None, None]);
1718
1719 let with_data = metadata_projection_task_with_first_row_id(
1720 write_parquet_with_wide_column(dir, "row_id_null_ref.parquet", vec![], vec![]),
1721 id_and_wide_schema(),
1722 vec![2, RESERVED_FIELD_ID_ROW_ID],
1723 None,
1724 );
1725 let (_, with_data_bytes) = scan_task(with_data).await;
1726
1727 assert!(
1728 meta_only_bytes < with_data_bytes,
1729 "_row_id-only scan with a null first_row_id should read fewer bytes than a scan \
1730 of the wide column: {meta_only_bytes} vs {with_data_bytes}"
1731 );
1732 }
1733
1734 #[tokio::test]
1735 async fn test_last_updated_seq_only_reads_no_data_columns() {
1736 use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1737
1738 let tmp_dir = TempDir::new().unwrap();
1743 let dir = tmp_dir.path().to_str().unwrap();
1744
1745 let lusn_only_task = |file_path: String, project_field_ids: Vec<i32>| {
1746 FileScanTask::builder()
1747 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1748 .with_start(0)
1749 .with_length(0)
1750 .with_data_file_path(file_path)
1751 .with_data_file_format(DataFileFormat::Parquet)
1752 .with_schema(id_and_wide_schema())
1753 .with_project_field_ids(project_field_ids)
1754 .with_first_row_id(Some(42))
1755 .with_data_sequence_number(Some(7))
1756 .with_case_sensitive(false)
1757 .build()
1758 .unwrap()
1759 };
1760
1761 let meta_only = lusn_only_task(
1762 write_parquet_with_wide_column(dir, "lusn_only.parquet", vec![], vec![]),
1763 vec![RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER],
1764 );
1765 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1766
1767 assert_eq!(batches[0].num_columns(), 1);
1768 assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]);
1769
1770 let with_data = lusn_only_task(
1771 write_parquet_with_wide_column(dir, "lusn_only_ref.parquet", vec![], vec![]),
1772 vec![2, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER],
1773 );
1774 let (_, with_data_bytes) = scan_task(with_data).await;
1775
1776 assert!(
1777 meta_only_bytes < with_data_bytes,
1778 "_last_updated_sequence_number-only scan should read fewer bytes than a scan of \
1779 the wide column: {meta_only_bytes} vs {with_data_bytes}"
1780 );
1781 }
1782
1783 #[tokio::test]
1784 async fn test_partition_only_reads_no_data_columns() {
1785 use arrow_array::StructArray;
1786
1787 use crate::metadata_columns::{RESERVED_COL_NAME_PARTITION, RESERVED_FIELD_ID_PARTITION};
1788 use crate::spec::{Literal, PartitionSpec, Struct, Transform};
1789
1790 let tmp_dir = TempDir::new().unwrap();
1794 let dir = tmp_dir.path().to_str().unwrap();
1795 let schema = id_and_wide_schema();
1796 let spec = Arc::new(
1797 PartitionSpec::builder(schema.clone())
1798 .with_spec_id(7)
1799 .add_partition_field("id", "id", Transform::Identity)
1800 .unwrap()
1801 .build()
1802 .unwrap(),
1803 );
1804 let unified_type = Arc::new(spec.partition_type(&schema).unwrap());
1805 let partition_data = Struct::from_iter(vec![Some(Literal::int(42))]);
1806
1807 let partition_task = |file_path: String, project_field_ids: Vec<i32>| {
1808 FileScanTask::builder()
1809 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1810 .with_start(0)
1811 .with_length(0)
1812 .with_data_file_path(file_path)
1813 .with_data_file_format(DataFileFormat::Parquet)
1814 .with_schema(schema.clone())
1815 .with_project_field_ids(project_field_ids)
1816 .with_partition_spec(Some(spec.clone()))
1817 .with_partition(Some(partition_data.clone()))
1818 .with_unified_partition_type(Some(unified_type.clone()))
1819 .with_case_sensitive(false)
1820 .build()
1821 .unwrap()
1822 };
1823
1824 let meta_only = partition_task(
1825 write_parquet_with_wide_column(dir, "partition_only.parquet", vec![], vec![]),
1826 vec![RESERVED_FIELD_ID_PARTITION],
1827 );
1828 let (batches, meta_only_bytes) = scan_task(meta_only).await;
1829
1830 assert_eq!(batches[0].num_columns(), 1);
1831 let partition_col = batches[0]
1832 .column_by_name(RESERVED_COL_NAME_PARTITION)
1833 .expect("_partition column should be present")
1834 .as_any()
1835 .downcast_ref::<StructArray>()
1836 .unwrap();
1837 assert_eq!(partition_col.len(), 3);
1838 let inner = partition_col
1839 .column(0)
1840 .as_any()
1841 .downcast_ref::<Int32Array>()
1842 .unwrap();
1843 assert_eq!(inner.values(), &[42, 42, 42]);
1844
1845 let with_data = partition_task(
1846 write_parquet_with_wide_column(dir, "partition_only_ref.parquet", vec![], vec![]),
1847 vec![2, RESERVED_FIELD_ID_PARTITION],
1848 );
1849 let (_, with_data_bytes) = scan_task(with_data).await;
1850
1851 assert!(
1852 meta_only_bytes < with_data_bytes,
1853 "_partition-only scan should read fewer bytes than a scan of the wide column: \
1854 {meta_only_bytes} vs {with_data_bytes}"
1855 );
1856 }
1857
1858 #[tokio::test]
1859 async fn test_row_id_resolves_alongside_id_less_leaf() {
1860 let tmp_dir = TempDir::new().unwrap();
1865 let dir = tmp_dir.path().to_str().unwrap();
1866 let idless_field = Field::new("variant_internal", DataType::Utf8, true);
1867 let idless_col = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
1868 let row_id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1869 let file_path = write_plain_parquet(
1870 dir,
1871 "row_id_with_idless_leaf.parquet",
1872 vec![idless_field, physical_row_id_field()],
1873 vec![idless_col, row_id_col],
1874 );
1875
1876 let task = row_id_task(file_path, Some(100));
1877
1878 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1879 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1880 let batches: Vec<RecordBatch> = reader
1881 .read(tasks)
1882 .unwrap()
1883 .stream()
1884 .try_collect()
1885 .await
1886 .unwrap();
1887
1888 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
1890 }
1891
1892 #[tokio::test]
1893 async fn test_row_id_and_last_updated_seq_co_projected() {
1894 use crate::metadata_columns::RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER;
1895
1896 let tmp_dir = TempDir::new().unwrap();
1900 let dir = tmp_dir.path().to_str().unwrap();
1901 let row_id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1902 let seq_col = Arc::new(Int64Array::from(vec![Some(50), None, Some(70)])) as ArrayRef;
1903 let file_path = write_plain_parquet(
1904 dir,
1905 "row_id_and_seq.parquet",
1906 vec![physical_row_id_field(), physical_last_updated_seq_field()],
1907 vec![row_id_col, seq_col],
1908 );
1909
1910 let schema = Arc::new(
1911 Schema::builder()
1912 .with_schema_id(1)
1913 .with_fields(vec![
1914 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1915 ])
1916 .build()
1917 .unwrap(),
1918 );
1919 let task = FileScanTask::builder()
1920 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
1921 .with_start(0)
1922 .with_length(0)
1923 .with_data_file_path(file_path)
1924 .with_data_file_format(DataFileFormat::Parquet)
1925 .with_schema(schema)
1926 .with_project_field_ids(vec![
1927 1,
1928 RESERVED_FIELD_ID_ROW_ID,
1929 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
1930 ])
1931 .with_first_row_id(Some(100))
1932 .with_data_sequence_number(Some(9))
1933 .with_case_sensitive(false)
1934 .build()
1935 .unwrap();
1936
1937 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1938 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1939 let batches: Vec<RecordBatch> = reader
1940 .read(tasks)
1941 .unwrap()
1942 .stream()
1943 .try_collect()
1944 .await
1945 .unwrap();
1946
1947 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
1949 assert_last_updated_seq_column(&batches, &[Some(50), Some(9), Some(70)]);
1951 }
1952
1953 #[tokio::test]
1954 async fn test_row_id_null_when_no_first_row_id() {
1955 let tmp_dir = TempDir::new().unwrap();
1956 let dir = tmp_dir.path().to_str().unwrap();
1957 let id_col = Arc::new(Int64Array::from(vec![Some(5), Some(6), Some(7)])) as ArrayRef;
1959 let file_path = write_plain_parquet(
1960 dir,
1961 "row_id_no_first.parquet",
1962 vec![physical_row_id_field()],
1963 vec![id_col],
1964 );
1965
1966 let task = row_id_task(file_path, None);
1968
1969 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
1970 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1971 let batches: Vec<RecordBatch> = reader
1972 .read(tasks)
1973 .unwrap()
1974 .stream()
1975 .try_collect()
1976 .await
1977 .unwrap();
1978
1979 assert_row_id_column(&batches, &[None, None, None]);
1980 }
1981
1982 #[tokio::test]
1983 async fn test_row_id_with_pos_column() {
1984 use crate::metadata_columns::RESERVED_COL_NAME_POS;
1985
1986 let tmp_dir = TempDir::new().unwrap();
1987 let dir = tmp_dir.path().to_str().unwrap();
1988 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
1989 let file_path = write_plain_parquet(
1990 dir,
1991 "row_id_and_pos.parquet",
1992 vec![physical_row_id_field()],
1993 vec![id_col],
1994 );
1995
1996 let schema = Arc::new(
2000 Schema::builder()
2001 .with_schema_id(1)
2002 .with_fields(vec![
2003 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2004 ])
2005 .build()
2006 .unwrap(),
2007 );
2008 let task = FileScanTask::builder()
2009 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2010 .with_start(0)
2011 .with_length(0)
2012 .with_data_file_path(file_path)
2013 .with_data_file_format(DataFileFormat::Parquet)
2014 .with_schema(schema)
2015 .with_project_field_ids(vec![1, RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID])
2016 .with_first_row_id(Some(100))
2017 .with_case_sensitive(false)
2018 .build()
2019 .unwrap();
2020
2021 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2022 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2023 let batches: Vec<RecordBatch> = reader
2024 .read(tasks)
2025 .unwrap()
2026 .stream()
2027 .try_collect()
2028 .await
2029 .unwrap();
2030
2031 assert_row_id_column(&batches, &[Some(5), Some(101), Some(8)]);
2033 let pos_col = batches[0]
2035 .column_by_name(RESERVED_COL_NAME_POS)
2036 .expect("_pos column should be present")
2037 .as_primitive::<arrow_array::types::Int64Type>();
2038 assert_eq!(pos_col.values(), &[0, 1, 2]);
2039 }
2040
2041 #[tokio::test]
2042 async fn test_row_id_mixed_files_share_schema() {
2043 use arrow_select::concat::concat_batches;
2044
2045 let tmp_dir = TempDir::new().unwrap();
2046 let dir = tmp_dir.path().to_str().unwrap();
2047
2048 let synth = row_id_task(
2054 write_plain_parquet(dir, "row_id_synth2.parquet", vec![], vec![]),
2055 Some(42),
2056 );
2057 let nulled = row_id_task(
2058 write_plain_parquet(dir, "row_id_null2.parquet", vec![], vec![]),
2059 None,
2060 );
2061 let coalesced = row_id_task(
2062 write_plain_parquet(
2063 dir,
2064 "row_id_coalesced2.parquet",
2065 vec![physical_row_id_field()],
2066 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
2067 ),
2068 Some(50),
2069 );
2070
2071 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2072 let tasks = Box::pin(futures::stream::iter(vec![
2073 Ok(synth),
2074 Ok(nulled),
2075 Ok(coalesced),
2076 ])) as FileScanTaskStream;
2077 let batches: Vec<RecordBatch> = reader
2078 .read(tasks)
2079 .unwrap()
2080 .stream()
2081 .try_collect()
2082 .await
2083 .unwrap();
2084
2085 assert_eq!(batches.len(), 3);
2086 let schema = batches[0].schema();
2087 concat_batches(&schema, &batches)
2088 .expect("synthesis, null and coalesce files must share one column type");
2089 }
2090
2091 #[tokio::test]
2092 async fn test_row_id_present_by_name_without_id_unsupported() {
2093 let tmp_dir = TempDir::new().unwrap();
2094 let dir = tmp_dir.path().to_str().unwrap();
2095 let id_field = Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true);
2098 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2099 let file_path =
2100 write_plain_parquet(dir, "row_id_by_name.parquet", vec![id_field], vec![id_col]);
2101
2102 let task = row_id_task(file_path, Some(100));
2103
2104 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2105 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2106 let result: Result<Vec<RecordBatch>, _> =
2107 reader.read(tasks).unwrap().stream().try_collect().await;
2108
2109 let err = result.unwrap_err();
2110 assert_eq!(err.kind(), crate::ErrorKind::FeatureUnsupported);
2111 assert!(
2112 format!("{err}").contains("without an embedded field id"),
2113 "unexpected error: {err}"
2114 );
2115 }
2116
2117 #[tokio::test]
2118 async fn test_row_id_present_by_name_without_id_nulls_when_no_lineage() {
2119 let tmp_dir = TempDir::new().unwrap();
2120 let dir = tmp_dir.path().to_str().unwrap();
2121 let id_field = Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true);
2126 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2127 let file_path = write_plain_parquet(
2128 dir,
2129 "row_id_by_name_no_lineage.parquet",
2130 vec![id_field],
2131 vec![id_col],
2132 );
2133
2134 let task = row_id_task(file_path, None);
2135
2136 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2137 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2138 let batches: Vec<RecordBatch> = reader
2139 .read(tasks)
2140 .unwrap()
2141 .stream()
2142 .try_collect()
2143 .await
2144 .unwrap();
2145
2146 assert_row_id_column(&batches, &[None, None, None]);
2147 }
2148
2149 fn row_id_task_with_predicate(
2152 file_path: String,
2153 first_row_id: Option<i64>,
2154 extra_project_field_ids: Vec<i32>,
2155 predicate: crate::expr::Predicate,
2156 ) -> FileScanTask {
2157 use crate::expr::Bind;
2158
2159 let schema = Arc::new(
2160 Schema::builder()
2161 .with_schema_id(1)
2162 .with_fields(vec![
2163 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2164 ])
2165 .build()
2166 .unwrap(),
2167 );
2168 let bound = predicate.bind(Arc::clone(&schema), false).unwrap();
2169
2170 let mut project_field_ids = vec![1];
2171 project_field_ids.extend(extra_project_field_ids);
2172 project_field_ids.push(RESERVED_FIELD_ID_ROW_ID);
2173
2174 FileScanTask::builder()
2175 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2176 .with_start(0)
2177 .with_length(0)
2178 .with_data_file_path(file_path)
2179 .with_data_file_format(DataFileFormat::Parquet)
2180 .with_schema(schema)
2181 .with_project_field_ids(project_field_ids)
2182 .with_predicate(Some(bound))
2183 .with_first_row_id(first_row_id)
2184 .with_case_sensitive(false)
2185 .build()
2186 .unwrap()
2187 }
2188
2189 #[tokio::test]
2190 async fn test_row_id_stable_under_row_selection() {
2191 use crate::expr::Reference;
2192 use crate::spec::Datum;
2193
2194 let tmp_dir = TempDir::new().unwrap();
2195 let dir = tmp_dir.path().to_str().unwrap();
2196 let file_path = write_plain_parquet(dir, "row_id_selection.parquet", vec![], vec![]);
2198
2199 let task = row_id_task_with_predicate(
2200 file_path,
2201 Some(100),
2202 vec![],
2203 Reference::new("id").not_equal_to(Datum::int(2)),
2204 );
2205
2206 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
2208 .with_row_selection_enabled(true)
2209 .build();
2210 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2211 let batches: Vec<RecordBatch> = reader
2212 .read(tasks)
2213 .unwrap()
2214 .stream()
2215 .try_collect()
2216 .await
2217 .unwrap();
2218
2219 assert_row_id_column(&batches, &[Some(100), Some(102)]);
2222 }
2223
2224 #[tokio::test]
2225 async fn test_row_id_coalesce_stable_under_row_selection() {
2226 use crate::expr::Reference;
2227 use crate::spec::Datum;
2228
2229 let tmp_dir = TempDir::new().unwrap();
2230 let dir = tmp_dir.path().to_str().unwrap();
2231 let id_col = Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef;
2234 let file_path = write_plain_parquet(
2235 dir,
2236 "row_id_coalesce_selection.parquet",
2237 vec![physical_row_id_field()],
2238 vec![id_col],
2239 );
2240
2241 let task = row_id_task_with_predicate(
2242 file_path,
2243 Some(100),
2244 vec![],
2245 Reference::new("id").not_equal_to(Datum::int(2)),
2246 );
2247
2248 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
2249 .with_row_selection_enabled(true)
2250 .build();
2251 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2252 let batches: Vec<RecordBatch> = reader
2253 .read(tasks)
2254 .unwrap()
2255 .stream()
2256 .try_collect()
2257 .await
2258 .unwrap();
2259
2260 assert_row_id_column(&batches, &[Some(5), Some(8)]);
2264 }
2265
2266 #[tokio::test]
2267 async fn test_row_id_global_across_row_groups() {
2268 let tmp_dir = TempDir::new().unwrap();
2269 let dir = tmp_dir.path().to_str().unwrap();
2270
2271 let file_path = format!("{dir}/row_id_multi_rg.parquet");
2275 let field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2276 PARQUET_FIELD_ID_META_KEY.to_string(),
2277 "1".to_string(),
2278 )]));
2279 let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
2280 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int32Array::from(
2281 vec![1, 2, 3, 4, 5],
2282 ))])
2283 .unwrap();
2284 let props = WriterProperties::builder()
2285 .set_compression(Compression::SNAPPY)
2286 .set_max_row_group_row_count(Some(2))
2287 .build();
2288 let file = File::create(&file_path).unwrap();
2289 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2290 writer.write(&batch).unwrap();
2291 writer.close().unwrap();
2292
2293 let task = row_id_task(file_path, Some(0));
2294 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2295 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2296 let batches: Vec<RecordBatch> = reader
2297 .read(tasks)
2298 .unwrap()
2299 .stream()
2300 .try_collect()
2301 .await
2302 .unwrap();
2303
2304 assert_row_id_column(&batches, &[Some(0), Some(1), Some(2), Some(3), Some(4)]);
2305 }
2306
2307 #[tokio::test]
2308 async fn test_row_id_global_when_first_row_group_pruned() {
2309 use parquet::file::reader::{FileReader, SerializedFileReader};
2310
2311 let tmp_dir = TempDir::new().unwrap();
2312 let dir = tmp_dir.path().to_str().unwrap();
2313
2314 let file_path = format!("{dir}/row_id_prune_rg0.parquet");
2319 let field = Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2320 PARQUET_FIELD_ID_META_KEY.to_string(),
2321 "1".to_string(),
2322 )]));
2323 let arrow_schema = Arc::new(ArrowSchema::new(vec![field]));
2324 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(Int32Array::from(
2325 vec![1, 2, 3, 4, 5, 6],
2326 ))])
2327 .unwrap();
2328 let props = WriterProperties::builder()
2329 .set_compression(Compression::SNAPPY)
2330 .set_max_row_group_row_count(Some(2))
2331 .build();
2332 let file = File::create(&file_path).unwrap();
2333 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2334 writer.write(&batch).unwrap();
2335 writer.close().unwrap();
2336
2337 let metadata = SerializedFileReader::new(File::open(&file_path).unwrap())
2340 .unwrap()
2341 .metadata()
2342 .clone();
2343 assert_eq!(metadata.num_row_groups(), 3);
2344 let start = 4 + metadata.row_group(0).compressed_size() as u64;
2345 let file_size = std::fs::metadata(&file_path).unwrap().len();
2346
2347 let task = row_id_task_with_options(file_path, Some(100), start, file_size - start, vec![]);
2348
2349 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2350 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2351 let batches: Vec<RecordBatch> = reader
2352 .read(tasks)
2353 .unwrap()
2354 .stream()
2355 .try_collect()
2356 .await
2357 .unwrap();
2358
2359 assert_row_id_column(&batches, &[Some(102), Some(103), Some(104), Some(105)]);
2362 }
2363
2364 fn write_positional_delete(
2367 dir: &str,
2368 name: &str,
2369 data_file_path: &str,
2370 positions: &[i64],
2371 ) -> String {
2372 use arrow_array::StringArray;
2373
2374 let file_path_field =
2375 Field::new("file_path", DataType::Utf8, false).with_metadata(HashMap::from([(
2376 PARQUET_FIELD_ID_META_KEY.to_string(),
2377 "2147483546".to_string(),
2378 )]));
2379 let pos_field = Field::new("pos", DataType::Int64, false).with_metadata(HashMap::from([(
2380 PARQUET_FIELD_ID_META_KEY.to_string(),
2381 "2147483545".to_string(),
2382 )]));
2383 let arrow_schema = Arc::new(ArrowSchema::new(vec![file_path_field, pos_field]));
2384 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![
2385 Arc::new(StringArray::from(vec![data_file_path; positions.len()])),
2386 Arc::new(Int64Array::from(positions.to_vec())),
2387 ])
2388 .unwrap();
2389
2390 let path = format!("{dir}/{name}");
2391 let file = File::create(&path).unwrap();
2392 let props = WriterProperties::builder()
2393 .set_compression(Compression::SNAPPY)
2394 .build();
2395 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2396 writer.write(&batch).unwrap();
2397 writer.close().unwrap();
2398 path
2399 }
2400
2401 #[tokio::test]
2402 async fn test_row_id_survives_positional_delete() {
2403 use crate::spec::DataContentType;
2404
2405 let tmp_dir = TempDir::new().unwrap();
2406 let dir = tmp_dir.path().to_str().unwrap();
2407 let data_path = write_plain_parquet(dir, "row_id_posdel_data.parquet", vec![], vec![]);
2409 let del_path = write_positional_delete(dir, "row_id_posdel.parquet", &data_path, &[1]);
2410
2411 let delete = FileScanTaskDeleteFile::builder()
2412 .with_file_path(del_path.clone())
2413 .with_file_format(DataFileFormat::Parquet)
2414 .with_file_size_in_bytes(std::fs::metadata(&del_path).unwrap().len())
2415 .with_file_type(DataContentType::PositionDeletes)
2416 .with_partition_spec_id(0)
2417 .build();
2418 let task = row_id_task_with_options(data_path, Some(100), 0, 0, vec![delete]);
2419
2420 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2421 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2422 let batches: Vec<RecordBatch> = reader
2423 .read(tasks)
2424 .unwrap()
2425 .stream()
2426 .try_collect()
2427 .await
2428 .unwrap();
2429
2430 assert_row_id_column(&batches, &[Some(100), Some(102)]);
2434 }
2435
2436 #[tokio::test]
2437 async fn test_row_id_name_collision_under_positional_fallback() {
2438 let tmp_dir = TempDir::new().unwrap();
2439 let dir = tmp_dir.path().to_str().unwrap();
2440
2441 let file_path = format!("{dir}/fallback_row_id_name.parquet");
2446 let arrow_schema = Arc::new(ArrowSchema::new(vec![
2447 Field::new("id", DataType::Int32, false),
2448 Field::new(RESERVED_COL_NAME_ROW_ID, DataType::Int64, true),
2449 ]));
2450 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![
2451 Arc::new(Int32Array::from(vec![1, 2, 3])),
2452 Arc::new(Int64Array::from(vec![7i64, 8, 9])),
2453 ])
2454 .unwrap();
2455 let props = WriterProperties::builder()
2456 .set_compression(Compression::SNAPPY)
2457 .build();
2458 let file = File::create(&file_path).unwrap();
2459 let mut writer = ArrowWriter::try_new(file, arrow_schema, Some(props)).unwrap();
2460 writer.write(&batch).unwrap();
2461 writer.close().unwrap();
2462
2463 let task = row_id_task(file_path, Some(100));
2464 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
2465 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2466 let batches: Vec<RecordBatch> = reader
2467 .read(tasks)
2468 .unwrap()
2469 .stream()
2470 .try_collect()
2471 .await
2472 .unwrap();
2473
2474 assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]);
2476 }
2477
2478 #[tokio::test]
2479 async fn test_read_encrypted_parquet_with_wrong_key_fails() {
2480 let encryption_key = b"0123456789abcdef";
2481 let wrong_key = b"fedcba9876543210";
2482
2483 let schema = Arc::new(
2484 Schema::builder()
2485 .with_schema_id(1)
2486 .with_fields(vec![
2487 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2488 ])
2489 .build()
2490 .unwrap(),
2491 );
2492
2493 let arrow_schema = Arc::new(ArrowSchema::new(vec![
2494 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2495 PARQUET_FIELD_ID_META_KEY.to_string(),
2496 "1".to_string(),
2497 )])),
2498 ]));
2499
2500 let tmp_dir = TempDir::new().unwrap();
2501 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2502 let file_io = FileIO::new_with_fs();
2503
2504 let id_data = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
2505 let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_data]).unwrap();
2506
2507 let file_path = format!("{table_location}/encrypted_wrong_key.parquet");
2508 write_encrypted_parquet(&file_path, &batch, encryption_key, None);
2509
2510 let wrong_key_metadata = crate::encryption::StandardKeyMetadata::try_new(wrong_key)
2511 .unwrap()
2512 .encode()
2513 .unwrap();
2514
2515 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
2516
2517 let task = FileScanTask::builder()
2518 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
2519 .with_start(0)
2520 .with_length(0)
2521 .with_data_file_path(file_path)
2522 .with_data_file_format(DataFileFormat::Parquet)
2523 .with_schema(schema)
2524 .with_project_field_ids(vec![1])
2525 .with_case_sensitive(false)
2526 .with_key_metadata(Some(wrong_key_metadata))
2527 .build()
2528 .unwrap();
2529
2530 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
2531 let result: Result<Vec<RecordBatch>, _> =
2532 reader.read(tasks).unwrap().stream().try_collect().await;
2533
2534 let err = result.unwrap_err();
2535 assert_eq!(err.kind(), crate::ErrorKind::Unexpected);
2536 let err_str = format!("{err}");
2537 assert!(
2538 err_str.contains("unable to decrypt parquet footer"),
2539 "Expected error about decryption failure, got: {err_str}"
2540 );
2541 }
2542
2543 #[tokio::test]
2546 async fn test_read_with_concurrency_one() {
2547 use arrow_array::Int32Array;
2548
2549 let schema = Arc::new(
2550 Schema::builder()
2551 .with_schema_id(1)
2552 .with_fields(vec![
2553 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
2554 NestedField::required(2, "file_num", Type::Primitive(PrimitiveType::Int))
2555 .into(),
2556 ])
2557 .build()
2558 .unwrap(),
2559 );
2560
2561 let arrow_schema = Arc::new(ArrowSchema::new(vec![
2562 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
2563 PARQUET_FIELD_ID_META_KEY.to_string(),
2564 "1".to_string(),
2565 )])),
2566 Field::new("file_num", DataType::Int32, false).with_metadata(HashMap::from([(
2567 PARQUET_FIELD_ID_META_KEY.to_string(),
2568 "2".to_string(),
2569 )])),
2570 ]));
2571
2572 let tmp_dir = TempDir::new().unwrap();
2573 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2574 let file_io = FileIO::new_with_fs();
2575
2576 let props = WriterProperties::builder()
2578 .set_compression(Compression::SNAPPY)
2579 .build();
2580
2581 for file_num in 0..3 {
2582 let id_data = Arc::new(Int32Array::from_iter_values(
2583 file_num * 10..(file_num + 1) * 10,
2584 )) as ArrayRef;
2585 let file_num_data = Arc::new(Int32Array::from(vec![file_num; 10])) as ArrayRef;
2586
2587 let to_write =
2588 RecordBatch::try_new(arrow_schema.clone(), vec![id_data, file_num_data]).unwrap();
2589
2590 let file = File::create(format!("{table_location}/file_{file_num}.parquet")).unwrap();
2591 let mut writer =
2592 ArrowWriter::try_new(file, to_write.schema(), Some(props.clone())).unwrap();
2593 writer.write(&to_write).expect("Writing batch");
2594 writer.close().unwrap();
2595 }
2596
2597 let reader = ArrowReaderBuilder::new(file_io, Runtime::current())
2599 .with_data_file_concurrency_limit(1)
2600 .build();
2601
2602 let tasks = vec![
2604 Ok(FileScanTask::builder()
2605 .with_file_size_in_bytes(
2606 std::fs::metadata(format!("{table_location}/file_0.parquet"))
2607 .unwrap()
2608 .len(),
2609 )
2610 .with_start(0)
2611 .with_length(0)
2612 .with_data_file_path(format!("{table_location}/file_0.parquet"))
2613 .with_data_file_format(DataFileFormat::Parquet)
2614 .with_schema(schema.clone())
2615 .with_project_field_ids(vec![1, 2])
2616 .with_case_sensitive(false)
2617 .build()
2618 .unwrap()),
2619 Ok(FileScanTask::builder()
2620 .with_file_size_in_bytes(
2621 std::fs::metadata(format!("{table_location}/file_1.parquet"))
2622 .unwrap()
2623 .len(),
2624 )
2625 .with_start(0)
2626 .with_length(0)
2627 .with_data_file_path(format!("{table_location}/file_1.parquet"))
2628 .with_data_file_format(DataFileFormat::Parquet)
2629 .with_schema(schema.clone())
2630 .with_project_field_ids(vec![1, 2])
2631 .with_case_sensitive(false)
2632 .build()
2633 .unwrap()),
2634 Ok(FileScanTask::builder()
2635 .with_file_size_in_bytes(
2636 std::fs::metadata(format!("{table_location}/file_2.parquet"))
2637 .unwrap()
2638 .len(),
2639 )
2640 .with_start(0)
2641 .with_length(0)
2642 .with_data_file_path(format!("{table_location}/file_2.parquet"))
2643 .with_data_file_format(DataFileFormat::Parquet)
2644 .with_schema(schema.clone())
2645 .with_project_field_ids(vec![1, 2])
2646 .with_case_sensitive(false)
2647 .build()
2648 .unwrap()),
2649 ];
2650
2651 let tasks_stream = Box::pin(futures::stream::iter(tasks)) as FileScanTaskStream;
2652
2653 let result = reader
2654 .read(tasks_stream)
2655 .unwrap()
2656 .stream()
2657 .try_collect::<Vec<RecordBatch>>()
2658 .await
2659 .unwrap();
2660
2661 let total_rows: usize = result.iter().map(|b| b.num_rows()).sum();
2663 assert_eq!(total_rows, 30, "Should have 30 total rows");
2664
2665 let mut all_ids = Vec::new();
2667 let mut all_file_nums = Vec::new();
2668
2669 for batch in &result {
2670 let id_col = batch
2671 .column(0)
2672 .as_primitive::<arrow_array::types::Int32Type>();
2673 let file_num_col = batch
2674 .column(1)
2675 .as_primitive::<arrow_array::types::Int32Type>();
2676
2677 for i in 0..batch.num_rows() {
2678 all_ids.push(id_col.value(i));
2679 all_file_nums.push(file_num_col.value(i));
2680 }
2681 }
2682
2683 assert_eq!(all_ids.len(), 30);
2684 assert_eq!(all_file_nums.len(), 30);
2685
2686 for i in 0..10 {
2691 assert_eq!(all_file_nums[i], 0, "First 10 rows should be from file_0");
2692 assert_eq!(all_ids[i], i as i32, "IDs should be 0-9");
2693 }
2694 for i in 10..20 {
2695 assert_eq!(all_file_nums[i], 1, "Next 10 rows should be from file_1");
2696 assert_eq!(all_ids[i], i as i32, "IDs should be 10-19");
2697 }
2698 for i in 20..30 {
2699 assert_eq!(all_file_nums[i], 2, "Last 10 rows should be from file_2");
2700 assert_eq!(all_ids[i], i as i32, "IDs should be 20-29");
2701 }
2702 }
2703
2704 #[tokio::test]
2705 async fn test_read_int96_timestamps_with_field_ids() {
2706 let schema = Arc::new(
2707 Schema::builder()
2708 .with_schema_id(1)
2709 .with_fields(vec![
2710 NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2711 .into(),
2712 NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2713 ])
2714 .build()
2715 .unwrap(),
2716 );
2717
2718 let tmp_dir = TempDir::new().unwrap();
2719 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2720 let (file_path, expected_micros) =
2721 write_int96_parquet_file(&table_location, "with_ids.parquet", true);
2722
2723 assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await;
2724 }
2725
2726 #[tokio::test]
2727 async fn test_read_int96_timestamps_without_field_ids() {
2728 let schema = Arc::new(
2729 Schema::builder()
2730 .with_schema_id(1)
2731 .with_fields(vec![
2732 NestedField::optional(1, "ts", Type::Primitive(PrimitiveType::Timestamp))
2733 .into(),
2734 NestedField::required(2, "id", Type::Primitive(PrimitiveType::Int)).into(),
2735 ])
2736 .build()
2737 .unwrap(),
2738 );
2739
2740 let tmp_dir = TempDir::new().unwrap();
2741 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2742 let (file_path, expected_micros) =
2743 write_int96_parquet_file(&table_location, "no_ids.parquet", false);
2744
2745 assert_int96_read_matches(&file_path, schema, vec![1, 2], &expected_micros).await;
2746 }
2747
2748 #[tokio::test]
2749 async fn test_read_int96_timestamps_in_struct() {
2750 use arrow_array::{StructArray, TimestampMicrosecondArray};
2751 use parquet::basic::{Repetition, Type as PhysicalType};
2752 use parquet::data_type::Int96Type;
2753 use parquet::file::writer::SerializedFileWriter;
2754 use parquet::schema::types::Type as SchemaType;
2755
2756 let tmp_dir = TempDir::new().unwrap();
2757 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2758 let file_path = format!("{table_location}/struct_int96.parquet");
2759
2760 let ts_type = SchemaType::primitive_type_builder("ts", PhysicalType::INT96)
2761 .with_repetition(Repetition::OPTIONAL)
2762 .with_id(Some(2))
2763 .build()
2764 .unwrap();
2765
2766 let struct_type = SchemaType::group_type_builder("data")
2767 .with_repetition(Repetition::REQUIRED)
2768 .with_id(Some(1))
2769 .with_fields(vec![Arc::new(ts_type)])
2770 .build()
2771 .unwrap();
2772
2773 let parquet_schema = SchemaType::group_type_builder("schema")
2774 .with_fields(vec![Arc::new(struct_type)])
2775 .build()
2776 .unwrap();
2777
2778 let (int96_val, expected_micros) = make_int96_test_value();
2779
2780 let file = File::create(&file_path).unwrap();
2781 let mut writer =
2782 SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
2783
2784 let mut row_group = writer.next_row_group().unwrap();
2787 {
2788 let mut col = row_group.next_column().unwrap().unwrap();
2789 col.typed::<Int96Type>()
2790 .write_batch(&[int96_val], Some(&[1]), None)
2791 .unwrap();
2792 col.close().unwrap();
2793 }
2794 row_group.close().unwrap();
2795 writer.close().unwrap();
2796
2797 let iceberg_schema = Arc::new(
2798 Schema::builder()
2799 .with_schema_id(1)
2800 .with_fields(vec![
2801 NestedField::required(
2802 1,
2803 "data",
2804 Type::Struct(crate::spec::StructType::new(vec![
2805 NestedField::optional(
2806 2,
2807 "ts",
2808 Type::Primitive(PrimitiveType::Timestamp),
2809 )
2810 .into(),
2811 ])),
2812 )
2813 .into(),
2814 ])
2815 .build()
2816 .unwrap(),
2817 );
2818
2819 let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
2820
2821 assert_eq!(batches.len(), 1);
2822 let struct_array = batches[0]
2823 .column(0)
2824 .as_any()
2825 .downcast_ref::<StructArray>()
2826 .expect("Expected StructArray");
2827 let ts_array = struct_array
2828 .column(0)
2829 .as_any()
2830 .downcast_ref::<TimestampMicrosecondArray>()
2831 .expect("Expected TimestampMicrosecondArray inside struct");
2832
2833 assert_eq!(
2834 ts_array.value(0),
2835 expected_micros,
2836 "INT96 in struct: got {}, expected {expected_micros}",
2837 ts_array.value(0)
2838 );
2839 }
2840
2841 #[tokio::test]
2842 async fn test_read_int96_timestamps_in_list() {
2843 use arrow_array::{ListArray, TimestampMicrosecondArray};
2844 use parquet::basic::{Repetition, Type as PhysicalType};
2845 use parquet::data_type::Int96Type;
2846 use parquet::file::writer::SerializedFileWriter;
2847 use parquet::schema::types::Type as SchemaType;
2848
2849 let tmp_dir = TempDir::new().unwrap();
2850 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2851 let file_path = format!("{table_location}/list_int96.parquet");
2852
2853 let element_type = SchemaType::primitive_type_builder("element", PhysicalType::INT96)
2860 .with_repetition(Repetition::OPTIONAL)
2861 .with_id(Some(2))
2862 .build()
2863 .unwrap();
2864
2865 let list_group = SchemaType::group_type_builder("list")
2866 .with_repetition(Repetition::REPEATED)
2867 .with_fields(vec![Arc::new(element_type)])
2868 .build()
2869 .unwrap();
2870
2871 let list_type = SchemaType::group_type_builder("timestamps")
2872 .with_repetition(Repetition::OPTIONAL)
2873 .with_id(Some(1))
2874 .with_logical_type(Some(parquet::basic::LogicalType::List))
2875 .with_fields(vec![Arc::new(list_group)])
2876 .build()
2877 .unwrap();
2878
2879 let parquet_schema = SchemaType::group_type_builder("schema")
2880 .with_fields(vec![Arc::new(list_type)])
2881 .build()
2882 .unwrap();
2883
2884 let (int96_val, expected_micros) = make_int96_test_value();
2885
2886 let file = File::create(&file_path).unwrap();
2887 let mut writer =
2888 SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
2889
2890 let mut row_group = writer.next_row_group().unwrap();
2894 {
2895 let mut col = row_group.next_column().unwrap().unwrap();
2896 col.typed::<Int96Type>()
2897 .write_batch(&[int96_val], Some(&[3]), Some(&[0]))
2898 .unwrap();
2899 col.close().unwrap();
2900 }
2901 row_group.close().unwrap();
2902 writer.close().unwrap();
2903
2904 let iceberg_schema = Arc::new(
2905 Schema::builder()
2906 .with_schema_id(1)
2907 .with_fields(vec![
2908 NestedField::optional(
2909 1,
2910 "timestamps",
2911 Type::List(crate::spec::ListType {
2912 element_field: NestedField::optional(
2913 2,
2914 "element",
2915 Type::Primitive(PrimitiveType::Timestamp),
2916 )
2917 .into(),
2918 }),
2919 )
2920 .into(),
2921 ])
2922 .build()
2923 .unwrap(),
2924 );
2925
2926 let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
2927
2928 assert_eq!(batches.len(), 1);
2929 let list_array = batches[0]
2930 .column(0)
2931 .as_any()
2932 .downcast_ref::<ListArray>()
2933 .expect("Expected ListArray");
2934 let ts_array = list_array
2935 .values()
2936 .as_any()
2937 .downcast_ref::<TimestampMicrosecondArray>()
2938 .expect("Expected TimestampMicrosecondArray inside list");
2939
2940 assert_eq!(
2941 ts_array.value(0),
2942 expected_micros,
2943 "INT96 in list: got {}, expected {expected_micros}",
2944 ts_array.value(0)
2945 );
2946 }
2947
2948 #[tokio::test]
2949 async fn test_read_int96_timestamps_in_map() {
2950 use arrow_array::{MapArray, TimestampMicrosecondArray};
2951 use parquet::basic::{Repetition, Type as PhysicalType};
2952 use parquet::data_type::{ByteArrayType, Int96Type};
2953 use parquet::file::writer::SerializedFileWriter;
2954 use parquet::schema::types::Type as SchemaType;
2955
2956 let tmp_dir = TempDir::new().unwrap();
2957 let table_location = tmp_dir.path().to_str().unwrap().to_string();
2958 let file_path = format!("{table_location}/map_int96.parquet");
2959
2960 let key_type = SchemaType::primitive_type_builder("key", PhysicalType::BYTE_ARRAY)
2968 .with_repetition(Repetition::REQUIRED)
2969 .with_logical_type(Some(parquet::basic::LogicalType::String))
2970 .with_id(Some(2))
2971 .build()
2972 .unwrap();
2973
2974 let value_type = SchemaType::primitive_type_builder("value", PhysicalType::INT96)
2975 .with_repetition(Repetition::OPTIONAL)
2976 .with_id(Some(3))
2977 .build()
2978 .unwrap();
2979
2980 let key_value_group = SchemaType::group_type_builder("key_value")
2981 .with_repetition(Repetition::REPEATED)
2982 .with_fields(vec![Arc::new(key_type), Arc::new(value_type)])
2983 .build()
2984 .unwrap();
2985
2986 let map_type = SchemaType::group_type_builder("ts_map")
2987 .with_repetition(Repetition::OPTIONAL)
2988 .with_id(Some(1))
2989 .with_logical_type(Some(parquet::basic::LogicalType::Map))
2990 .with_fields(vec![Arc::new(key_value_group)])
2991 .build()
2992 .unwrap();
2993
2994 let parquet_schema = SchemaType::group_type_builder("schema")
2995 .with_fields(vec![Arc::new(map_type)])
2996 .build()
2997 .unwrap();
2998
2999 let (int96_val, expected_micros) = make_int96_test_value();
3000
3001 let file = File::create(&file_path).unwrap();
3002 let mut writer =
3003 SerializedFileWriter::new(file, Arc::new(parquet_schema), Default::default()).unwrap();
3004
3005 let mut row_group = writer.next_row_group().unwrap();
3010 {
3011 let mut col = row_group.next_column().unwrap().unwrap();
3012 col.typed::<ByteArrayType>()
3013 .write_batch(
3014 &[parquet::data_type::ByteArray::from("event_time")],
3015 Some(&[2]),
3016 Some(&[0]),
3017 )
3018 .unwrap();
3019 col.close().unwrap();
3020 }
3021 {
3022 let mut col = row_group.next_column().unwrap().unwrap();
3023 col.typed::<Int96Type>()
3024 .write_batch(&[int96_val], Some(&[3]), Some(&[0]))
3025 .unwrap();
3026 col.close().unwrap();
3027 }
3028 row_group.close().unwrap();
3029 writer.close().unwrap();
3030
3031 let iceberg_schema = Arc::new(
3032 Schema::builder()
3033 .with_schema_id(1)
3034 .with_fields(vec![
3035 NestedField::optional(
3036 1,
3037 "ts_map",
3038 Type::Map(crate::spec::MapType {
3039 key_field: NestedField::required(
3040 2,
3041 "key",
3042 Type::Primitive(PrimitiveType::String),
3043 )
3044 .into(),
3045 value_field: NestedField::optional(
3046 3,
3047 "value",
3048 Type::Primitive(PrimitiveType::Timestamp),
3049 )
3050 .into(),
3051 }),
3052 )
3053 .into(),
3054 ])
3055 .build()
3056 .unwrap(),
3057 );
3058
3059 let batches = read_int96_batches(&file_path, iceberg_schema, vec![1]).await;
3060
3061 assert_eq!(batches.len(), 1);
3062 let map_array = batches[0]
3063 .column(0)
3064 .as_any()
3065 .downcast_ref::<MapArray>()
3066 .expect("Expected MapArray");
3067 let ts_array = map_array
3068 .values()
3069 .as_any()
3070 .downcast_ref::<TimestampMicrosecondArray>()
3071 .expect("Expected TimestampMicrosecondArray as map values");
3072
3073 assert_eq!(
3074 ts_array.value(0),
3075 expected_micros,
3076 "INT96 in map: got {}, expected {expected_micros}",
3077 ts_array.value(0)
3078 );
3079 }
3080
3081 fn write_parquet_with_wide_column(
3087 dir: &str,
3088 name: &str,
3089 extra_fields: Vec<Field>,
3090 extra_columns: Vec<ArrayRef>,
3091 ) -> String {
3092 let wide_field =
3093 Field::new("wide", DataType::Utf8, false).with_metadata(HashMap::from([(
3094 PARQUET_FIELD_ID_META_KEY.to_string(),
3095 "2".to_string(),
3096 )]));
3097 let wide_values: Vec<String> = (0..3)
3100 .map(|i| {
3101 (0..2048)
3102 .map(|j| ((i * 2048 + j) % 251) as u8 as char)
3103 .collect()
3104 })
3105 .collect();
3106
3107 let mut fields = vec![wide_field];
3108 fields.extend(extra_fields);
3109 let mut columns: Vec<ArrayRef> = vec![Arc::new(StringArray::from(wide_values))];
3110 columns.extend(extra_columns);
3111 write_plain_parquet(dir, name, fields, columns)
3112 }
3113
3114 fn id_and_wide_schema() -> SchemaRef {
3117 Arc::new(
3118 Schema::builder()
3119 .with_schema_id(1)
3120 .with_fields(vec![
3121 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3122 NestedField::required(2, "wide", Type::Primitive(PrimitiveType::String)).into(),
3123 ])
3124 .build()
3125 .unwrap(),
3126 )
3127 }
3128
3129 fn metadata_projection_task(
3131 file_path: String,
3132 schema: SchemaRef,
3133 project_field_ids: Vec<i32>,
3134 ) -> FileScanTask {
3135 metadata_projection_task_with_first_row_id(file_path, schema, project_field_ids, None)
3136 }
3137
3138 fn metadata_projection_task_with_first_row_id(
3139 file_path: String,
3140 schema: SchemaRef,
3141 project_field_ids: Vec<i32>,
3142 first_row_id: Option<i64>,
3143 ) -> FileScanTask {
3144 FileScanTask::builder()
3145 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3146 .with_start(0)
3147 .with_length(0)
3148 .with_data_file_path(file_path)
3149 .with_data_file_format(DataFileFormat::Parquet)
3150 .with_schema(schema)
3151 .with_project_field_ids(project_field_ids)
3152 .with_first_row_id(first_row_id)
3153 .with_case_sensitive(false)
3154 .build()
3155 .unwrap()
3156 }
3157
3158 async fn scan_task(task: FileScanTask) -> (Vec<RecordBatch>, u64) {
3160 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build();
3161 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
3162 let scan = reader.read(tasks).unwrap();
3163 let metrics = scan.metrics().clone();
3164 let batches = scan.stream().try_collect().await.unwrap();
3165 (batches, metrics.bytes_read())
3166 }
3167
3168 #[tokio::test]
3169 async fn test_pos_only_projection_reads_no_data_columns() {
3170 let tmp_dir = TempDir::new().unwrap();
3171 let dir = tmp_dir.path().to_str().unwrap();
3172
3173 let pos_only = metadata_projection_task(
3174 write_parquet_with_wide_column(dir, "pos_only.parquet", vec![], vec![]),
3175 id_and_wide_schema(),
3176 vec![RESERVED_FIELD_ID_POS],
3177 );
3178 let (batches, pos_only_bytes) = scan_task(pos_only).await;
3179
3180 assert_eq!(batches[0].num_columns(), 1);
3182 let pos_col = batches[0]
3183 .column_by_name(RESERVED_COL_NAME_POS)
3184 .expect("_pos column should be present")
3185 .as_primitive::<arrow_array::types::Int64Type>();
3186 assert_eq!(pos_col.values(), &[0, 1, 2]);
3187
3188 let with_data = metadata_projection_task(
3191 write_parquet_with_wide_column(dir, "pos_only_ref.parquet", vec![], vec![]),
3192 id_and_wide_schema(),
3193 vec![2, RESERVED_FIELD_ID_POS],
3194 );
3195 let (_, with_data_bytes) = scan_task(with_data).await;
3196
3197 assert!(
3198 pos_only_bytes < with_data_bytes,
3199 "_pos-only scan should read fewer bytes than a scan of the wide column: \
3200 {pos_only_bytes} vs {with_data_bytes}"
3201 );
3202 }
3203
3204 #[tokio::test]
3205 async fn test_pos_only_projection_keeps_absolute_pos_under_predicate() {
3206 use crate::expr::{Bind, Reference};
3207 use crate::spec::Datum;
3208
3209 let tmp_dir = TempDir::new().unwrap();
3210 let dir = tmp_dir.path().to_str().unwrap();
3211 let file_path = write_plain_parquet(dir, "pos_only_predicate.parquet", vec![], vec![]);
3213
3214 let schema = Arc::new(
3215 Schema::builder()
3216 .with_schema_id(1)
3217 .with_fields(vec![
3218 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3219 ])
3220 .build()
3221 .unwrap(),
3222 );
3223 let bound = Reference::new("id")
3224 .not_equal_to(Datum::int(2))
3225 .bind(Arc::clone(&schema), false)
3226 .unwrap();
3227 let task = FileScanTask::builder()
3228 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3229 .with_start(0)
3230 .with_length(0)
3231 .with_data_file_path(file_path)
3232 .with_data_file_format(DataFileFormat::Parquet)
3233 .with_schema(schema)
3234 .with_project_field_ids(vec![RESERVED_FIELD_ID_POS])
3235 .with_predicate(Some(bound))
3236 .with_case_sensitive(false)
3237 .build()
3238 .unwrap();
3239
3240 let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current())
3244 .with_row_selection_enabled(true)
3245 .build();
3246 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
3247 let batches: Vec<RecordBatch> = reader
3248 .read(tasks)
3249 .unwrap()
3250 .stream()
3251 .try_collect()
3252 .await
3253 .unwrap();
3254
3255 let pos: Vec<i64> = batches
3256 .iter()
3257 .flat_map(|b| {
3258 b.column_by_name(RESERVED_COL_NAME_POS)
3259 .expect("_pos column should be present")
3260 .as_primitive::<arrow_array::types::Int64Type>()
3261 .values()
3262 .to_vec()
3263 })
3264 .collect();
3265 assert_eq!(pos, vec![0, 2]);
3266 }
3267
3268 #[tokio::test]
3269 async fn test_pos_and_file_projection() {
3270 use crate::metadata_columns::RESERVED_COL_NAME_FILE;
3271
3272 let tmp_dir = TempDir::new().unwrap();
3273 let dir = tmp_dir.path().to_str().unwrap();
3274 let file_path = write_parquet_with_wide_column(dir, "pos_and_file.parquet", vec![], vec![]);
3277 let task = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![
3278 RESERVED_FIELD_ID_POS,
3279 RESERVED_FIELD_ID_FILE,
3280 ]);
3281 let (batches, _) = scan_task(task).await;
3282
3283 assert_eq!(batches[0].num_columns(), 2);
3285 let pos_col = batches[0]
3286 .column_by_name(RESERVED_COL_NAME_POS)
3287 .expect("_pos column should be present")
3288 .as_primitive::<arrow_array::types::Int64Type>();
3289 assert_eq!(pos_col.values(), &[0, 1, 2]);
3290 let file_col = batches[0]
3291 .column_by_name(RESERVED_COL_NAME_FILE)
3292 .expect("_file column should be present");
3293 let file_col = cast(file_col, &DataType::Utf8).unwrap();
3294 let file_col = file_col.as_any().downcast_ref::<StringArray>().unwrap();
3295 assert_eq!(file_col.value(0), file_path);
3296 }
3297
3298 #[tokio::test]
3299 async fn test_pos_and_physical_seq_projection_reads_only_the_leaf() {
3300 use crate::metadata_columns::{
3301 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3302 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3303 };
3304
3305 let tmp_dir = TempDir::new().unwrap();
3309 let dir = tmp_dir.path().to_str().unwrap();
3310
3311 let write = |name: &str| {
3313 write_parquet_with_wide_column(
3314 dir,
3315 name,
3316 vec![physical_last_updated_seq_field()],
3317 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
3318 )
3319 };
3320
3321 let seq_task = |path: String, ids: Vec<i32>| {
3322 FileScanTask::builder()
3323 .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len())
3324 .with_start(0)
3325 .with_length(0)
3326 .with_data_file_path(path)
3327 .with_data_file_format(DataFileFormat::Parquet)
3328 .with_schema(id_and_wide_schema())
3329 .with_project_field_ids(ids)
3330 .with_first_row_id(Some(100))
3331 .with_data_sequence_number(Some(9))
3332 .with_case_sensitive(false)
3333 .build()
3334 .unwrap()
3335 };
3336
3337 let meta_only = seq_task(write("pos_seq.parquet"), vec![
3338 RESERVED_FIELD_ID_POS,
3339 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3340 ]);
3341 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3342
3343 let pos_col = batches[0]
3345 .column_by_name(RESERVED_COL_NAME_POS)
3346 .expect("_pos column should be present")
3347 .as_primitive::<arrow_array::types::Int64Type>();
3348 assert_eq!(pos_col.values(), &[0, 1, 2]);
3349 let seq_col = batches[0]
3350 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3351 .expect("_last_updated_sequence_number column should be present");
3352 let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3353 let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3354 assert_eq!(seq_col.value(0), 5);
3356 assert_eq!(seq_col.value(1), 9);
3357 assert_eq!(seq_col.value(2), 8);
3358 assert!(batches[0].column_by_name("wide").is_none());
3359
3360 let with_data = seq_task(write("pos_seq_ref.parquet"), vec![
3363 2,
3364 RESERVED_FIELD_ID_POS,
3365 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3366 ]);
3367 let (_, with_data_bytes) = scan_task(with_data).await;
3368
3369 assert!(
3370 meta_only_bytes < with_data_bytes,
3371 "_pos + physical sequence scan should read fewer bytes than one that also \
3372 reads the wide column: {meta_only_bytes} vs {with_data_bytes}"
3373 );
3374 }
3375
3376 #[tokio::test]
3377 async fn test_seq_only_projection_reads_only_the_leaf() {
3378 use crate::metadata_columns::{
3379 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3380 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3381 };
3382
3383 let tmp_dir = TempDir::new().unwrap();
3387 let dir = tmp_dir.path().to_str().unwrap();
3388 let write = |name: &str| {
3389 write_parquet_with_wide_column(
3390 dir,
3391 name,
3392 vec![physical_last_updated_seq_field()],
3393 vec![Arc::new(Int64Array::from(vec![Some(5), None, Some(8)])) as ArrayRef],
3394 )
3395 };
3396 let seq_task = |path: String, ids: Vec<i32>| {
3397 FileScanTask::builder()
3398 .with_file_size_in_bytes(std::fs::metadata(&path).unwrap().len())
3399 .with_start(0)
3400 .with_length(0)
3401 .with_data_file_path(path)
3402 .with_data_file_format(DataFileFormat::Parquet)
3403 .with_schema(id_and_wide_schema())
3404 .with_project_field_ids(ids)
3405 .with_first_row_id(Some(100))
3406 .with_data_sequence_number(Some(9))
3407 .with_case_sensitive(false)
3408 .build()
3409 .unwrap()
3410 };
3411
3412 let meta_only = seq_task(write("seq_only.parquet"), vec![
3413 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3414 ]);
3415 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3416
3417 let seq_col = batches[0]
3418 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3419 .expect("_last_updated_sequence_number column should be present");
3420 let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3421 let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3422 assert_eq!(seq_col.value(0), 5);
3423 assert_eq!(seq_col.value(1), 9);
3424 assert_eq!(seq_col.value(2), 8);
3425 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3426 assert_eq!(total_rows, 3);
3427 assert!(batches[0].column_by_name("wide").is_none());
3428
3429 let with_data = seq_task(write("seq_only_ref.parquet"), vec![
3430 2,
3431 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3432 ]);
3433 let (_, with_data_bytes) = scan_task(with_data).await;
3434
3435 assert!(
3436 meta_only_bytes < with_data_bytes,
3437 "seq-only scan should read fewer bytes than one that also reads the wide \
3438 column: {meta_only_bytes} vs {with_data_bytes}"
3439 );
3440 }
3441
3442 #[tokio::test]
3443 async fn test_seq_only_null_first_row_id_reads_no_data_columns() {
3444 use crate::metadata_columns::{
3445 RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER,
3446 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3447 };
3448
3449 let tmp_dir = TempDir::new().unwrap();
3454 let dir = tmp_dir.path().to_str().unwrap();
3455 let seq_task = |name: &str, ids: Vec<i32>| {
3456 let file_path = write_parquet_with_wide_column(
3457 dir,
3458 name,
3459 vec![physical_last_updated_seq_field()],
3460 vec![Arc::new(Int64Array::from(vec![Some(5), Some(6), Some(7)])) as ArrayRef],
3461 );
3462 FileScanTask::builder()
3463 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3464 .with_start(0)
3465 .with_length(0)
3466 .with_data_file_path(file_path)
3467 .with_data_file_format(DataFileFormat::Parquet)
3468 .with_schema(id_and_wide_schema())
3469 .with_project_field_ids(ids)
3470 .with_first_row_id(None)
3471 .with_data_sequence_number(Some(9))
3472 .with_case_sensitive(false)
3473 .build()
3474 .unwrap()
3475 };
3476
3477 let meta_only = seq_task("seq_only_null_first.parquet", vec![
3478 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3479 ]);
3480 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3481
3482 assert_eq!(batches[0].num_columns(), 1);
3483 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3484 assert_eq!(total_rows, 3);
3485 let seq_col = batches[0]
3486 .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER)
3487 .expect("_last_updated_sequence_number column should be present");
3488 let seq_col = cast(seq_col, &DataType::Int64).unwrap();
3489 let seq_col = seq_col.as_any().downcast_ref::<Int64Array>().unwrap();
3490 assert!((0..3).all(|i| seq_col.is_null(i)));
3491
3492 let with_data = seq_task("seq_only_null_first_ref.parquet", vec![
3493 2,
3494 RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER,
3495 ]);
3496 let (_, with_data_bytes) = scan_task(with_data).await;
3497
3498 assert!(
3499 meta_only_bytes < with_data_bytes,
3500 "seq-only scan with a null first_row_id should read fewer bytes than a scan of \
3501 the wide column: {meta_only_bytes} vs {with_data_bytes}"
3502 );
3503 }
3504
3505 #[tokio::test]
3506 async fn test_file_only_reads_no_data_columns() {
3507 use crate::metadata_columns::RESERVED_COL_NAME_FILE;
3508
3509 let tmp_dir = TempDir::new().unwrap();
3513 let dir = tmp_dir.path().to_str().unwrap();
3514
3515 let file_path = write_parquet_with_wide_column(dir, "file_only.parquet", vec![], vec![]);
3516 let meta_only = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![
3517 RESERVED_FIELD_ID_FILE,
3518 ]);
3519 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3520
3521 assert_eq!(batches[0].num_columns(), 1);
3522 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3523 assert_eq!(total_rows, 3);
3524 let file_col = batches[0]
3525 .column_by_name(RESERVED_COL_NAME_FILE)
3526 .expect("_file column should be present");
3527 let file_col = cast(file_col, &DataType::Utf8).unwrap();
3528 let file_col = file_col.as_any().downcast_ref::<StringArray>().unwrap();
3529 assert_eq!(file_col.value(0), file_path);
3530
3531 let with_data = metadata_projection_task(
3532 write_parquet_with_wide_column(dir, "file_only_ref.parquet", vec![], vec![]),
3533 id_and_wide_schema(),
3534 vec![2, RESERVED_FIELD_ID_FILE],
3535 );
3536 let (_, with_data_bytes) = scan_task(with_data).await;
3537
3538 assert!(
3539 meta_only_bytes < with_data_bytes,
3540 "_file-only scan should read fewer bytes than a scan of the wide column: \
3541 {meta_only_bytes} vs {with_data_bytes}"
3542 );
3543 }
3544
3545 #[tokio::test]
3546 async fn test_spec_id_only_reads_no_data_columns() {
3547 use crate::metadata_columns::{RESERVED_COL_NAME_SPEC_ID, RESERVED_FIELD_ID_SPEC_ID};
3548 use crate::spec::{Literal, PartitionSpec, Struct, Transform};
3549
3550 let tmp_dir = TempDir::new().unwrap();
3554 let dir = tmp_dir.path().to_str().unwrap();
3555 let schema = id_and_wide_schema();
3556 let spec = Arc::new(
3557 PartitionSpec::builder(schema.clone())
3558 .with_spec_id(7)
3559 .add_partition_field("id", "id", Transform::Identity)
3560 .unwrap()
3561 .build()
3562 .unwrap(),
3563 );
3564
3565 let spec_id_task = |file_path: String, project_field_ids: Vec<i32>| {
3566 FileScanTask::builder()
3567 .with_file_size_in_bytes(std::fs::metadata(&file_path).unwrap().len())
3568 .with_start(0)
3569 .with_length(0)
3570 .with_data_file_path(file_path)
3571 .with_data_file_format(DataFileFormat::Parquet)
3572 .with_schema(schema.clone())
3573 .with_project_field_ids(project_field_ids)
3574 .with_partition(Some(Struct::from_iter([Some(Literal::int(42))])))
3575 .with_partition_spec(Some(spec.clone()))
3576 .with_case_sensitive(false)
3577 .build()
3578 .unwrap()
3579 };
3580
3581 let meta_only = spec_id_task(
3582 write_parquet_with_wide_column(dir, "spec_id_only.parquet", vec![], vec![]),
3583 vec![RESERVED_FIELD_ID_SPEC_ID],
3584 );
3585 let (batches, meta_only_bytes) = scan_task(meta_only).await;
3586
3587 assert_eq!(batches[0].num_columns(), 1);
3588 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3589 assert_eq!(total_rows, 3);
3590 let spec_id_col = batches[0]
3591 .column_by_name(RESERVED_COL_NAME_SPEC_ID)
3592 .expect("_spec_id column should be present");
3593 let spec_id_col = cast(spec_id_col, &DataType::Int32).unwrap();
3594 let spec_id_col = spec_id_col.as_any().downcast_ref::<Int32Array>().unwrap();
3595 assert_eq!(spec_id_col.values(), &[7, 7, 7]);
3596
3597 let with_data = spec_id_task(
3598 write_parquet_with_wide_column(dir, "spec_id_only_ref.parquet", vec![], vec![]),
3599 vec![2, RESERVED_FIELD_ID_SPEC_ID],
3600 );
3601 let (_, with_data_bytes) = scan_task(with_data).await;
3602
3603 assert!(
3604 meta_only_bytes < with_data_bytes,
3605 "_spec_id-only scan should read fewer bytes than a scan of the wide column: \
3606 {meta_only_bytes} vs {with_data_bytes}"
3607 );
3608 }
3609
3610 #[tokio::test]
3611 async fn test_empty_projection_preserves_row_count() {
3612 let tmp_dir = TempDir::new().unwrap();
3613 let dir = tmp_dir.path().to_str().unwrap();
3614 let file_path = write_plain_parquet(dir, "empty_projection.parquet", vec![], vec![]);
3615 let schema = Arc::new(
3616 Schema::builder()
3617 .with_schema_id(1)
3618 .with_fields(vec![
3619 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
3620 ])
3621 .build()
3622 .unwrap(),
3623 );
3624 let task = metadata_projection_task(file_path, schema, vec![]);
3625 let (batches, _) = scan_task(task).await;
3626
3627 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
3629 assert_eq!(total_rows, 3);
3630 }
3631}