diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 1c5af7cae0..a577a9d593 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -96,6 +96,11 @@ pub fn iceberg::arrow::RecordBatchPartitionSplitter::try_new_with_precomputed_va pub struct iceberg::arrow::ScanMetrics impl iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::bytes_read(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::data_bytes_read(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::data_files_opened(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::delete_bytes_read(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::delete_files_opened(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::rows_emitted(&self) -> u64 impl core::clone::Clone for iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::clone(&self) -> iceberg::scan::ScanMetrics impl core::fmt::Debug for iceberg::scan::ScanMetrics @@ -1427,6 +1432,11 @@ pub fn iceberg::scan::FileScanTaskDeleteFile::deserialize<__D>(__deserializer: _ pub struct iceberg::scan::ScanMetrics impl iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::bytes_read(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::data_bytes_read(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::data_files_opened(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::delete_bytes_read(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::delete_files_opened(&self) -> u64 +pub fn iceberg::scan::ScanMetrics::rows_emitted(&self) -> u64 impl core::clone::Clone for iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::clone(&self) -> iceberg::scan::ScanMetrics impl core::fmt::Debug for iceberg::scan::ScanMetrics diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 0711a58b35..4a53466598 100644 --- a/crates/iceberg/src/arrow/caching_delete_file_loader.rs +++ b/crates/iceberg/src/arrow/caching_delete_file_loader.rs @@ -25,14 +25,14 @@ use tokio::sync::oneshot::{Receiver, channel}; use super::delete_filter::{DeleteFilter, PosDelLoadAction, PosDelLoadGuard}; use crate::arrow::delete_file_loader::BasicDeleteFileLoader; -use crate::arrow::scan_metrics::ScanMetrics; +use crate::arrow::scan_metrics::{CountingFileRead, ScanMetrics}; use crate::arrow::{arrow_primitive_to_literal, arrow_schema_to_schema}; use crate::delete_vector::DeleteVector; use crate::encryption::{EncryptedInputFile, StandardKeyMetadata}; use crate::error::invalid_data; use crate::expr::Predicate::AlwaysTrue; use crate::expr::{Predicate, Reference}; -use crate::io::FileIO; +use crate::io::{FileIO, FileRead}; use crate::runtime::Runtime; use crate::scan::{ArrowRecordBatchStream, FileScanTaskDeleteFile}; use crate::spec::{ @@ -289,6 +289,7 @@ impl CachingDeleteFileLoader { .parquet_to_batch_stream( task.file_path(), task.file_size_in_bytes(), + task.file_type(), task.key_metadata(), ) .await @@ -321,6 +322,7 @@ impl CachingDeleteFileLoader { .parquet_to_batch_stream( task.file_path(), task.file_size_in_bytes(), + task.file_type(), task.key_metadata(), ) .await?, @@ -421,18 +423,33 @@ impl CachingDeleteFileLoader { let input_file = basic_delete_file_loader .file_io() .new_input(task.file_path())?; - let blob = match task.key_metadata() { + let key_metadata = task + .key_metadata() + .map(StandardKeyMetadata::decode) + .transpose()?; + let counted_reader: Box = Box::new(CountingFileRead::new( + input_file.reader().await?, + basic_delete_file_loader.scan_metrics().clone(), + task.file_type(), + )); + let reader = match key_metadata { Some(key_metadata) => { - let key_metadata = StandardKeyMetadata::decode(key_metadata)?; - EncryptedInputFile::new(input_file, key_metadata) - .reader() - .await? - .read(start..start + len) - .await? + /* + * EncryptedInputFile reads bytes from the `input_file`, decrypts them and then reads the provided byte range + * Passing a counted_reader using the decrypt_reader helper manually + * allows us to count the number of bytes read from storage *before* decryption + */ + EncryptedInputFile::new(input_file, key_metadata).decrypt_reader(counted_reader)? } - None => input_file.reader().await?.read(start..start + len).await?, + None => counted_reader, }; + basic_delete_file_loader + .scan_metrics() + .add_delete_files_opened(1); + + let blob = reader.read(start..start + len).await?; + Ok(DeleteFileContext::DelVec { data_file_path, blob, @@ -891,6 +908,7 @@ mod tests { .parquet_to_batch_stream( &eq_delete_file_path, std::fs::metadata(&eq_delete_file_path).unwrap().len(), + DataContentType::EqualityDeletes, None, ) .await @@ -1330,6 +1348,7 @@ mod tests { .parquet_to_batch_stream( &delete_file_path, std::fs::metadata(&delete_file_path).unwrap().len(), + DataContentType::EqualityDeletes, None, ) .await @@ -1455,8 +1474,10 @@ mod tests { .unwrap(); // Load the deletes - should handle both types without error + let scan_metrics = ScanMetrics::new(); let delete_file_loader = - CachingDeleteFileLoader::new(file_io.clone(), 10, Runtime::current()); + CachingDeleteFileLoader::new(file_io.clone(), 10, Runtime::current()) + .with_scan_metrics(scan_metrics.clone()); let delete_filter = delete_file_loader .load_deletes(file_scan_task.deletes(), file_scan_task.schema_ref()) .await @@ -1472,6 +1493,12 @@ mod tests { "Failed to build equality delete predicate: {:?}", result.err() ); + assert_eq!(scan_metrics.data_bytes_read(), 0); + assert!(scan_metrics.delete_bytes_read() > 0); + assert_eq!(scan_metrics.bytes_read(), scan_metrics.delete_bytes_read()); + assert_eq!(scan_metrics.data_files_opened(), 0); + assert_eq!(scan_metrics.delete_files_opened(), 2); + assert_eq!(scan_metrics.rows_emitted(), 0); } #[tokio::test] @@ -1507,7 +1534,12 @@ mod tests { let basic_delete_file_loader = BasicDeleteFileLoader::new(file_io.clone(), ScanMetrics::new()); let record_batch_stream = basic_delete_file_loader - .parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len(), None) + .parquet_to_batch_stream( + &path, + std::fs::metadata(&path).unwrap().len(), + DataContentType::EqualityDeletes, + None, + ) .await .expect("could not get batch stream"); @@ -1627,7 +1659,9 @@ mod tests { .unwrap(), ); - let loader = CachingDeleteFileLoader::new(file_io, 10, Runtime::current()); + let scan_metrics = ScanMetrics::new(); + let loader = CachingDeleteFileLoader::new(file_io, 10, Runtime::current()) + .with_scan_metrics(scan_metrics.clone()); let delete_filter = loader.load_deletes(&[dv], schema).await.unwrap().unwrap(); let delete_vector = delete_filter @@ -1636,11 +1670,17 @@ mod tests { let mut positions: Vec = delete_vector.lock().unwrap().iter().collect(); positions.sort_unstable(); assert_eq!(positions, vec![0, 1, 5]); + assert_eq!(scan_metrics.data_bytes_read(), 0); + assert_eq!(scan_metrics.delete_bytes_read(), content_size); + assert_eq!(scan_metrics.bytes_read(), content_size); + assert_eq!(scan_metrics.data_files_opened(), 0); + assert_eq!(scan_metrics.delete_files_opened(), 1); + assert_eq!(scan_metrics.rows_emitted(), 0); } #[tokio::test] async fn test_load_deletes_decrypts_deletion_vector() { - use crate::encryption::{EncryptedOutputFile, StandardKeyMetadata}; + use crate::encryption::{AesGcmCipher, EncryptedOutputFile, StandardKeyMetadata}; let tmp_dir = TempDir::new().unwrap(); let table_location = tmp_dir.path().to_str().unwrap().to_string(); @@ -1651,29 +1691,6 @@ mod tests { .with_aad_prefix(b"test-aad-prefix!"); let blob = encode_dv_blob([2u64, 4]); - let plaintext_size = blob.len() as u64; - let dv_path = format!("{table_location}/deletes.puffin"); - let output = EncryptedOutputFile::new(file_io.new_output(&dv_path).unwrap(), key_metadata); - let file_metadata = output.write(Bytes::from(blob)).await.unwrap(); - let encoded_key_metadata = output - .key_metadata_with_saved_file_metadata(&file_metadata) - .encode() - .unwrap(); - - // content_offset / content_size_in_bytes are in the plaintext coordinate space, distinct - // from the ciphertext's on-disk size (header, nonce, and tag overhead). - let file_size = std::fs::metadata(&dv_path).unwrap().len(); - let data_file_path = format!("{table_location}/data-1.parquet"); - let dv = dv_task( - dv_path.clone(), - file_size, - data_file_path.clone(), - 0, - plaintext_size, - 2, - Some(encoded_key_metadata), - ); - let schema = Arc::new( Schema::builder() .with_fields(vec![ @@ -1683,15 +1700,75 @@ mod tests { .unwrap(), ); - let loader = CachingDeleteFileLoader::new(file_io, 10, Runtime::current()); - let delete_filter = loader.load_deletes(&[dv], schema).await.unwrap().unwrap(); + let block_size = 1024 * 1024; + let block_overhead = AesGcmCipher::NONCE_LEN + AesGcmCipher::TAG_LEN; + // AGS1 reads fetch complete encrypted blocks, excluding the 8-byte stream header. + for (offset, plaintext_size, storage_bytes) in [ + (0, blob.len(), blob.len() + block_overhead), + (12, blob.len() + 28, blob.len() + 28 + block_overhead), + ( + block_size - 4, + 2 * block_size + 64, + 2 * (block_size + block_overhead), + ), + ] { + let mut plaintext = vec![0; plaintext_size]; + plaintext[offset..offset + blob.len()].copy_from_slice(&blob); + let dv_path = format!("{table_location}/deletes.puffin"); + let output = EncryptedOutputFile::new( + file_io.new_output(&dv_path).unwrap(), + key_metadata.clone(), + ); + let file_metadata = output.write(Bytes::from(plaintext)).await.unwrap(); - let delete_vector = delete_filter - .get_delete_vector_for_path(&data_file_path) - .expect("a delete vector should be indexed for the referenced data file"); - let mut positions: Vec = delete_vector.lock().unwrap().iter().collect(); - positions.sort_unstable(); - assert_eq!(positions, vec![2, 4]); + for wrong_aad in [false, true] { + let mut metadata = output.key_metadata_with_saved_file_metadata(&file_metadata); + if wrong_aad { + metadata = metadata.with_aad_prefix(b"wrong-aad"); + } + let data_file_path = format!("{table_location}/data-1.parquet"); + let dv = dv_task( + dv_path.clone(), + file_metadata.size, + data_file_path.clone(), + offset as u64, + blob.len() as u64, + 2, + Some(metadata.encode().unwrap()), + ); + let scan_metrics = ScanMetrics::new(); + let loader = CachingDeleteFileLoader::new(file_io.clone(), 10, Runtime::current()) + .with_scan_metrics(scan_metrics.clone()); + let result = loader.load_deletes(&[dv], schema.clone()).await.unwrap(); + + if wrong_aad { + let err = result.unwrap_err(); + assert_eq!(err.kind(), ErrorKind::Unexpected); + assert!(err.message().contains("AES-GCM decryption failed")); + } else { + let delete_filter = result.unwrap(); + let delete_vector = delete_filter + .get_delete_vector_for_path(&data_file_path) + .expect("a delete vector should be indexed for the referenced data file"); + let mut positions: Vec = delete_vector.lock().unwrap().iter().collect(); + positions.sort_unstable(); + assert_eq!(positions, vec![2, 4]); + } + + assert_eq!( + ( + scan_metrics.data_bytes_read(), + scan_metrics.delete_bytes_read(), + scan_metrics.bytes_read(), + scan_metrics.data_files_opened(), + scan_metrics.delete_files_opened(), + scan_metrics.rows_emitted(), + ), + (0, storage_bytes as u64, storage_bytes as u64, 0, 1, 0), + "offset={offset}, wrong_aad={wrong_aad}", + ); + } + } } #[tokio::test] diff --git a/crates/iceberg/src/arrow/delete_file_loader.rs b/crates/iceberg/src/arrow/delete_file_loader.rs index ec61710a42..26729a0943 100644 --- a/crates/iceberg/src/arrow/delete_file_loader.rs +++ b/crates/iceberg/src/arrow/delete_file_loader.rs @@ -26,7 +26,7 @@ use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; use crate::arrow::scan_metrics::ScanMetrics; use crate::io::FileIO; use crate::scan::{ArrowRecordBatchStream, FileScanTaskDeleteFile}; -use crate::spec::{Schema, SchemaRef}; +use crate::spec::{DataContentType, Schema, SchemaRef}; use crate::{Error, ErrorKind, Result}; /// Delete File Loader @@ -62,11 +62,16 @@ impl BasicDeleteFileLoader { &self.file_io } + pub(crate) fn scan_metrics(&self) -> &ScanMetrics { + &self.scan_metrics + } + /// Loads a RecordBatchStream for a given datafile. pub(crate) async fn parquet_to_batch_stream( &self, data_file_path: &str, file_size_in_bytes: u64, + file_type: DataContentType, key_metadata: Option<&[u8]>, ) -> Result { /* @@ -80,11 +85,14 @@ impl BasicDeleteFileLoader { &self.file_io, file_size_in_bytes, parquet_read_options, - self.scan_metrics.bytes_read_counter(), + self.scan_metrics.clone(), + file_type, key_metadata, ) .await?; + self.scan_metrics().add_delete_files_opened(1); + let record_batch_stream = ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_file_reader, arrow_metadata) .build()? @@ -126,6 +134,7 @@ impl DeleteFileLoader for BasicDeleteFileLoader { .parquet_to_batch_stream( task.file_path(), task.file_size_in_bytes(), + task.file_type(), task.key_metadata(), ) .await?; @@ -156,7 +165,7 @@ mod tests { let file_io = FileIO::new_with_fs(); let scan_metrics = ScanMetrics::new(); - let delete_file_loader = BasicDeleteFileLoader::new(file_io.clone(), scan_metrics); + let delete_file_loader = BasicDeleteFileLoader::new(file_io.clone(), scan_metrics.clone()); let file_scan_tasks = setup(table_location); @@ -171,6 +180,12 @@ mod tests { let result = result.try_collect::>().await.unwrap(); assert_eq!(result.len(), 1); + assert_eq!(scan_metrics.data_bytes_read(), 0); + assert!(scan_metrics.delete_bytes_read() > 0); + assert_eq!(scan_metrics.bytes_read(), scan_metrics.delete_bytes_read()); + assert_eq!(scan_metrics.data_files_opened(), 0); + assert_eq!(scan_metrics.delete_files_opened(), 1); + assert_eq!(scan_metrics.rows_emitted(), 0); } #[tokio::test] diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 1ddee493c4..ba07403007 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -23,7 +23,6 @@ use std::borrow::Cow; use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::AtomicU64; use arrow_schema::{DataType, Field}; use futures::{StreamExt, TryStreamExt}; @@ -58,7 +57,7 @@ use crate::metadata_columns::{ }; use crate::runtime::Runtime; use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream}; -use crate::spec::{Datum, PartitionSpec, Struct}; +use crate::spec::{DataContentType, Datum, PartitionSpec, Struct}; use crate::{Error, ErrorKind}; impl ArrowReader { @@ -167,7 +166,7 @@ impl FileScanTaskReader { let path = task.data_file_path().to_string(); let file_io = self.file_io.clone(); let file_size_in_bytes = task.file_size_in_bytes(); - let bytes_read = Arc::clone(self.scan_metrics.bytes_read_counter()); + let scan_metrics = self.scan_metrics.clone(); let key_metadata = task.key_metadata().map(|km| km.to_vec()); let opening = self .runtime @@ -178,7 +177,8 @@ impl FileScanTaskReader { &file_io, file_size_in_bytes, parquet_read_options, - &bytes_read, + scan_metrics, + DataContentType::Data, key_metadata.as_deref(), ) .await @@ -186,6 +186,8 @@ impl FileScanTaskReader { .abort_on_drop(); let (parquet_file_reader, arrow_metadata) = opening.await??; + self.scan_metrics.add_data_files_opened(1); + // Check if Parquet file has embedded field IDs // Corresponds to Java's ParquetSchemaUtil.hasIds() // Reference: parquet/src/main/java/org/apache/iceberg/parquet/ParquetSchemaUtil.java:118 @@ -659,13 +661,19 @@ impl FileScanTaskReader { // batches (using the reader-produced `_pos` position) before the transformer, which // then passes it through as a virtual field. let first_row_id = task.first_row_id(); + let scan_metrics = self.scan_metrics; let record_batch_stream = record_batch_stream_builder.build()?.map(move |batch| { let mut batch = batch.map_err(|err| -> Error { err.into() })?; if project_row_id { batch = synthesize_row_id_column(batch, first_row_id)?; } // Process the record batch (type promotion, column reordering, virtual fields, etc.) - record_batch_transformer.process_record_batch(batch) + let record_batch = record_batch_transformer.process_record_batch(batch); + if let Ok(records) = &record_batch { + scan_metrics.add_rows_emitted(records.num_rows() as u64); + } + + record_batch }); Ok(Box::pin(record_batch_stream) as ArrowRecordBatchStream) @@ -811,18 +819,20 @@ impl FileScanTaskReader { impl ArrowReader { /// Opens a Parquet file and loads its metadata, wrapping the reader with - /// [`CountingFileRead`] so all I/O is accumulated into `bytes_read`. + /// [`CountingFileRead`] so all I/O is accumulated into the scan metrics + /// for the specified file type. pub(crate) async fn open_parquet_file( data_file_path: &str, file_io: &FileIO, file_size_in_bytes: u64, parquet_read_options: ParquetReadOptions, - bytes_read: &Arc, + scan_metrics: ScanMetrics, + file_type: DataContentType, key_metadata: Option<&[u8]>, ) -> Result<(ArrowFileReader, ArrowReaderMetadata)> { let parquet_file = file_io.new_input(data_file_path)?; let counting_reader = - CountingFileRead::new(parquet_file.reader().await?, Arc::clone(bytes_read)); + CountingFileRead::new(parquet_file.reader().await?, scan_metrics, file_type); Self::build_parquet_reader( Box::new(counting_reader), file_size_in_bytes, @@ -910,7 +920,7 @@ mod tests { RESERVED_COL_NAME_POS, RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE, RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_ROW_ID, }; - use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream}; + use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream, ScanMetrics}; use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, SchemaRef, Type}; // INT96 encoding: [nanos_low_u32, nanos_high_u32, julian_day_u32] @@ -1800,7 +1810,8 @@ mod tests { vec![RESERVED_FIELD_ID_ROW_ID], Some(100), ); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); assert_row_id_column(&batches, &[Some(100), Some(101), Some(102)]); @@ -1812,7 +1823,8 @@ mod tests { vec![2, RESERVED_FIELD_ID_ROW_ID], Some(100), ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -1836,7 +1848,8 @@ mod tests { vec![RESERVED_FIELD_ID_ROW_ID], None, ); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); assert_row_id_column(&batches, &[None, None, None]); @@ -1847,7 +1860,8 @@ mod tests { vec![2, RESERVED_FIELD_ID_ROW_ID], None, ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -1887,7 +1901,8 @@ mod tests { write_parquet_with_wide_column(dir, "lusn_only.parquet", vec![], vec![]), vec![RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER], ); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); assert_last_updated_seq_column(&batches, &[Some(7), Some(7), Some(7)]); @@ -1896,7 +1911,8 @@ mod tests { write_parquet_with_wide_column(dir, "lusn_only_ref.parquet", vec![], vec![]), vec![2, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER], ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -1950,7 +1966,8 @@ mod tests { write_parquet_with_wide_column(dir, "partition_only.parquet", vec![], vec![]), vec![RESERVED_FIELD_ID_PARTITION], ); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); let partition_col = batches[0] @@ -1971,7 +1988,8 @@ mod tests { write_parquet_with_wide_column(dir, "partition_only_ref.parquet", vec![], vec![]), vec![2, RESERVED_FIELD_ID_PARTITION], ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -3372,14 +3390,14 @@ mod tests { .unwrap() } - /// Runs a single-task scan and returns the batches plus the bytes read from storage. - async fn scan_task(task: FileScanTask) -> (Vec, u64) { + /// Runs a single-task scan and returns the batches plus its metrics. + async fn scan_task(task: FileScanTask) -> (Vec, ScanMetrics) { let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()).build(); let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; let scan = reader.read(tasks).unwrap(); let metrics = scan.metrics().clone(); let batches = scan.stream().try_collect().await.unwrap(); - (batches, metrics.bytes_read()) + (batches, metrics) } #[tokio::test] @@ -3392,7 +3410,8 @@ mod tests { id_and_wide_schema(), vec![RESERVED_FIELD_ID_POS], ); - let (batches, pos_only_bytes) = scan_task(pos_only).await; + let (batches, pos_only_metrics) = scan_task(pos_only).await; + let pos_only_bytes = pos_only_metrics.bytes_read(); // Only `_pos` is materialized -- no data columns. assert_eq!(batches[0].num_columns(), 1); @@ -3409,7 +3428,8 @@ mod tests { id_and_wide_schema(), vec![2, RESERVED_FIELD_ID_POS], ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( pos_only_bytes < with_data_bytes, @@ -3428,7 +3448,8 @@ mod tests { id_and_wide_schema(), vec![], ); - let (batches, empty_bytes) = scan_task(empty).await; + let (batches, empty_metrics) = scan_task(empty).await; + let empty_bytes = empty_metrics.bytes_read(); // No data columns come out, but the fixture's 3 rows are still counted. assert!(batches.iter().all(|b| b.num_columns() == 0)); @@ -3439,7 +3460,8 @@ mod tests { id_and_wide_schema(), vec![2], ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( empty_bytes < with_data_bytes, @@ -3477,9 +3499,18 @@ mod tests { .build() .unwrap(); - let (batches, _) = scan_task(task).await; + let (batches, metrics) = scan_task(task).await; assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + assert!(metrics.data_bytes_read() > 0); + assert!(metrics.delete_bytes_read() > 0); + assert_eq!( + metrics.bytes_read(), + metrics.data_bytes_read() + metrics.delete_bytes_read() + ); + assert_eq!(metrics.data_files_opened(), 1); + assert_eq!(metrics.delete_files_opened(), 1); + assert_eq!(metrics.rows_emitted(), 2); } #[tokio::test] @@ -3514,6 +3545,31 @@ mod tests { assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); } + #[tokio::test] + async fn test_scan_metrics_for_data_file() { + let tmp_dir = TempDir::new().unwrap(); + let dir = tmp_dir.path().to_str().unwrap(); + + let task = metadata_projection_task( + write_parquet_with_wide_column(dir, "scan_metrics.parquet", vec![], vec![]), + id_and_wide_schema(), + vec![1, 2], + ); + + let (batches, metrics) = scan_task(task).await; + let rows = batches + .iter() + .map(|batch| batch.num_rows() as u64) + .sum::(); + + assert!(metrics.data_bytes_read() > 0); + assert_eq!(metrics.delete_bytes_read(), 0); + assert_eq!(metrics.bytes_read(), metrics.data_bytes_read()); + assert_eq!(metrics.data_files_opened(), 1); + assert_eq!(metrics.delete_files_opened(), 0); + assert_eq!(metrics.rows_emitted(), rows); + } + #[tokio::test] async fn test_pos_only_projection_keeps_absolute_pos_under_predicate() { use crate::expr::{Bind, Reference}; @@ -3651,7 +3707,8 @@ mod tests { RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, ]); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); // `_pos` and the coalesced sequence column materialize; the wide column does not. let pos_col = batches[0] @@ -3677,7 +3734,8 @@ mod tests { RESERVED_FIELD_ID_POS, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, ]); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -3725,7 +3783,8 @@ mod tests { let meta_only = seq_task(write("seq_only.parquet"), vec![ RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, ]); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); let seq_col = batches[0] .column_by_name(RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER) @@ -3743,7 +3802,8 @@ mod tests { 2, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, ]); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -3790,7 +3850,8 @@ mod tests { let meta_only = seq_task("seq_only_null_first.parquet", vec![ RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, ]); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); @@ -3806,7 +3867,8 @@ mod tests { 2, RESERVED_FIELD_ID_LAST_UPDATED_SEQUENCE_NUMBER, ]); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -3829,7 +3891,8 @@ mod tests { let meta_only = metadata_projection_task(file_path.clone(), id_and_wide_schema(), vec![ RESERVED_FIELD_ID_FILE, ]); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); @@ -3846,7 +3909,8 @@ mod tests { id_and_wide_schema(), vec![2, RESERVED_FIELD_ID_FILE], ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, @@ -3895,7 +3959,8 @@ mod tests { write_parquet_with_wide_column(dir, "spec_id_only.parquet", vec![], vec![]), vec![RESERVED_FIELD_ID_SPEC_ID], ); - let (batches, meta_only_bytes) = scan_task(meta_only).await; + let (batches, meta_only_metrics) = scan_task(meta_only).await; + let meta_only_bytes = meta_only_metrics.bytes_read(); assert_eq!(batches[0].num_columns(), 1); let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); @@ -3911,7 +3976,8 @@ mod tests { write_parquet_with_wide_column(dir, "spec_id_only_ref.parquet", vec![], vec![]), vec![2, RESERVED_FIELD_ID_SPEC_ID], ); - let (_, with_data_bytes) = scan_task(with_data).await; + let (_, with_data_metrics) = scan_task(with_data).await; + let with_data_bytes = with_data_metrics.bytes_read(); assert!( meta_only_bytes < with_data_bytes, diff --git a/crates/iceberg/src/arrow/scan_metrics.rs b/crates/iceberg/src/arrow/scan_metrics.rs index 642190c57d..b4b35cc230 100644 --- a/crates/iceberg/src/arrow/scan_metrics.rs +++ b/crates/iceberg/src/arrow/scan_metrics.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Scan metrics and I/O counting for Parquet data file reads. +//! Scan metrics and I/O counting for data and delete file reads. use std::ops::Range; use std::sync::Arc; @@ -26,16 +26,22 @@ use bytes::Bytes; use crate::error::Result; use crate::io::FileRead; use crate::scan::ArrowRecordBatchStream; +use crate::spec::DataContentType; -/// Wraps a [`FileRead`] to count bytes read via a shared atomic counter. +/// Wraps a [`FileRead`] to count bytes read in the shared scan metrics. pub(crate) struct CountingFileRead { inner: F, - bytes_read: Arc, + scan_metrics: ScanMetrics, + file_type: DataContentType, } impl CountingFileRead { - pub(crate) fn new(inner: F, bytes_read: Arc) -> Self { - Self { inner, bytes_read } + pub(crate) fn new(inner: F, scan_metrics: ScanMetrics, file_type: DataContentType) -> Self { + Self { + inner, + scan_metrics, + file_type, + } } } @@ -43,8 +49,13 @@ impl CountingFileRead { impl FileRead for CountingFileRead { async fn read(&self, range: Range) -> Result { debug_assert!(range.end >= range.start); - self.bytes_read - .fetch_add(range.end - range.start, Ordering::Relaxed); + let bytes_read = range.end - range.start; + match self.file_type { + DataContentType::Data => self.scan_metrics.add_data_bytes_read(bytes_read), + DataContentType::PositionDeletes | DataContentType::EqualityDeletes => { + self.scan_metrics.add_delete_bytes_read(bytes_read) + } + } self.inner.read(range).await } } @@ -52,23 +63,83 @@ impl FileRead for CountingFileRead { /// Metrics collected during an Iceberg scan. #[derive(Clone, Debug)] pub struct ScanMetrics { - bytes_read: Arc, + inner: Arc, +} + +#[derive(Debug, Default)] +struct ScanMetricsInner { + data_bytes_read: AtomicU64, + delete_bytes_read: AtomicU64, + data_files_opened: AtomicU64, + delete_files_opened: AtomicU64, + rows_emitted: AtomicU64, } impl ScanMetrics { pub(crate) fn new() -> Self { Self { - bytes_read: Arc::new(AtomicU64::new(0)), + inner: Arc::new(ScanMetricsInner::default()), } } - pub(crate) fn bytes_read_counter(&self) -> &Arc { - &self.bytes_read + pub(crate) fn add_data_bytes_read(&self, bytes: u64) { + self.inner + .data_bytes_read + .fetch_add(bytes, Ordering::Relaxed); + } + + pub(crate) fn add_delete_bytes_read(&self, bytes: u64) { + self.inner + .delete_bytes_read + .fetch_add(bytes, Ordering::Relaxed); + } + + pub(crate) fn add_data_files_opened(&self, files: u64) { + self.inner + .data_files_opened + .fetch_add(files, Ordering::Relaxed); + } + + pub(crate) fn add_delete_files_opened(&self, files: u64) { + self.inner + .delete_files_opened + .fetch_add(files, Ordering::Relaxed); + } + + pub(crate) fn add_rows_emitted(&self, rows: u64) { + self.inner.rows_emitted.fetch_add(rows, Ordering::Relaxed); } /// Total bytes read from storage during this scan, including data files and delete files. pub fn bytes_read(&self) -> u64 { - self.bytes_read.load(Ordering::Relaxed) + self.data_bytes_read() + self.delete_bytes_read() + } + + /// Total bytes read from data files in storage during this scan. + pub fn data_bytes_read(&self) -> u64 { + self.inner.data_bytes_read.load(Ordering::Relaxed) + } + + /// Total bytes read from delete files in storage during this scan. + pub fn delete_bytes_read(&self) -> u64 { + self.inner.delete_bytes_read.load(Ordering::Relaxed) + } + + /// Total number of data files opened during this scan + pub fn data_files_opened(&self) -> u64 { + self.inner.data_files_opened.load(Ordering::Relaxed) + } + + /// Total number of delete files opened during this scan + /// Delete files can be reused/shared across multiple data file scans + /// This does *not* include delete files that were reused from cache + pub fn delete_files_opened(&self) -> u64 { + self.inner.delete_files_opened.load(Ordering::Relaxed) + } + + /// Total number of rows emitted from this scan + pub fn rows_emitted(&self) -> u64 { + self.inner.rows_emitted.load(Ordering::Relaxed) } } @@ -94,3 +165,35 @@ impl ScanResult { &self.metrics } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cloned_metrics_share_concurrent_updates() { + let metrics = ScanMetrics::new(); + + std::thread::scope(|scope| { + for _ in 0..4 { + let metrics = metrics.clone(); + scope.spawn(move || { + for _ in 0..100 { + metrics.add_data_bytes_read(10); + metrics.add_delete_bytes_read(20); + metrics.add_data_files_opened(1); + metrics.add_delete_files_opened(2); + metrics.add_rows_emitted(3); + } + }); + } + }); + + assert_eq!(metrics.data_bytes_read(), 4000); + assert_eq!(metrics.delete_bytes_read(), 8000); + assert_eq!(metrics.bytes_read(), 12000); + assert_eq!(metrics.data_files_opened(), 400); + assert_eq!(metrics.delete_files_opened(), 800); + assert_eq!(metrics.rows_emitted(), 1200); + } +} diff --git a/crates/iceberg/src/encryption/io.rs b/crates/iceberg/src/encryption/io.rs index 17fc1a1f23..fe0560857d 100644 --- a/crates/iceberg/src/encryption/io.rs +++ b/crates/iceberg/src/encryption/io.rs @@ -74,8 +74,16 @@ impl EncryptedInputFile { /// Creates a reader that transparently decrypts on each read. pub async fn reader(&self) -> Result> { - let encrypted_length = self.encrypted_length()?; let raw_reader = self.inner.reader().await?; + self.decrypt_reader(raw_reader) + } + + /// Wraps a transparently decrypting reader over the given `raw_reader` + pub(crate) fn decrypt_reader( + &self, + raw_reader: Box, + ) -> Result> { + let encrypted_length = self.encrypted_length()?; let cipher = build_cipher(&self.key_metadata)?; let aad_prefix: Box<[u8]> = self.key_metadata.aad_prefix().unwrap_or_default().into(); let decrypting = AesGcmFileRead::new(raw_reader, cipher, aad_prefix, encrypted_length)?; @@ -273,25 +281,30 @@ mod tests { #[tokio::test] async fn test_invalid_file_length_is_rejected() { let fileio = FileIO::new_with_memory(); + let path = "memory:///test/invalid_length.bin"; + let output = EncryptedOutputFile::new(fileio.new_output(path).unwrap(), key_metadata()); + output.write(Bytes::from_static(b"data")).await.unwrap(); + for length in [ 0, u64::from(GCM_STREAM_HEADER_LENGTH), u64::from(MIN_STREAM_LENGTH - 1), ] { let input = EncryptedInputFile::new( - fileio - .new_input("memory:///test/invalid_length.bin") - .unwrap(), + fileio.new_input(path).unwrap(), key_metadata().with_file_length(length), ); - assert_eq!( - input.metadata().err().unwrap().kind(), - ErrorKind::DataInvalid - ); - assert_eq!( - input.reader().await.err().unwrap().kind(), - ErrorKind::DataInvalid - ); + let expected_message = + format!("Invalid encrypted file length: {length} is less than {MIN_STREAM_LENGTH}"); + for err in [ + input.metadata().err().unwrap(), + input.reader().await.err().unwrap(), + ] { + assert_eq!( + (err.kind(), err.message()), + (ErrorKind::DataInvalid, expected_message.as_str()) + ); + } } }