From c2df6575d160aa0ca5f419f27ed0724409d231be Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 16:06:29 -0700 Subject: [PATCH 01/12] feat: expose native Parquet scan I/O metrics --- native/core/src/execution/metrics/utils.rs | 5 +- .../eager_page_index_reader_factory.rs | 413 +++++++++++++++++- native/core/src/parquet/parquet_exec.rs | 276 +++++++++++- .../spark/sql/comet/CometMetricNode.scala | 40 +- .../apache/comet/exec/CometExecSuite.scala | 32 ++ 5 files changed, 744 insertions(+), 22 deletions(-) diff --git a/native/core/src/execution/metrics/utils.rs b/native/core/src/execution/metrics/utils.rs index 50a35393496..4104aaa2e02 100644 --- a/native/core/src/execution/metrics/utils.rs +++ b/native/core/src/execution/metrics/utils.rs @@ -65,9 +65,8 @@ pub(crate) fn to_native_metric_node( let children = spark_plan.children(); let mut native_metric_node = NativeMetricNode { - // Most operator metric maps are well under 20 entries (e.g. hash-join: 9, - // native-scan: ~20). Pre-sizing to 16 avoids the default-capacity rehash. - metrics: HashMap::with_capacity(16), + // Native scans expose the largest metric map, including scan I/O breakdowns. + metrics: HashMap::with_capacity(32), children: Vec::with_capacity(children.len()), }; diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 278814c4bf9..ec686b9b5a1 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -46,6 +46,7 @@ //! Filed upstream as apache/datafusion#23978. Revert this once the opener merges its deferred //! page-index load back into `FileMetadataCache` instead of bypassing it. +use async_trait::async_trait; use bytes::Bytes; use datafusion::common::Result as DFResult; use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; @@ -53,29 +54,235 @@ use datafusion::datasource::physical_plan::parquet::{ ParquetFileMetrics, ParquetFileReaderFactory, }; use datafusion::execution::cache::cache_manager::FileMetadataCache; -use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion::physical_plan::metrics::{ + Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricType, +}; use datafusion_datasource::PartitionedFile; use futures::future::BoxFuture; -use futures::FutureExt; -use object_store::ObjectStore; +use futures::{FutureExt, StreamExt, TryStreamExt}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + Result as ObjectStoreResult, +}; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; -use std::fmt::Debug; +use std::fmt::{Debug, Display, Formatter}; use std::ops::Range; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; +/// Whether the reader's storage API represents a local file or a non-local object store. +/// +/// The source metrics intentionally describe the API boundary, not network wire bytes. A remote +/// backend can coalesce ranges, retry requests, or serve bytes from a cache below ObjectStore, +/// none of which this reader can observe without backend-specific hooks. +#[derive(Debug, Clone, Copy)] +pub(crate) enum ScanIoSource { + ObjectStore, + Local, +} + +/// Scan I/O metrics at the boundaries this reader can observe without guessing. +/// +/// requested is the sum of byte ranges requested by the Parquet reader or metadata loader. +/// returned is the length of successfully returned buffers at the same boundary. Metadata +/// includes footer prefetches, footer decode follow-up reads, and page-index ranges; DataFusion +/// does not identify those subranges separately, so we report them together instead of claiming +/// unsupported footer-versus-page-index precision. +/// +/// scan_io_object_store metrics count non-file storage API bytes and scan_io_local metrics count +/// file bytes. Parsed metadata cache entries do not have a meaningful raw-byte size, so cache +/// metrics report successful cache-eligible loads that did or did not require storage I/O rather +/// than inventing cache-byte counts. +#[derive(Debug, Clone)] +struct ScanIoMetrics { + bytes_requested: Count, + bytes_returned: Count, + data_bytes_requested: Count, + data_bytes_returned: Count, + metadata_bytes_requested: Count, + metadata_bytes_returned: Count, + object_store_bytes_requested: Count, + object_store_bytes_returned: Count, + local_bytes_requested: Count, + local_bytes_returned: Count, + metadata_cache_hits: Count, + metadata_cache_misses: Count, + source: ScanIoSource, +} + +impl ScanIoMetrics { + fn new( + partition: usize, + filename: &str, + metrics: &ExecutionPlanMetricsSet, + source: ScanIoSource, + ) -> Self { + Self { + bytes_requested: byte_counter(metrics, partition, filename, "scan_io_bytes_requested"), + bytes_returned: byte_counter(metrics, partition, filename, "scan_io_bytes_returned"), + data_bytes_requested: byte_counter( + metrics, + partition, + filename, + "scan_io_data_bytes_requested", + ), + data_bytes_returned: byte_counter( + metrics, + partition, + filename, + "scan_io_data_bytes_returned", + ), + metadata_bytes_requested: byte_counter( + metrics, + partition, + filename, + "scan_io_metadata_bytes_requested", + ), + metadata_bytes_returned: byte_counter( + metrics, + partition, + filename, + "scan_io_metadata_bytes_returned", + ), + object_store_bytes_requested: byte_counter( + metrics, + partition, + filename, + "scan_io_object_store_bytes_requested", + ), + object_store_bytes_returned: byte_counter( + metrics, + partition, + filename, + "scan_io_object_store_bytes_returned", + ), + local_bytes_requested: byte_counter( + metrics, + partition, + filename, + "scan_io_local_bytes_requested", + ), + local_bytes_returned: byte_counter( + metrics, + partition, + filename, + "scan_io_local_bytes_returned", + ), + metadata_cache_hits: count_counter( + metrics, + partition, + filename, + "scan_io_metadata_cache_hits", + ), + metadata_cache_misses: count_counter( + metrics, + partition, + filename, + "scan_io_metadata_cache_misses", + ), + source, + } + } + + fn add_data_requested(&self, bytes: usize) { + self.data_bytes_requested.add(bytes); + self.add_requested(bytes); + } + + fn add_data_returned(&self, bytes: usize) { + self.data_bytes_returned.add(bytes); + self.add_returned(bytes); + } + + fn add_metadata_requested(&self, bytes: usize) { + self.metadata_bytes_requested.add(bytes); + self.add_requested(bytes); + } + + fn add_metadata_returned(&self, bytes: usize) { + self.metadata_bytes_returned.add(bytes); + self.add_returned(bytes); + } + + fn record_metadata_cache_result(&self, storage_reads: usize) { + if storage_reads == 0 { + self.metadata_cache_hits.add(1); + } else { + self.metadata_cache_misses.add(1); + } + } + + fn add_requested(&self, bytes: usize) { + self.bytes_requested.add(bytes); + match self.source { + ScanIoSource::ObjectStore => self.object_store_bytes_requested.add(bytes), + ScanIoSource::Local => self.local_bytes_requested.add(bytes), + } + } + + fn add_returned(&self, bytes: usize) { + self.bytes_returned.add(bytes); + match self.source { + ScanIoSource::ObjectStore => self.object_store_bytes_returned.add(bytes), + ScanIoSource::Local => self.local_bytes_returned.add(bytes), + } + } +} + +fn byte_counter( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + filename: &str, + name: &'static str, +) -> Count { + MetricBuilder::new(metrics) + .with_new_label("filename", filename.to_string()) + .with_type(MetricType::Summary) + .with_category(MetricCategory::Bytes) + .counter(name, partition) +} + +fn count_counter( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + filename: &str, + name: &'static str, +) -> Count { + MetricBuilder::new(metrics) + .with_new_label("filename", filename.to_string()) + .with_type(MetricType::Summary) + .counter(name, partition) +} + +fn range_bytes(range: &Range) -> usize { + (range.end - range.start) as usize +} + +fn ranges_bytes(ranges: &[Range]) -> usize { + ranges.iter().map(range_bytes).sum() +} + #[derive(Debug)] pub struct EagerPageIndexReaderFactory { store: Arc, metadata_cache: Arc, + source: ScanIoSource, } impl EagerPageIndexReaderFactory { - pub fn new(store: Arc, metadata_cache: Arc) -> Self { + pub(crate) fn new( + store: Arc, + metadata_cache: Arc, + source: ScanIoSource, + ) -> Self { Self { store, metadata_cache, + source, } } } @@ -93,6 +300,12 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { partitioned_file.object_meta.location.as_ref(), metrics, ); + let scan_io_metrics = ScanIoMetrics::new( + partition_index, + partitioned_file.object_meta.location.as_ref(), + metrics, + self.source, + ); let mut inner = ParquetObjectReader::new( Arc::clone(&self.store), partitioned_file.object_meta.location.clone(), @@ -104,6 +317,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { Ok(Box::new(EagerPageIndexReader { file_metrics, + scan_io_metrics, store: Arc::clone(&self.store), inner, partitioned_file, @@ -115,6 +329,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { struct EagerPageIndexReader { file_metrics: ParquetFileMetrics, + scan_io_metrics: ScanIoMetrics, store: Arc, inner: ParquetObjectReader, partitioned_file: PartitionedFile, @@ -124,9 +339,17 @@ struct EagerPageIndexReader { impl AsyncFileReader for EagerPageIndexReader { fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { - let bytes_scanned = range.end - range.start; - self.file_metrics.bytes_scanned.add(bytes_scanned as usize); - self.inner.get_bytes(range) + let requested = range_bytes(&range); + self.file_metrics.bytes_scanned.add(requested); + self.scan_io_metrics.add_data_requested(requested); + let scan_io_metrics = self.scan_io_metrics.clone(); + let future = self.inner.get_bytes(range); + async move { + let bytes = future.await?; + scan_io_metrics.add_data_returned(bytes.len()); + Ok(bytes) + } + .boxed() } fn get_byte_ranges( @@ -136,9 +359,17 @@ impl AsyncFileReader for EagerPageIndexReader { where Self: Send, { - let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); - self.file_metrics.bytes_scanned.add(total as usize); - self.inner.get_byte_ranges(ranges) + let requested = ranges_bytes(&ranges); + self.file_metrics.bytes_scanned.add(requested); + self.scan_io_metrics.add_data_requested(requested); + let scan_io_metrics = self.scan_io_metrics.clone(); + let future = self.inner.get_byte_ranges(ranges); + async move { + let bytes = future.await?; + scan_io_metrics.add_data_returned(bytes.iter().map(Bytes::len).sum()); + Ok(bytes) + } + .boxed() } fn get_metadata<'a>( @@ -151,17 +382,25 @@ impl AsyncFileReader for EagerPageIndexReader { let metadata_cache = Arc::clone(&self.metadata_cache); let store = Arc::clone(&self.store); let metadata_size_hint = self.metadata_size_hint; + let scan_io_metrics = self.scan_io_metrics.clone(); async move { let file_decryption_properties = options .and_then(|o| o.file_decryption_properties()) .map(Arc::clone); + let cache_enabled = file_decryption_properties.is_none(); let page_index_policy = if file_decryption_properties.is_none() { Some(PageIndexPolicy::Optional) } else { options.map(|o| o.column_index_policy()) }; + let metadata_storage_reads = Arc::new(AtomicUsize::new(0)); + let metadata_store = MetadataIoObjectStore { + inner: store, + scan_io_metrics: scan_io_metrics.clone(), + storage_reads: Arc::clone(&metadata_storage_reads), + }; - DFParquetMetadata::new(store.as_ref(), &object_meta) + let metadata = DFParquetMetadata::new(&metadata_store, &object_meta) .with_decryption_properties(file_decryption_properties) .with_file_metadata_cache(Some(metadata_cache)) .with_metadata_size_hint(metadata_size_hint) @@ -173,12 +412,160 @@ impl AsyncFileReader for EagerPageIndexReader { "Failed to fetch metadata for file {}: {e}", object_meta.location, )) - }) + }); + + if cache_enabled && metadata.is_ok() { + scan_io_metrics + .record_metadata_cache_result(metadata_storage_reads.load(Ordering::Relaxed)); + } + + metadata } .boxed() } } +/// Counts metadata storage calls while forwarding every operation to the configured store. +/// +/// Data reads are already visible at the AsyncFileReader data methods. Metadata reads bypass +/// those methods inside DFParquetMetadata, so only metadata uses this wrapper. +#[derive(Debug)] +struct MetadataIoObjectStore { + inner: Arc, + scan_io_metrics: ScanIoMetrics, + storage_reads: Arc, +} + +impl MetadataIoObjectStore { + fn record_request(&self, bytes: usize) { + if bytes > 0 { + self.storage_reads.fetch_add(1, Ordering::Relaxed); + self.scan_io_metrics.add_metadata_requested(bytes); + } + } +} + +impl Display for MetadataIoObjectStore { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + write!(formatter, "metadata-io({})", self.inner) + } +} + +#[async_trait] +#[deny(clippy::missing_trait_methods)] +impl ObjectStore for MetadataIoObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + options: PutOptions, + ) -> ObjectStoreResult { + self.inner.put_opts(location, payload, options).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + options: PutMultipartOptions, + ) -> ObjectStoreResult> { + self.inner.put_multipart_opts(location, options).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> ObjectStoreResult { + if options.head { + return self.inner.get_opts(location, options).await; + } + + let requested = match options.range.as_ref() { + Some(GetRange::Bounded(range)) => Some(range_bytes(range)), + Some(GetRange::Suffix(bytes)) => Some(*bytes as usize), + Some(GetRange::Offset(_)) | None => None, + }; + if let Some(requested) = requested { + self.record_request(requested); + } + + let result = self.inner.get_opts(location, options).await?; + if requested.is_none() { + self.record_request(range_bytes(&result.range)); + } + + let meta = result.meta.clone(); + let range = result.range.clone(); + let attributes = result.attributes.clone(); + let scan_io_metrics = self.scan_io_metrics.clone(); + let payload = GetResultPayload::Stream( + result + .into_stream() + .inspect_ok(move |bytes| scan_io_metrics.add_metadata_returned(bytes.len())) + .boxed(), + ); + + Ok(GetResult { + payload, + meta, + range, + attributes, + }) + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> ObjectStoreResult> { + self.record_request(ranges_bytes(ranges)); + let bytes = self.inner.get_ranges(location, ranges).await?; + self.scan_io_metrics + .add_metadata_returned(bytes.iter().map(Bytes::len).sum()); + Ok(bytes) + } + + fn delete_stream( + &self, + locations: futures::stream::BoxStream<'static, ObjectStoreResult>, + ) -> futures::stream::BoxStream<'static, ObjectStoreResult> { + self.inner.delete_stream(locations) + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> futures::stream::BoxStream<'static, ObjectStoreResult> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> futures::stream::BoxStream<'static, ObjectStoreResult> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> ObjectStoreResult { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> ObjectStoreResult<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + options: RenameOptions, + ) -> ObjectStoreResult<()> { + self.inner.rename_opts(from, to, options).await + } +} + impl Drop for EagerPageIndexReader { fn drop(&mut self) { self.file_metrics diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 1308ce97fca..7ab0c1d751b 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -16,7 +16,7 @@ // under the License. use crate::execution::operators::ExecutionError; -use crate::parquet::eager_page_index_reader_factory::EagerPageIndexReaderFactory; +use crate::parquet::eager_page_index_reader_factory::{EagerPageIndexReaderFactory, ScanIoSource}; use crate::parquet::encryption_support::{CometEncryptionConfig, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_support::SparkParquetOptions; use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; @@ -158,14 +158,16 @@ pub(crate) fn init_datasource_exec( // cached with the footer, at the cost of losing the skip's benefit when it would have // applied. Filed upstream as apache/datafusion#23978; revert this once that's fixed. // - // TODO: metadata I/O is invisible in metrics. `fetch_metadata` reads via `ObjectStore::get_ranges`, - // bypassing the `get_bytes` path where `bytes_scanned` is counted. A byte-counting ObjectStore - // wrapper would surface it. let runtime_env = session_ctx.runtime_env(); let store = runtime_env.object_store(&object_store_url)?; let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let scan_io_source = if object_store_url == ObjectStoreUrl::local_filesystem() { + ScanIoSource::Local + } else { + ScanIoSource::ObjectStore + }; parquet_source = parquet_source.with_parquet_file_reader_factory(Arc::new( - EagerPageIndexReaderFactory::new(store, metadata_cache), + EagerPageIndexReaderFactory::new(store, metadata_cache, scan_io_source), )); // Route data filters through `try_pushdown_filters` rather than calling @@ -295,6 +297,10 @@ mod tests { use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use datafusion::datasource::physical_plan::parquet::metadata::CachedParquetMetaData; + use datafusion::datasource::physical_plan::parquet::ParquetFileReaderFactory; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::ExecutionPlan; use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename; use futures::StreamExt; @@ -302,6 +308,101 @@ mod tests { use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::fs::File; + fn write_scan_io_fixture() -> (String, SchemaRef) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let first = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from((0..500).collect::>())), + Arc::new(Int32Array::from((10_000..10_500).collect::>())), + ], + ) + .unwrap(); + let second = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from((1_000..1_500).collect::>())), + Arc::new(Int32Array::from((11_000..11_500).collect::>())), + ], + ) + .unwrap(); + + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_data_page_row_count_limit(100) + .set_max_row_group_row_count(Some(500)) + .build(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), Some(props)).unwrap(); + writer.write(&first).unwrap(); + writer.write(&second).unwrap(); + writer.close().unwrap(); + + (filename, schema) + } + + fn init_test_scan( + required_schema: SchemaRef, + data_schema: SchemaRef, + partitioned_file: PartitionedFile, + projection: Option>, + filters: Option>>, + session_ctx: &Arc, + ) -> Arc { + init_datasource_exec( + required_schema, + Some(data_schema), + None, + ObjectStoreUrl::local_filesystem(), + vec![vec![partitioned_file]], + projection, + filters, + None, + "UTC", + true, + false, + false, + false, + session_ctx, + false, + false, + false, + ) + .unwrap() + } + + async fn drain_scan(scan: &Arc, session_ctx: &Arc) { + let mut stream = scan.execute(0, session_ctx.task_ctx()).unwrap(); + while let Some(batch) = stream.next().await { + batch.unwrap(); + } + } + + fn scan_metric(scan: &Arc, name: &str) -> usize { + scan.metrics() + .unwrap() + .sum_by_name(name) + .unwrap_or_else(|| panic!("missing metric {name}")) + .as_usize() + } + + fn reader_metric(metrics: &ExecutionPlanMetricsSet, name: &str) -> usize { + metrics + .clone_inner() + .sum_by_name(name) + .unwrap_or_else(|| panic!("missing metric {name}")) + .as_usize() + } + // Regression test for #4990: a fresh `TableParquetOptions::new()` ignored session-level // `datafusion.execution.parquet.*` settings entirely, so `spark.comet.datafusion. // execution.parquet.*` (behind `respectDataFusionConfigs`) and `spark.comet.parquet. @@ -432,4 +533,169 @@ mod tests { "cached metadata must include the page index" ); } + + #[tokio::test] + async fn reports_cold_and_warm_metadata_io_without_data_reads() { + let (filename, _schema) = write_scan_io_fixture(); + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let runtime_env = session_ctx.runtime_env(); + let store = runtime_env + .object_store(ObjectStoreUrl::local_filesystem()) + .unwrap(); + let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let factory = EagerPageIndexReaderFactory::new(store, metadata_cache, ScanIoSource::Local); + + let cold_metrics = ExecutionPlanMetricsSet::new(); + let mut cold_reader = factory + .create_reader(0, partitioned_file.clone(), Some(512 * 1024), &cold_metrics) + .unwrap(); + cold_reader.get_metadata(None).await.unwrap(); + + let cold_metadata_requested = + reader_metric(&cold_metrics, "scan_io_metadata_bytes_requested"); + let cold_metadata_returned = + reader_metric(&cold_metrics, "scan_io_metadata_bytes_returned"); + assert!(cold_metadata_requested > 0); + assert_eq!(cold_metadata_returned, cold_metadata_requested); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_data_bytes_requested"), + 0 + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_bytes_requested"), + cold_metadata_requested + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_local_bytes_requested"), + cold_metadata_requested + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_local_bytes_returned"), + cold_metadata_returned + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_object_store_bytes_requested"), + 0 + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_object_store_bytes_returned"), + 0 + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_metadata_cache_hits"), + 0 + ); + assert_eq!( + reader_metric(&cold_metrics, "scan_io_metadata_cache_misses"), + 1 + ); + + let warm_metrics = ExecutionPlanMetricsSet::new(); + let mut warm_reader = factory + .create_reader(0, partitioned_file, Some(512 * 1024), &warm_metrics) + .unwrap(); + warm_reader.get_metadata(None).await.unwrap(); + + assert_eq!( + reader_metric(&warm_metrics, "scan_io_metadata_bytes_requested"), + 0 + ); + assert_eq!( + reader_metric(&warm_metrics, "scan_io_metadata_bytes_returned"), + 0 + ); + assert_eq!(reader_metric(&warm_metrics, "scan_io_bytes_requested"), 0); + assert_eq!( + reader_metric(&warm_metrics, "scan_io_metadata_cache_hits"), + 1 + ); + assert_eq!( + reader_metric(&warm_metrics, "scan_io_metadata_cache_misses"), + 0 + ); + } + + #[tokio::test] + async fn reports_projected_and_predicate_pruned_data_io() { + let (filename, schema) = write_scan_io_fixture(); + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + + let full_scan = init_test_scan( + Arc::clone(&schema), + Arc::clone(&schema), + partitioned_file.clone(), + None, + None, + &session_ctx, + ); + drain_scan(&full_scan, &session_ctx).await; + + let projected_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let projected_scan = init_test_scan( + Arc::clone(&projected_schema), + Arc::clone(&schema), + partitioned_file.clone(), + Some(vec![0]), + None, + &session_ctx, + ); + drain_scan(&projected_scan, &session_ctx).await; + + let filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(500)))), + )); + let filtered_scan = init_test_scan( + projected_schema, + schema, + partitioned_file, + Some(vec![0]), + Some(vec![filter]), + &session_ctx, + ); + drain_scan(&filtered_scan, &session_ctx).await; + + let full_data = scan_metric(&full_scan, "scan_io_data_bytes_requested"); + let projected_data = scan_metric(&projected_scan, "scan_io_data_bytes_requested"); + let filtered_data = scan_metric(&filtered_scan, "scan_io_data_bytes_requested"); + assert!( + projected_data < full_data, + "projection should request fewer data bytes: projected={projected_data}, full={full_data}" + ); + assert!( + filtered_data < projected_data, + "row-group pruning should request fewer data bytes: filtered={filtered_data}, projected={projected_data}" + ); + + for scan in [&full_scan, &projected_scan, &filtered_scan] { + let data_requested = scan_metric(scan, "scan_io_data_bytes_requested"); + let data_returned = scan_metric(scan, "scan_io_data_bytes_returned"); + let metadata_requested = scan_metric(scan, "scan_io_metadata_bytes_requested"); + let metadata_returned = scan_metric(scan, "scan_io_metadata_bytes_returned"); + assert_eq!(scan_metric(scan, "bytes_scanned"), data_requested); + assert_eq!(data_returned, data_requested); + assert_eq!(metadata_returned, metadata_requested); + assert_eq!( + scan_metric(scan, "scan_io_bytes_requested"), + data_requested + metadata_requested + ); + assert_eq!( + scan_metric(scan, "scan_io_bytes_returned"), + data_returned + metadata_returned + ); + assert_eq!( + scan_metric(scan, "scan_io_local_bytes_requested"), + data_requested + metadata_requested + ); + assert_eq!( + scan_metric(scan, "scan_io_local_bytes_returned"), + data_returned + metadata_returned + ); + assert_eq!(scan_metric(scan, "scan_io_object_store_bytes_requested"), 0); + assert_eq!(scan_metric(scan, "scan_io_object_store_bytes_returned"), 0); + } + } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 6075dcc34dd..fe77694cdb3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -281,7 +281,45 @@ object CometMetricNode { "limit_matched_row_groups" -> SQLMetrics.createMetric(sc, "Number of row groups matched by limit pruning (not pruned)"), "bytes_scanned" -> - SQLMetrics.createSizeMetric(sc, "Number of bytes scanned"), + SQLMetrics.createSizeMetric( + sc, + "Legacy data-page bytes requested by native Parquet scan"), + "scan_io_bytes_requested" -> + SQLMetrics.createSizeMetric( + sc, + "Total data-page plus metadata byte ranges requested by native Parquet scan"), + "scan_io_bytes_returned" -> + SQLMetrics.createSizeMetric( + sc, + "Total data-page plus metadata bytes returned to native Parquet scan"), + "scan_io_data_bytes_requested" -> + SQLMetrics.createSizeMetric(sc, "Data-page byte ranges requested by native Parquet scan"), + "scan_io_data_bytes_returned" -> + SQLMetrics.createSizeMetric(sc, "Data-page bytes returned to native Parquet scan"), + "scan_io_metadata_bytes_requested" -> + SQLMetrics.createSizeMetric( + sc, + "Footer and page-index byte ranges requested through the storage API"), + "scan_io_metadata_bytes_returned" -> + SQLMetrics.createSizeMetric( + sc, + "Footer and page-index bytes returned through the storage API"), + "scan_io_object_store_bytes_requested" -> + SQLMetrics.createSizeMetric( + sc, + "Non-local ObjectStore API bytes requested, excluding metadata cache hits"), + "scan_io_object_store_bytes_returned" -> + SQLMetrics.createSizeMetric( + sc, + "Non-local ObjectStore API bytes returned, excluding metadata cache hits"), + "scan_io_local_bytes_requested" -> + SQLMetrics.createSizeMetric(sc, "Local file bytes requested by native Parquet scan"), + "scan_io_local_bytes_returned" -> + SQLMetrics.createSizeMetric(sc, "Local file bytes returned to native Parquet scan"), + "scan_io_metadata_cache_hits" -> + SQLMetrics.createMetric(sc, "Metadata loads served without storage I/O"), + "scan_io_metadata_cache_misses" -> + SQLMetrics.createMetric(sc, "Metadata loads requiring storage I/O"), "pushdown_rows_pruned" -> SQLMetrics.createMetric(sc, "Rows filtered out by predicates pushed into parquet scan"), "pushdown_rows_matched" -> diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index f8c325c35c6..1cd497a5d21 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -2416,10 +2416,42 @@ class CometExecSuite extends CometTestBase { assert(metrics.contains("time_elapsed_opening")) assert(metrics.contains("time_elapsed_processing")) assert(metrics.contains("time_elapsed_scanning_until_data")) + Seq( + "scan_io_bytes_requested", + "scan_io_bytes_returned", + "scan_io_data_bytes_requested", + "scan_io_data_bytes_returned", + "scan_io_metadata_bytes_requested", + "scan_io_metadata_bytes_returned", + "scan_io_object_store_bytes_requested", + "scan_io_object_store_bytes_returned", + "scan_io_local_bytes_requested", + "scan_io_local_bytes_returned", + "scan_io_metadata_cache_hits", + "scan_io_metadata_cache_misses").foreach { name => + assert(metrics.contains(name), s"Missing $name. Available: ${metrics.keys}") + } assert( metrics("time_elapsed_scanning_total").value > 0, "time_elapsed_scanning_total should be > 0") assert(metrics("bytes_scanned").value > 0, "bytes_scanned should be > 0") + assert(metrics("scan_io_data_bytes_requested").value > 0) + assert(metrics("scan_io_data_bytes_returned").value > 0) + assert(metrics("scan_io_metadata_bytes_requested").value > 0) + assert(metrics("scan_io_metadata_bytes_returned").value > 0) + assert( + metrics("scan_io_bytes_requested").value == + metrics("scan_io_data_bytes_requested").value + + metrics("scan_io_metadata_bytes_requested").value) + assert( + metrics("scan_io_bytes_returned").value == + metrics("scan_io_data_bytes_returned").value + + metrics("scan_io_metadata_bytes_returned").value) + assert(metrics("scan_io_local_bytes_requested").value > 0) + assert(metrics("scan_io_local_bytes_returned").value > 0) + assert(metrics("scan_io_object_store_bytes_requested").value == 0) + assert(metrics("scan_io_object_store_bytes_returned").value == 0) + assert(metrics("scan_io_metadata_cache_misses").value > 0) assert(metrics("output_rows").value > 0, "output_rows should be > 0") assert(metrics("time_elapsed_opening").value > 0, "time_elapsed_opening should be > 0") assert( From 3b91b65e4f717322cad0acdde5ac77c5bad9d4e0 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 18:13:43 -0700 Subject: [PATCH 02/12] Fix native Parquet scan I/O metric correctness and overhead --- .../eager_page_index_reader_factory.rs | 298 +++++++++++------- native/core/src/parquet/parquet_exec.rs | 183 ++++++++++- .../spark/sql/comet/CometMetricNode.scala | 8 +- .../apache/comet/exec/CometExecSuite.scala | 1 + 4 files changed, 368 insertions(+), 122 deletions(-) diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index ec686b9b5a1..58ead5f893e 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -89,15 +89,15 @@ pub(crate) enum ScanIoSource { /// /// requested is the sum of byte ranges requested by the Parquet reader or metadata loader. /// returned is the length of successfully returned buffers at the same boundary. Metadata -/// includes footer prefetches, footer decode follow-up reads, and page-index ranges; DataFusion -/// does not identify those subranges separately, so we report them together instead of claiming -/// unsupported footer-versus-page-index precision. +/// includes footer prefetches, footer decode follow-up reads, page-index ranges, and Bloom +/// filters; DataFusion does not identify those subranges separately, so we report them together +/// instead of claiming unsupported per-subtype precision. /// /// scan_io_object_store metrics count non-file storage API bytes and scan_io_local metrics count /// file bytes. Parsed metadata cache entries do not have a meaningful raw-byte size, so cache /// metrics report successful cache-eligible loads that did or did not require storage I/O rather /// than inventing cache-byte counts. -#[derive(Debug, Clone)] +#[derive(Debug)] struct ScanIoMetrics { bytes_requested: Count, bytes_returned: Count, @@ -115,79 +115,48 @@ struct ScanIoMetrics { } impl ScanIoMetrics { - fn new( - partition: usize, - filename: &str, - metrics: &ExecutionPlanMetricsSet, - source: ScanIoSource, - ) -> Self { + fn new(metrics: &ExecutionPlanMetricsSet, source: ScanIoSource) -> Self { Self { - bytes_requested: byte_counter(metrics, partition, filename, "scan_io_bytes_requested"), - bytes_returned: byte_counter(metrics, partition, filename, "scan_io_bytes_returned"), - data_bytes_requested: byte_counter( - metrics, - partition, - filename, - "scan_io_data_bytes_requested", - ), - data_bytes_returned: byte_counter( - metrics, - partition, - filename, - "scan_io_data_bytes_returned", - ), - metadata_bytes_requested: byte_counter( - metrics, - partition, - filename, - "scan_io_metadata_bytes_requested", - ), - metadata_bytes_returned: byte_counter( - metrics, - partition, - filename, - "scan_io_metadata_bytes_returned", - ), + bytes_requested: byte_counter(metrics, "scan_io_bytes_requested"), + bytes_returned: byte_counter(metrics, "scan_io_bytes_returned"), + data_bytes_requested: byte_counter(metrics, "scan_io_data_bytes_requested"), + data_bytes_returned: byte_counter(metrics, "scan_io_data_bytes_returned"), + metadata_bytes_requested: byte_counter(metrics, "scan_io_metadata_bytes_requested"), + metadata_bytes_returned: byte_counter(metrics, "scan_io_metadata_bytes_returned"), object_store_bytes_requested: byte_counter( metrics, - partition, - filename, "scan_io_object_store_bytes_requested", ), object_store_bytes_returned: byte_counter( metrics, - partition, - filename, "scan_io_object_store_bytes_returned", ), - local_bytes_requested: byte_counter( - metrics, - partition, - filename, - "scan_io_local_bytes_requested", - ), - local_bytes_returned: byte_counter( - metrics, - partition, - filename, - "scan_io_local_bytes_returned", - ), - metadata_cache_hits: count_counter( - metrics, - partition, - filename, - "scan_io_metadata_cache_hits", - ), - metadata_cache_misses: count_counter( - metrics, - partition, - filename, - "scan_io_metadata_cache_misses", - ), + local_bytes_requested: byte_counter(metrics, "scan_io_local_bytes_requested"), + local_bytes_returned: byte_counter(metrics, "scan_io_local_bytes_returned"), + metadata_cache_hits: count_counter(metrics, "scan_io_metadata_cache_hits"), + metadata_cache_misses: count_counter(metrics, "scan_io_metadata_cache_misses"), source, } } + fn add_reader_requested(&self, bytes: ScanIoBytes) { + if bytes.data > 0 { + self.add_data_requested(bytes.data); + } + if bytes.metadata > 0 { + self.add_metadata_requested(bytes.metadata); + } + } + + fn add_reader_returned(&self, bytes: ScanIoBytes) { + if bytes.data > 0 { + self.add_data_returned(bytes.data); + } + if bytes.metadata > 0 { + self.add_metadata_returned(bytes.metadata); + } + } + fn add_data_requested(&self, bytes: usize) { self.data_bytes_requested.add(bytes); self.add_requested(bytes); @@ -233,29 +202,30 @@ impl ScanIoMetrics { } } -fn byte_counter( - metrics: &ExecutionPlanMetricsSet, - partition: usize, - filename: &str, - name: &'static str, -) -> Count { +fn byte_counter(metrics: &ExecutionPlanMetricsSet, name: &'static str) -> Count { MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) .with_type(MetricType::Summary) .with_category(MetricCategory::Bytes) - .counter(name, partition) + .global_counter(name) } -fn count_counter( - metrics: &ExecutionPlanMetricsSet, - partition: usize, - filename: &str, - name: &'static str, -) -> Count { +fn count_counter(metrics: &ExecutionPlanMetricsSet, name: &'static str) -> Count { MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) .with_type(MetricType::Summary) - .counter(name, partition) + .global_counter(name) +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +struct ScanIoBytes { + data: usize, + metadata: usize, +} + +impl ScanIoBytes { + fn add(&mut self, other: Self) { + self.data += other.data; + self.metadata += other.metadata; + } } fn range_bytes(range: &Range) -> usize { @@ -266,11 +236,80 @@ fn ranges_bytes(ranges: &[Range]) -> usize { ranges.iter().map(range_bytes).sum() } +fn column_data_ranges(metadata: &ParquetMetaData) -> Arc<[Range]> { + let mut ranges = metadata + .row_groups() + .iter() + .flat_map(|row_group| row_group.columns()) + .filter_map(|column| { + let start = u64::try_from( + column + .dictionary_page_offset() + .unwrap_or_else(|| column.data_page_offset()), + ) + .ok()?; + let length = u64::try_from(column.compressed_size()).ok()?; + let end = start.checked_add(length)?; + (start < end).then_some(start..end) + }) + .collect::>(); + ranges.sort_unstable_by_key(|range| range.start); + + let mut merged: Vec> = Vec::with_capacity(ranges.len()); + for range in ranges { + if let Some(previous) = merged.last_mut() { + if range.start <= previous.end { + previous.end = previous.end.max(range.end); + continue; + } + } + merged.push(range); + } + Arc::from(merged) +} + +fn classify_reader_range(range: &Range, data_ranges: Option<&[Range]>) -> ScanIoBytes { + let requested = range_bytes(range); + let Some(data_ranges) = data_ranges else { + return ScanIoBytes { + data: requested, + metadata: 0, + }; + }; + + let first = data_ranges.partition_point(|data_range| data_range.end <= range.start); + let mut data = 0; + for data_range in &data_ranges[first..] { + if data_range.start >= range.end { + break; + } + let start = range.start.max(data_range.start); + let end = range.end.min(data_range.end); + if start < end { + data += range_bytes(&(start..end)); + } + } + + ScanIoBytes { + data, + metadata: requested - data, + } +} + +fn classify_returned_range( + requested: &Range, + returned: usize, + data_ranges: Option<&[Range]>, +) -> ScanIoBytes { + let end = requested.start.saturating_add(returned as u64); + classify_reader_range(&(requested.start..end), data_ranges) +} + #[derive(Debug)] pub struct EagerPageIndexReaderFactory { store: Arc, metadata_cache: Arc, - source: ScanIoSource, + scan_io_metrics: Arc, } impl EagerPageIndexReaderFactory { @@ -278,11 +317,12 @@ impl EagerPageIndexReaderFactory { store: Arc, metadata_cache: Arc, source: ScanIoSource, + metrics: &ExecutionPlanMetricsSet, ) -> Self { Self { store, metadata_cache, - source, + scan_io_metrics: Arc::new(ScanIoMetrics::new(metrics, source)), } } } @@ -300,12 +340,6 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { partitioned_file.object_meta.location.as_ref(), metrics, ); - let scan_io_metrics = ScanIoMetrics::new( - partition_index, - partitioned_file.object_meta.location.as_ref(), - metrics, - self.source, - ); let mut inner = ParquetObjectReader::new( Arc::clone(&self.store), partitioned_file.object_meta.location.clone(), @@ -317,36 +351,45 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { Ok(Box::new(EagerPageIndexReader { file_metrics, - scan_io_metrics, + scan_io_metrics: Arc::clone(&self.scan_io_metrics), store: Arc::clone(&self.store), inner, partitioned_file, metadata_cache: Arc::clone(&self.metadata_cache), metadata_size_hint, + data_ranges: None, })) } } struct EagerPageIndexReader { file_metrics: ParquetFileMetrics, - scan_io_metrics: ScanIoMetrics, + scan_io_metrics: Arc, store: Arc, inner: ParquetObjectReader, partitioned_file: PartitionedFile, metadata_cache: Arc, metadata_size_hint: Option, + data_ranges: Option]>>, } impl AsyncFileReader for EagerPageIndexReader { fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { let requested = range_bytes(&range); self.file_metrics.bytes_scanned.add(requested); - self.scan_io_metrics.add_data_requested(requested); - let scan_io_metrics = self.scan_io_metrics.clone(); + self.scan_io_metrics + .add_reader_requested(classify_reader_range(&range, self.data_ranges.as_deref())); + let scan_io_metrics = Arc::clone(&self.scan_io_metrics); + let data_ranges = self.data_ranges.clone(); + let requested_range = range.clone(); let future = self.inner.get_bytes(range); async move { let bytes = future.await?; - scan_io_metrics.add_data_returned(bytes.len()); + scan_io_metrics.add_reader_returned(classify_returned_range( + &requested_range, + bytes.len(), + data_ranges.as_deref(), + )); Ok(bytes) } .boxed() @@ -361,12 +404,38 @@ impl AsyncFileReader for EagerPageIndexReader { { let requested = ranges_bytes(&ranges); self.file_metrics.bytes_scanned.add(requested); - self.scan_io_metrics.add_data_requested(requested); - let scan_io_metrics = self.scan_io_metrics.clone(); + let requested_bytes = ranges + .iter() + .fold(ScanIoBytes::default(), |mut total, range| { + total.add(classify_reader_range(range, self.data_ranges.as_deref())); + total + }); + self.scan_io_metrics.add_reader_requested(requested_bytes); + let returned_ranges = (requested_bytes.metadata > 0).then(|| ranges.clone()); + let scan_io_metrics = Arc::clone(&self.scan_io_metrics); + let data_ranges = self.data_ranges.clone(); let future = self.inner.get_byte_ranges(ranges); async move { let bytes = future.await?; - scan_io_metrics.add_data_returned(bytes.iter().map(Bytes::len).sum()); + let returned = if let Some(ranges) = returned_ranges { + ranges.iter().zip(&bytes).fold( + ScanIoBytes::default(), + |mut total, (range, bytes)| { + total.add(classify_returned_range( + range, + bytes.len(), + data_ranges.as_deref(), + )); + total + }, + ) + } else { + ScanIoBytes { + data: bytes.iter().map(Bytes::len).sum(), + metadata: 0, + } + }; + scan_io_metrics.add_reader_returned(returned); Ok(bytes) } .boxed() @@ -382,7 +451,8 @@ impl AsyncFileReader for EagerPageIndexReader { let metadata_cache = Arc::clone(&self.metadata_cache); let store = Arc::clone(&self.store); let metadata_size_hint = self.metadata_size_hint; - let scan_io_metrics = self.scan_io_metrics.clone(); + let scan_io_metrics = Arc::clone(&self.scan_io_metrics); + let data_ranges = &mut self.data_ranges; async move { let file_decryption_properties = options .and_then(|o| o.file_decryption_properties()) @@ -396,7 +466,7 @@ impl AsyncFileReader for EagerPageIndexReader { let metadata_storage_reads = Arc::new(AtomicUsize::new(0)); let metadata_store = MetadataIoObjectStore { inner: store, - scan_io_metrics: scan_io_metrics.clone(), + scan_io_metrics: Arc::clone(&scan_io_metrics), storage_reads: Arc::clone(&metadata_storage_reads), }; @@ -414,9 +484,13 @@ impl AsyncFileReader for EagerPageIndexReader { )) }); - if cache_enabled && metadata.is_ok() { - scan_io_metrics - .record_metadata_cache_result(metadata_storage_reads.load(Ordering::Relaxed)); + if let Ok(metadata) = &metadata { + *data_ranges = Some(column_data_ranges(metadata)); + if cache_enabled { + scan_io_metrics.record_metadata_cache_result( + metadata_storage_reads.load(Ordering::Relaxed), + ); + } } metadata @@ -432,7 +506,7 @@ impl AsyncFileReader for EagerPageIndexReader { #[derive(Debug)] struct MetadataIoObjectStore { inner: Arc, - scan_io_metrics: ScanIoMetrics, + scan_io_metrics: Arc, storage_reads: Arc, } @@ -493,13 +567,19 @@ impl ObjectStore for MetadataIoObjectStore { let meta = result.meta.clone(); let range = result.range.clone(); let attributes = result.attributes.clone(); - let scan_io_metrics = self.scan_io_metrics.clone(); - let payload = GetResultPayload::Stream( - result - .into_stream() - .inspect_ok(move |bytes| scan_io_metrics.add_metadata_returned(bytes.len())) - .boxed(), - ); + let payload = if matches!(&result.payload, GetResultPayload::File(..)) { + let bytes = result.bytes().await?; + self.scan_io_metrics.add_metadata_returned(bytes.len()); + GetResultPayload::Stream(futures::stream::once(async move { Ok(bytes) }).boxed()) + } else { + let scan_io_metrics = Arc::clone(&self.scan_io_metrics); + GetResultPayload::Stream( + result + .into_stream() + .inspect_ok(move |bytes| scan_io_metrics.add_metadata_returned(bytes.len())) + .boxed(), + ) + }; Ok(GetResult { payload, diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 7ab0c1d751b..29d4c97ebc7 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -166,9 +166,13 @@ pub(crate) fn init_datasource_exec( } else { ScanIoSource::ObjectStore }; - parquet_source = parquet_source.with_parquet_file_reader_factory(Arc::new( - EagerPageIndexReaderFactory::new(store, metadata_cache, scan_io_source), + let reader_factory = Arc::new(EagerPageIndexReaderFactory::new( + store, + metadata_cache, + scan_io_source, + parquet_source.metrics(), )); + parquet_source = parquet_source.with_parquet_file_reader_factory(reader_factory); // Route data filters through `try_pushdown_filters` rather than calling // `with_predicate` directly. This is the contract DataFusion's optimizer @@ -296,8 +300,11 @@ mod tests { use arrow::array::Int32Array; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; - use datafusion::datasource::physical_plan::parquet::metadata::CachedParquetMetaData; + use datafusion::datasource::physical_plan::parquet::metadata::{ + CachedParquetMetaData, DFParquetMetadata, + }; use datafusion::datasource::physical_plan::parquet::ParquetFileReaderFactory; + use datafusion::execution::cache::cache_manager::CachedFileMetadataEntry; use datafusion::logical_expr::Operator; use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; @@ -305,6 +312,7 @@ mod tests { use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename; use futures::StreamExt; use parquet::arrow::ArrowWriter; + use parquet::file::metadata::PageIndexPolicy; use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::fs::File; @@ -543,11 +551,15 @@ mod tests { let store = runtime_env .object_store(ObjectStoreUrl::local_filesystem()) .unwrap(); - let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); - let factory = EagerPageIndexReaderFactory::new(store, metadata_cache, ScanIoSource::Local); - let cold_metrics = ExecutionPlanMetricsSet::new(); - let mut cold_reader = factory + let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let cold_factory = EagerPageIndexReaderFactory::new( + Arc::clone(&store), + Arc::clone(&metadata_cache), + ScanIoSource::Local, + &cold_metrics, + ); + let mut cold_reader = cold_factory .create_reader(0, partitioned_file.clone(), Some(512 * 1024), &cold_metrics) .unwrap(); cold_reader.get_metadata(None).await.unwrap(); @@ -592,7 +604,13 @@ mod tests { ); let warm_metrics = ExecutionPlanMetricsSet::new(); - let mut warm_reader = factory + let warm_factory = EagerPageIndexReaderFactory::new( + store, + metadata_cache, + ScanIoSource::Local, + &warm_metrics, + ); + let mut warm_reader = warm_factory .create_reader(0, partitioned_file, Some(512 * 1024), &warm_metrics) .unwrap(); warm_reader.get_metadata(None).await.unwrap(); @@ -616,6 +634,155 @@ mod tests { ); } + #[tokio::test] + async fn upgrades_partial_cached_metadata_with_local_direct_reads() { + let (filename, _schema) = write_scan_io_fixture(); + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let runtime_env = session_ctx.runtime_env(); + let store = runtime_env + .object_store(ObjectStoreUrl::local_filesystem()) + .unwrap(); + let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let footer_metadata = DFParquetMetadata::new(store.as_ref(), &partitioned_file.object_meta) + .with_page_index_policy(Some(PageIndexPolicy::Skip)) + .fetch_metadata() + .await + .unwrap(); + assert!(footer_metadata.column_index().is_none()); + assert!(footer_metadata.offset_index().is_none()); + metadata_cache.put( + &partitioned_file.object_meta.location, + CachedFileMetadataEntry::new( + partitioned_file.object_meta.clone(), + Arc::new(CachedParquetMetaData::new(footer_metadata)), + ), + ); + + let metrics = ExecutionPlanMetricsSet::new(); + let factory = + EagerPageIndexReaderFactory::new(store, metadata_cache, ScanIoSource::Local, &metrics); + let mut reader = factory + .create_reader(0, partitioned_file, Some(512 * 1024), &metrics) + .unwrap(); + let metadata = reader.get_metadata(None).await.unwrap(); + + assert!(metadata.column_index().is_some()); + assert!(metadata.offset_index().is_some()); + let requested = reader_metric(&metrics, "scan_io_metadata_bytes_requested"); + assert!(requested > 0); + assert_eq!( + reader_metric(&metrics, "scan_io_metadata_bytes_returned"), + requested + ); + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes_requested"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_hits"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_misses"), 1); + + let column = metadata.row_group(0).column(0); + let page_index_offset = u64::try_from(column.column_index_offset().unwrap()).unwrap(); + let page_index_length = u64::try_from(column.column_index_length().unwrap()).unwrap(); + reader + .get_bytes(page_index_offset..page_index_offset + page_index_length) + .await + .unwrap(); + assert_eq!( + reader_metric(&metrics, "scan_io_metadata_bytes_requested"), + requested + page_index_length as usize + ); + assert_eq!( + reader_metric(&metrics, "scan_io_metadata_bytes_returned"), + requested + page_index_length as usize + ); + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes_requested"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes_returned"), 0); + } + + #[test] + fn registers_scan_io_metrics_once_per_execution_plan() { + let (filename, _schema) = write_scan_io_fixture(); + let partitioned_file = PartitionedFile::from_path(filename).unwrap(); + let session_ctx = Arc::new(SessionContext::new()); + let runtime_env = session_ctx.runtime_env(); + let store = runtime_env + .object_store(ObjectStoreUrl::local_filesystem()) + .unwrap(); + let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = + EagerPageIndexReaderFactory::new(store, metadata_cache, ScanIoSource::Local, &metrics); + + for partition in 0..8 { + factory + .create_reader(partition, partitioned_file.clone(), None, &metrics) + .unwrap(); + } + + let registered_metrics = metrics.clone_inner(); + let scan_io_metrics = registered_metrics + .iter() + .filter(|metric| metric.value().name().starts_with("scan_io_")) + .collect::>(); + assert_eq!(scan_io_metrics.len(), 12); + for metric in scan_io_metrics { + assert!(metric.labels().is_empty()); + assert_eq!(metric.partition(), None); + } + } + + #[tokio::test] + async fn classifies_bloom_filter_reads_as_metadata() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from( + (0..500).map(|value| value * 2).collect::>(), + ))], + ) + .unwrap(); + let filename = get_temp_filename() + .as_path() + .as_os_str() + .to_str() + .unwrap() + .to_string(); + let props = WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_bloom_filter_ndv(500) + .set_bloom_filter_fpp(0.0001) + .build(); + let file = File::create(&filename).unwrap(); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(501)))), + )); + let session_ctx = Arc::new(SessionContext::new()); + let scan = init_test_scan( + Arc::clone(&schema), + schema, + PartitionedFile::from_path(filename).unwrap(), + None, + Some(vec![filter]), + &session_ctx, + ); + drain_scan(&scan, &session_ctx).await; + + let bloom_bytes = scan_metric(&scan, "bytes_scanned"); + assert!(bloom_bytes > 0); + assert_eq!(scan_metric(&scan, "scan_io_data_bytes_requested"), 0); + assert_eq!(scan_metric(&scan, "scan_io_data_bytes_returned"), 0); + assert!(scan_metric(&scan, "scan_io_metadata_bytes_requested") >= bloom_bytes); + assert_eq!( + scan_metric(&scan, "scan_io_metadata_bytes_requested"), + scan_metric(&scan, "scan_io_metadata_bytes_returned") + ); + } + #[tokio::test] async fn reports_projected_and_predicate_pruned_data_io() { let (filename, schema) = write_scan_io_fixture(); diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index fe77694cdb3..7e759bfca05 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -281,9 +281,7 @@ object CometMetricNode { "limit_matched_row_groups" -> SQLMetrics.createMetric(sc, "Number of row groups matched by limit pruning (not pruned)"), "bytes_scanned" -> - SQLMetrics.createSizeMetric( - sc, - "Legacy data-page bytes requested by native Parquet scan"), + SQLMetrics.createSizeMetric(sc, "Number of bytes scanned"), "scan_io_bytes_requested" -> SQLMetrics.createSizeMetric( sc, @@ -299,11 +297,11 @@ object CometMetricNode { "scan_io_metadata_bytes_requested" -> SQLMetrics.createSizeMetric( sc, - "Footer and page-index byte ranges requested through the storage API"), + "Footer, page-index, and Bloom-filter byte ranges requested through the storage API"), "scan_io_metadata_bytes_returned" -> SQLMetrics.createSizeMetric( sc, - "Footer and page-index bytes returned through the storage API"), + "Footer, page-index, and Bloom-filter bytes returned through the storage API"), "scan_io_object_store_bytes_requested" -> SQLMetrics.createSizeMetric( sc, diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 1cd497a5d21..035a6d1af77 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -2412,6 +2412,7 @@ class CometExecSuite extends CometTestBase { metrics.contains("time_elapsed_scanning_total"), s"Missing time_elapsed_scanning_total. Available: ${metrics.keys}") assert(metrics.contains("bytes_scanned")) + assert(metrics("bytes_scanned").name.contains("Number of bytes scanned")) assert(metrics.contains("output_rows")) assert(metrics.contains("time_elapsed_opening")) assert(metrics.contains("time_elapsed_processing")) From 9e3d0c5ae42f81bf4aba9198aaafc7097d0c0851 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 19:35:29 -0700 Subject: [PATCH 03/12] Fix scan metric shutdown and split-reader overhead --- native/core/src/execution/jni_api.rs | 37 ++++++++++- .../eager_page_index_reader_factory.rs | 65 ++++++++++++++++++- 2 files changed, 99 insertions(+), 3 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d80754736b7..048d7c63ff2 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -95,6 +95,7 @@ use std::time::{Duration, Instant}; use std::{sync::Arc, task::Poll}; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; +use tokio::task::JoinHandle; use crate::execution::memory_pools::{ create_memory_pool, handle_task_shared_pool_release, parse_memory_pool_config, MemoryPoolConfig, @@ -324,6 +325,7 @@ struct ExecutionContext { pub stream: Option, /// Receives batches from a spawned tokio task (async I/O path) pub batch_receiver: Option>>, + pub batch_producer: Option>, /// Native metrics pub metrics: Arc>>, // The interval in milliseconds to update metrics @@ -537,6 +539,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources, stream: None, batch_receiver: None, + batch_producer: None, metrics, metrics_update_interval, metrics_last_update_time: Instant::now(), @@ -835,7 +838,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // decreasing to 1 would serialize production and consumption. let (tx, rx) = mpsc::channel(2); let mut stream = stream; - get_runtime().spawn(async move { + let producer = get_runtime().spawn(async move { let result = std::panic::AssertUnwindSafe(async { while let Some(batch) = stream.next().await { if tx.send(batch).await.is_err() { @@ -862,6 +865,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } }); exec_context.batch_receiver = Some(rx); + exec_context.batch_producer = Some(producer); } else { exec_context.stream = Some(stream); } @@ -966,6 +970,11 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( try_unwrap_or_throw(&e, |env| unsafe { let execution_context = get_execution_context(exec_context); + execution_context.batch_receiver.take(); + if let Some(producer) = execution_context.batch_producer.take() { + stop_batch_producer(producer); + } + // Update metrics update_metrics(env, execution_context)?; @@ -989,6 +998,11 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( }) } +fn stop_batch_producer(producer: JoinHandle<()>) { + producer.abort(); + let _ = get_runtime().block_on(producer); +} + fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { if let Some(native_query) = &exec_context.root_op { let metrics = exec_context.metrics.as_obj(); @@ -1375,3 +1389,24 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( Ok(()) }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn waits_for_background_batch_producer_shutdown() { + let (sender, receiver) = std::sync::mpsc::channel::<()>(); + let producer = get_runtime().spawn(async move { + let _sender = sender; + futures::future::pending::<()>().await; + }); + + stop_batch_producer(producer); + + assert_eq!( + receiver.try_recv(), + Err(std::sync::mpsc::TryRecvError::Disconnected) + ); + } +} diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 58ead5f893e..d3871df904f 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -66,13 +66,15 @@ use object_store::{ ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as ObjectStoreResult, }; +use parking_lot::Mutex; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +use std::collections::HashMap; use std::fmt::{Debug, Display, Formatter}; use std::ops::Range; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Weak}; /// Whether the reader's storage API represents a local file or a non-local object store. /// @@ -85,6 +87,8 @@ pub(crate) enum ScanIoSource { Local, } +type ColumnDataRangeCache = HashMap, Arc<[Range]>)>; + /// Scan I/O metrics at the boundaries this reader can observe without guessing. /// /// requested is the sum of byte ranges requested by the Parquet reader or metadata loader. @@ -111,6 +115,7 @@ struct ScanIoMetrics { local_bytes_returned: Count, metadata_cache_hits: Count, metadata_cache_misses: Count, + column_data_ranges: Mutex, source: ScanIoSource, } @@ -135,10 +140,33 @@ impl ScanIoMetrics { local_bytes_returned: byte_counter(metrics, "scan_io_local_bytes_returned"), metadata_cache_hits: count_counter(metrics, "scan_io_metadata_cache_hits"), metadata_cache_misses: count_counter(metrics, "scan_io_metadata_cache_misses"), + column_data_ranges: Mutex::new(HashMap::new()), source, } } + fn column_data_ranges( + &self, + location: &Path, + metadata: &Arc, + ) -> Arc<[Range]> { + let metadata_identity = Arc::downgrade(metadata); + let mut cached_ranges = self.column_data_ranges.lock(); + if let Some((cached_metadata, ranges)) = cached_ranges.get(location) { + if cached_metadata.ptr_eq(&metadata_identity) { + return Arc::clone(ranges); + } + } + + if cached_ranges.len() == cached_ranges.capacity() { + cached_ranges.retain(|_, (cached_metadata, _)| cached_metadata.strong_count() > 0); + } + + let ranges = column_data_ranges(metadata); + cached_ranges.insert(location.clone(), (metadata_identity, Arc::clone(&ranges))); + ranges + } + fn add_reader_requested(&self, bytes: ScanIoBytes) { if bytes.data > 0 { self.add_data_requested(bytes.data); @@ -485,7 +513,8 @@ impl AsyncFileReader for EagerPageIndexReader { }); if let Ok(metadata) = &metadata { - *data_ranges = Some(column_data_ranges(metadata)); + *data_ranges = + Some(scan_io_metrics.column_data_ranges(&object_meta.location, metadata)); if cache_enabled { scan_io_metrics.record_metadata_cache_result( metadata_storage_reads.load(Ordering::Relaxed), @@ -658,3 +687,35 @@ impl Drop for EagerPageIndexReader { .set_total(self.partitioned_file.object_meta.size as usize); } } + +#[cfg(test)] +mod tests { + use super::*; + use parquet::file::metadata::FileMetaData; + use parquet::schema::types::{SchemaDescriptor, Type}; + + #[test] + fn shares_column_data_ranges_by_metadata_identity() { + let schema = Arc::new(SchemaDescriptor::new(Arc::new( + Type::group_type_builder("schema").build().unwrap(), + ))); + let metadata = Arc::new(ParquetMetaData::new( + FileMetaData::new(1, 0, None, None, schema, None), + vec![], + )); + let metrics = ScanIoMetrics::new(&ExecutionPlanMetricsSet::new(), ScanIoSource::Local); + let location = Path::from("test.parquet"); + + let first = metrics.column_data_ranges(&location, &metadata); + let second = metrics.column_data_ranges(&location, &metadata); + assert!(Arc::ptr_eq(&first, &second)); + + let replacement = Arc::new(metadata.as_ref().clone()); + let refreshed = metrics.column_data_ranges(&location, &replacement); + assert!(!Arc::ptr_eq(&first, &refreshed)); + assert!(Arc::ptr_eq( + &refreshed, + &metrics.column_data_ranges(&location, &replacement) + )); + } +} From 372ebb3cff63ccb820b7e41a23ea8b43ed0c6708 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 20:13:18 -0700 Subject: [PATCH 04/12] Simplify native scan I/O accounting boundaries --- native/core/src/execution/jni_api.rs | 37 +-- native/core/src/execution/metrics/utils.rs | 5 +- .../eager_page_index_reader_factory.rs | 217 +----------------- 3 files changed, 12 insertions(+), 247 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 048d7c63ff2..d80754736b7 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -95,7 +95,6 @@ use std::time::{Duration, Instant}; use std::{sync::Arc, task::Poll}; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; -use tokio::task::JoinHandle; use crate::execution::memory_pools::{ create_memory_pool, handle_task_shared_pool_release, parse_memory_pool_config, MemoryPoolConfig, @@ -325,7 +324,6 @@ struct ExecutionContext { pub stream: Option, /// Receives batches from a spawned tokio task (async I/O path) pub batch_receiver: Option>>, - pub batch_producer: Option>, /// Native metrics pub metrics: Arc>>, // The interval in milliseconds to update metrics @@ -539,7 +537,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources, stream: None, batch_receiver: None, - batch_producer: None, metrics, metrics_update_interval, metrics_last_update_time: Instant::now(), @@ -838,7 +835,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // decreasing to 1 would serialize production and consumption. let (tx, rx) = mpsc::channel(2); let mut stream = stream; - let producer = get_runtime().spawn(async move { + get_runtime().spawn(async move { let result = std::panic::AssertUnwindSafe(async { while let Some(batch) = stream.next().await { if tx.send(batch).await.is_err() { @@ -865,7 +862,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } }); exec_context.batch_receiver = Some(rx); - exec_context.batch_producer = Some(producer); } else { exec_context.stream = Some(stream); } @@ -970,11 +966,6 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( try_unwrap_or_throw(&e, |env| unsafe { let execution_context = get_execution_context(exec_context); - execution_context.batch_receiver.take(); - if let Some(producer) = execution_context.batch_producer.take() { - stop_batch_producer(producer); - } - // Update metrics update_metrics(env, execution_context)?; @@ -998,11 +989,6 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( }) } -fn stop_batch_producer(producer: JoinHandle<()>) { - producer.abort(); - let _ = get_runtime().block_on(producer); -} - fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { if let Some(native_query) = &exec_context.root_op { let metrics = exec_context.metrics.as_obj(); @@ -1389,24 +1375,3 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( Ok(()) }) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn waits_for_background_batch_producer_shutdown() { - let (sender, receiver) = std::sync::mpsc::channel::<()>(); - let producer = get_runtime().spawn(async move { - let _sender = sender; - futures::future::pending::<()>().await; - }); - - stop_batch_producer(producer); - - assert_eq!( - receiver.try_recv(), - Err(std::sync::mpsc::TryRecvError::Disconnected) - ); - } -} diff --git a/native/core/src/execution/metrics/utils.rs b/native/core/src/execution/metrics/utils.rs index 4104aaa2e02..50a35393496 100644 --- a/native/core/src/execution/metrics/utils.rs +++ b/native/core/src/execution/metrics/utils.rs @@ -65,8 +65,9 @@ pub(crate) fn to_native_metric_node( let children = spark_plan.children(); let mut native_metric_node = NativeMetricNode { - // Native scans expose the largest metric map, including scan I/O breakdowns. - metrics: HashMap::with_capacity(32), + // Most operator metric maps are well under 20 entries (e.g. hash-join: 9, + // native-scan: ~20). Pre-sizing to 16 avoids the default-capacity rehash. + metrics: HashMap::with_capacity(16), children: Vec::with_capacity(children.len()), }; diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index d3871df904f..61389e9a2f4 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -66,15 +66,13 @@ use object_store::{ ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as ObjectStoreResult, }; -use parking_lot::Mutex; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; -use std::collections::HashMap; use std::fmt::{Debug, Display, Formatter}; use std::ops::Range; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Weak}; +use std::sync::Arc; /// Whether the reader's storage API represents a local file or a non-local object store. /// @@ -87,8 +85,6 @@ pub(crate) enum ScanIoSource { Local, } -type ColumnDataRangeCache = HashMap, Arc<[Range]>)>; - /// Scan I/O metrics at the boundaries this reader can observe without guessing. /// /// requested is the sum of byte ranges requested by the Parquet reader or metadata loader. @@ -115,7 +111,6 @@ struct ScanIoMetrics { local_bytes_returned: Count, metadata_cache_hits: Count, metadata_cache_misses: Count, - column_data_ranges: Mutex, source: ScanIoSource, } @@ -140,51 +135,10 @@ impl ScanIoMetrics { local_bytes_returned: byte_counter(metrics, "scan_io_local_bytes_returned"), metadata_cache_hits: count_counter(metrics, "scan_io_metadata_cache_hits"), metadata_cache_misses: count_counter(metrics, "scan_io_metadata_cache_misses"), - column_data_ranges: Mutex::new(HashMap::new()), source, } } - fn column_data_ranges( - &self, - location: &Path, - metadata: &Arc, - ) -> Arc<[Range]> { - let metadata_identity = Arc::downgrade(metadata); - let mut cached_ranges = self.column_data_ranges.lock(); - if let Some((cached_metadata, ranges)) = cached_ranges.get(location) { - if cached_metadata.ptr_eq(&metadata_identity) { - return Arc::clone(ranges); - } - } - - if cached_ranges.len() == cached_ranges.capacity() { - cached_ranges.retain(|_, (cached_metadata, _)| cached_metadata.strong_count() > 0); - } - - let ranges = column_data_ranges(metadata); - cached_ranges.insert(location.clone(), (metadata_identity, Arc::clone(&ranges))); - ranges - } - - fn add_reader_requested(&self, bytes: ScanIoBytes) { - if bytes.data > 0 { - self.add_data_requested(bytes.data); - } - if bytes.metadata > 0 { - self.add_metadata_requested(bytes.metadata); - } - } - - fn add_reader_returned(&self, bytes: ScanIoBytes) { - if bytes.data > 0 { - self.add_data_returned(bytes.data); - } - if bytes.metadata > 0 { - self.add_metadata_returned(bytes.metadata); - } - } - fn add_data_requested(&self, bytes: usize) { self.data_bytes_requested.add(bytes); self.add_requested(bytes); @@ -243,19 +197,6 @@ fn count_counter(metrics: &ExecutionPlanMetricsSet, name: &'static str) -> Count .global_counter(name) } -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -struct ScanIoBytes { - data: usize, - metadata: usize, -} - -impl ScanIoBytes { - fn add(&mut self, other: Self) { - self.data += other.data; - self.metadata += other.metadata; - } -} - fn range_bytes(range: &Range) -> usize { (range.end - range.start) as usize } @@ -264,75 +205,6 @@ fn ranges_bytes(ranges: &[Range]) -> usize { ranges.iter().map(range_bytes).sum() } -fn column_data_ranges(metadata: &ParquetMetaData) -> Arc<[Range]> { - let mut ranges = metadata - .row_groups() - .iter() - .flat_map(|row_group| row_group.columns()) - .filter_map(|column| { - let start = u64::try_from( - column - .dictionary_page_offset() - .unwrap_or_else(|| column.data_page_offset()), - ) - .ok()?; - let length = u64::try_from(column.compressed_size()).ok()?; - let end = start.checked_add(length)?; - (start < end).then_some(start..end) - }) - .collect::>(); - ranges.sort_unstable_by_key(|range| range.start); - - let mut merged: Vec> = Vec::with_capacity(ranges.len()); - for range in ranges { - if let Some(previous) = merged.last_mut() { - if range.start <= previous.end { - previous.end = previous.end.max(range.end); - continue; - } - } - merged.push(range); - } - Arc::from(merged) -} - -fn classify_reader_range(range: &Range, data_ranges: Option<&[Range]>) -> ScanIoBytes { - let requested = range_bytes(range); - let Some(data_ranges) = data_ranges else { - return ScanIoBytes { - data: requested, - metadata: 0, - }; - }; - - let first = data_ranges.partition_point(|data_range| data_range.end <= range.start); - let mut data = 0; - for data_range in &data_ranges[first..] { - if data_range.start >= range.end { - break; - } - let start = range.start.max(data_range.start); - let end = range.end.min(data_range.end); - if start < end { - data += range_bytes(&(start..end)); - } - } - - ScanIoBytes { - data, - metadata: requested - data, - } -} - -fn classify_returned_range( - requested: &Range, - returned: usize, - data_ranges: Option<&[Range]>, -) -> ScanIoBytes { - let end = requested.start.saturating_add(returned as u64); - classify_reader_range(&(requested.start..end), data_ranges) -} - #[derive(Debug)] pub struct EagerPageIndexReaderFactory { store: Arc, @@ -385,7 +257,6 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { partitioned_file, metadata_cache: Arc::clone(&self.metadata_cache), metadata_size_hint, - data_ranges: None, })) } } @@ -398,26 +269,18 @@ struct EagerPageIndexReader { partitioned_file: PartitionedFile, metadata_cache: Arc, metadata_size_hint: Option, - data_ranges: Option]>>, } impl AsyncFileReader for EagerPageIndexReader { fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { let requested = range_bytes(&range); self.file_metrics.bytes_scanned.add(requested); - self.scan_io_metrics - .add_reader_requested(classify_reader_range(&range, self.data_ranges.as_deref())); + self.scan_io_metrics.add_metadata_requested(requested); let scan_io_metrics = Arc::clone(&self.scan_io_metrics); - let data_ranges = self.data_ranges.clone(); - let requested_range = range.clone(); let future = self.inner.get_bytes(range); async move { let bytes = future.await?; - scan_io_metrics.add_reader_returned(classify_returned_range( - &requested_range, - bytes.len(), - data_ranges.as_deref(), - )); + scan_io_metrics.add_metadata_returned(bytes.len()); Ok(bytes) } .boxed() @@ -432,38 +295,12 @@ impl AsyncFileReader for EagerPageIndexReader { { let requested = ranges_bytes(&ranges); self.file_metrics.bytes_scanned.add(requested); - let requested_bytes = ranges - .iter() - .fold(ScanIoBytes::default(), |mut total, range| { - total.add(classify_reader_range(range, self.data_ranges.as_deref())); - total - }); - self.scan_io_metrics.add_reader_requested(requested_bytes); - let returned_ranges = (requested_bytes.metadata > 0).then(|| ranges.clone()); + self.scan_io_metrics.add_data_requested(requested); let scan_io_metrics = Arc::clone(&self.scan_io_metrics); - let data_ranges = self.data_ranges.clone(); let future = self.inner.get_byte_ranges(ranges); async move { let bytes = future.await?; - let returned = if let Some(ranges) = returned_ranges { - ranges.iter().zip(&bytes).fold( - ScanIoBytes::default(), - |mut total, (range, bytes)| { - total.add(classify_returned_range( - range, - bytes.len(), - data_ranges.as_deref(), - )); - total - }, - ) - } else { - ScanIoBytes { - data: bytes.iter().map(Bytes::len).sum(), - metadata: 0, - } - }; - scan_io_metrics.add_reader_returned(returned); + scan_io_metrics.add_data_returned(bytes.iter().map(Bytes::len).sum()); Ok(bytes) } .boxed() @@ -480,7 +317,6 @@ impl AsyncFileReader for EagerPageIndexReader { let store = Arc::clone(&self.store); let metadata_size_hint = self.metadata_size_hint; let scan_io_metrics = Arc::clone(&self.scan_io_metrics); - let data_ranges = &mut self.data_ranges; async move { let file_decryption_properties = options .and_then(|o| o.file_decryption_properties()) @@ -512,14 +348,9 @@ impl AsyncFileReader for EagerPageIndexReader { )) }); - if let Ok(metadata) = &metadata { - *data_ranges = - Some(scan_io_metrics.column_data_ranges(&object_meta.location, metadata)); - if cache_enabled { - scan_io_metrics.record_metadata_cache_result( - metadata_storage_reads.load(Ordering::Relaxed), - ); - } + if metadata.is_ok() && cache_enabled { + scan_io_metrics + .record_metadata_cache_result(metadata_storage_reads.load(Ordering::Relaxed)); } metadata @@ -687,35 +518,3 @@ impl Drop for EagerPageIndexReader { .set_total(self.partitioned_file.object_meta.size as usize); } } - -#[cfg(test)] -mod tests { - use super::*; - use parquet::file::metadata::FileMetaData; - use parquet::schema::types::{SchemaDescriptor, Type}; - - #[test] - fn shares_column_data_ranges_by_metadata_identity() { - let schema = Arc::new(SchemaDescriptor::new(Arc::new( - Type::group_type_builder("schema").build().unwrap(), - ))); - let metadata = Arc::new(ParquetMetaData::new( - FileMetaData::new(1, 0, None, None, schema, None), - vec![], - )); - let metrics = ScanIoMetrics::new(&ExecutionPlanMetricsSet::new(), ScanIoSource::Local); - let location = Path::from("test.parquet"); - - let first = metrics.column_data_ranges(&location, &metadata); - let second = metrics.column_data_ranges(&location, &metadata); - assert!(Arc::ptr_eq(&first, &second)); - - let replacement = Arc::new(metadata.as_ref().clone()); - let refreshed = metrics.column_data_ranges(&location, &replacement); - assert!(!Arc::ptr_eq(&first, &refreshed)); - assert!(Arc::ptr_eq( - &refreshed, - &metrics.column_data_ranges(&location, &replacement) - )); - } -} From f0f1d36f2d205037bb113f16d6cf847dd9dbba0b Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 21:49:18 -0700 Subject: [PATCH 05/12] Measure native Parquet read amplification across storage boundaries --- native/core/src/execution/planner.rs | 5 +- .../eager_page_index_reader_factory.rs | 254 +++++++------- native/core/src/parquet/mod.rs | 5 +- native/core/src/parquet/parquet_exec.rs | 314 +++++++++++++----- native/core/src/parquet/parquet_support.rs | 51 ++- .../spark/sql/comet/CometMetricNode.scala | 40 +-- .../apache/comet/exec/CometExecSuite.scala | 40 +-- 7 files changed, 458 insertions(+), 251 deletions(-) diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index d109627e825..76e5f725114 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -94,7 +94,7 @@ use iceberg::expr::Bind; use crate::execution::operators::ExecutionError::GeneralError; use crate::execution::shuffle::{CometPartitioning, CompressionCodec}; use crate::execution::spark_plan::SparkPlan; -use crate::parquet::parquet_support::prepare_object_store_with_configs; +use crate::parquet::parquet_support::{is_hdfs_scheme, prepare_object_store_with_configs}; use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, @@ -1631,6 +1631,8 @@ impl PhysicalPlanner { .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(); + let is_hdfs_object_store = url::Url::parse(&one_file) + .is_ok_and(|url| is_hdfs_scheme(&url, &object_store_options)); let (object_store_url, _) = prepare_object_store_with_configs( self.session_ctx.runtime_env(), one_file, @@ -1646,6 +1648,7 @@ impl PhysicalPlanner { Some(data_schema), Some(partition_schema), object_store_url, + is_hdfs_object_store, file_groups, Some(projection_vector), Some(data_filters?), diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 61389e9a2f4..bb2a18fa9fd 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -62,103 +62,60 @@ use futures::future::BoxFuture; use futures::{FutureExt, StreamExt, TryStreamExt}; use object_store::path::Path; use object_store::{ - CopyOptions, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload, - ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, - Result as ObjectStoreResult, + coalesce_ranges, CopyOptions, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, + MultipartUpload, ObjectMeta, ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, + PutPayload, PutResult, RenameOptions, Result as ObjectStoreResult, + OBJECT_STORE_COALESCE_DEFAULT, }; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; -use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{FooterTail, PageIndexPolicy, ParquetMetaData}; use std::fmt::{Debug, Display, Formatter}; use std::ops::Range; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -/// Whether the reader's storage API represents a local file or a non-local object store. -/// -/// The source metrics intentionally describe the API boundary, not network wire bytes. A remote -/// backend can coalesce ranges, retry requests, or serve bytes from a cache below ObjectStore, -/// none of which this reader can observe without backend-specific hooks. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum ScanIoSource { ObjectStore, Local, + OtherObjectStore, } -/// Scan I/O metrics at the boundaries this reader can observe without guessing. -/// -/// requested is the sum of byte ranges requested by the Parquet reader or metadata loader. -/// returned is the length of successfully returned buffers at the same boundary. Metadata -/// includes footer prefetches, footer decode follow-up reads, page-index ranges, and Bloom -/// filters; DataFusion does not identify those subranges separately, so we report them together -/// instead of claiming unsupported per-subtype precision. -/// -/// scan_io_object_store metrics count non-file storage API bytes and scan_io_local metrics count -/// file bytes. Parsed metadata cache entries do not have a meaningful raw-byte size, so cache -/// metrics report successful cache-eligible loads that did or did not require storage I/O rather -/// than inventing cache-byte counts. #[derive(Debug)] struct ScanIoMetrics { - bytes_requested: Count, - bytes_returned: Count, - data_bytes_requested: Count, - data_bytes_returned: Count, - metadata_bytes_requested: Count, - metadata_bytes_returned: Count, - object_store_bytes_requested: Count, - object_store_bytes_returned: Count, - local_bytes_requested: Count, - local_bytes_returned: Count, + data_bytes: Count, + metadata_bytes: Count, + footer_reads: Count, + footer_bytes: Count, + object_store_get_calls: Count, + object_store_get_requested_bytes: Count, + object_store_response_bytes_read: Count, metadata_cache_hits: Count, metadata_cache_misses: Count, - source: ScanIoSource, } impl ScanIoMetrics { - fn new(metrics: &ExecutionPlanMetricsSet, source: ScanIoSource) -> Self { + fn new(metrics: &ExecutionPlanMetricsSet) -> Self { Self { - bytes_requested: byte_counter(metrics, "scan_io_bytes_requested"), - bytes_returned: byte_counter(metrics, "scan_io_bytes_returned"), - data_bytes_requested: byte_counter(metrics, "scan_io_data_bytes_requested"), - data_bytes_returned: byte_counter(metrics, "scan_io_data_bytes_returned"), - metadata_bytes_requested: byte_counter(metrics, "scan_io_metadata_bytes_requested"), - metadata_bytes_returned: byte_counter(metrics, "scan_io_metadata_bytes_returned"), - object_store_bytes_requested: byte_counter( + data_bytes: byte_counter(metrics, "scan_io_data_bytes"), + metadata_bytes: byte_counter(metrics, "scan_io_metadata_bytes"), + footer_reads: count_counter(metrics, "scan_io_footer_reads"), + footer_bytes: byte_counter(metrics, "scan_io_footer_bytes"), + object_store_get_calls: count_counter(metrics, "scan_io_object_store_get_calls"), + object_store_get_requested_bytes: byte_counter( metrics, - "scan_io_object_store_bytes_requested", + "scan_io_object_store_get_requested_bytes", ), - object_store_bytes_returned: byte_counter( + object_store_response_bytes_read: byte_counter( metrics, - "scan_io_object_store_bytes_returned", + "scan_io_object_store_response_bytes_read", ), - local_bytes_requested: byte_counter(metrics, "scan_io_local_bytes_requested"), - local_bytes_returned: byte_counter(metrics, "scan_io_local_bytes_returned"), metadata_cache_hits: count_counter(metrics, "scan_io_metadata_cache_hits"), metadata_cache_misses: count_counter(metrics, "scan_io_metadata_cache_misses"), - source, } } - fn add_data_requested(&self, bytes: usize) { - self.data_bytes_requested.add(bytes); - self.add_requested(bytes); - } - - fn add_data_returned(&self, bytes: usize) { - self.data_bytes_returned.add(bytes); - self.add_returned(bytes); - } - - fn add_metadata_requested(&self, bytes: usize) { - self.metadata_bytes_requested.add(bytes); - self.add_requested(bytes); - } - - fn add_metadata_returned(&self, bytes: usize) { - self.metadata_bytes_returned.add(bytes); - self.add_returned(bytes); - } - fn record_metadata_cache_result(&self, storage_reads: usize) { if storage_reads == 0 { self.metadata_cache_hits.add(1); @@ -166,22 +123,6 @@ impl ScanIoMetrics { self.metadata_cache_misses.add(1); } } - - fn add_requested(&self, bytes: usize) { - self.bytes_requested.add(bytes); - match self.source { - ScanIoSource::ObjectStore => self.object_store_bytes_requested.add(bytes), - ScanIoSource::Local => self.local_bytes_requested.add(bytes), - } - } - - fn add_returned(&self, bytes: usize) { - self.bytes_returned.add(bytes); - match self.source { - ScanIoSource::ObjectStore => self.object_store_bytes_returned.add(bytes), - ScanIoSource::Local => self.local_bytes_returned.add(bytes), - } - } } fn byte_counter(metrics: &ExecutionPlanMetricsSet, name: &'static str) -> Count { @@ -219,10 +160,20 @@ impl EagerPageIndexReaderFactory { source: ScanIoSource, metrics: &ExecutionPlanMetricsSet, ) -> Self { + let scan_io_metrics = Arc::new(ScanIoMetrics::new(metrics)); + let store: Arc = if source == ScanIoSource::ObjectStore { + Arc::new(ScanIoObjectStore { + inner: store, + scan_io_metrics: Arc::clone(&scan_io_metrics), + role: ScanIoStoreRole::ObjectStore, + }) + } else { + store + }; Self { store, metadata_cache, - scan_io_metrics: Arc::new(ScanIoMetrics::new(metrics, source)), + scan_io_metrics, } } } @@ -275,12 +226,11 @@ impl AsyncFileReader for EagerPageIndexReader { fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { let requested = range_bytes(&range); self.file_metrics.bytes_scanned.add(requested); - self.scan_io_metrics.add_metadata_requested(requested); let scan_io_metrics = Arc::clone(&self.scan_io_metrics); let future = self.inner.get_bytes(range); async move { let bytes = future.await?; - scan_io_metrics.add_metadata_returned(bytes.len()); + scan_io_metrics.metadata_bytes.add(bytes.len()); Ok(bytes) } .boxed() @@ -295,12 +245,13 @@ impl AsyncFileReader for EagerPageIndexReader { { let requested = ranges_bytes(&ranges); self.file_metrics.bytes_scanned.add(requested); - self.scan_io_metrics.add_data_requested(requested); let scan_io_metrics = Arc::clone(&self.scan_io_metrics); let future = self.inner.get_byte_ranges(ranges); async move { let bytes = future.await?; - scan_io_metrics.add_data_returned(bytes.iter().map(Bytes::len).sum()); + scan_io_metrics + .data_bytes + .add(bytes.iter().map(Bytes::len).sum()); Ok(bytes) } .boxed() @@ -328,10 +279,15 @@ impl AsyncFileReader for EagerPageIndexReader { options.map(|o| o.column_index_policy()) }; let metadata_storage_reads = Arc::new(AtomicUsize::new(0)); - let metadata_store = MetadataIoObjectStore { + let footer_payload_bytes = Arc::new(AtomicUsize::new(0)); + let metadata_store = ScanIoObjectStore { inner: store, scan_io_metrics: Arc::clone(&scan_io_metrics), - storage_reads: Arc::clone(&metadata_storage_reads), + role: ScanIoStoreRole::Metadata { + storage_reads: Arc::clone(&metadata_storage_reads), + footer_payload_bytes: Arc::clone(&footer_payload_bytes), + file_size: object_meta.size, + }, }; let metadata = DFParquetMetadata::new(&metadata_store, &object_meta) @@ -348,9 +304,17 @@ impl AsyncFileReader for EagerPageIndexReader { )) }); - if metadata.is_ok() && cache_enabled { - scan_io_metrics - .record_metadata_cache_result(metadata_storage_reads.load(Ordering::Relaxed)); + if metadata.is_ok() { + let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); + if footer_bytes > 0 { + scan_io_metrics.footer_reads.add(1); + scan_io_metrics.footer_bytes.add(footer_bytes); + } + if cache_enabled { + scan_io_metrics.record_metadata_cache_result( + metadata_storage_reads.load(Ordering::Relaxed), + ); + } } metadata @@ -359,35 +323,77 @@ impl AsyncFileReader for EagerPageIndexReader { } } -/// Counts metadata storage calls while forwarding every operation to the configured store. -/// -/// Data reads are already visible at the AsyncFileReader data methods. Metadata reads bypass -/// those methods inside DFParquetMetadata, so only metadata uses this wrapper. #[derive(Debug)] -struct MetadataIoObjectStore { +enum ScanIoStoreRole { + ObjectStore, + Metadata { + storage_reads: Arc, + footer_payload_bytes: Arc, + file_size: u64, + }, +} + +#[derive(Debug)] +struct ScanIoObjectStore { inner: Arc, scan_io_metrics: Arc, - storage_reads: Arc, + role: ScanIoStoreRole, } -impl MetadataIoObjectStore { +impl ScanIoObjectStore { fn record_request(&self, bytes: usize) { - if bytes > 0 { - self.storage_reads.fetch_add(1, Ordering::Relaxed); - self.scan_io_metrics.add_metadata_requested(bytes); + if bytes == 0 { + return; + } + match &self.role { + ScanIoStoreRole::ObjectStore => { + self.scan_io_metrics.object_store_get_calls.add(1); + self.scan_io_metrics + .object_store_get_requested_bytes + .add(bytes); + } + ScanIoStoreRole::Metadata { storage_reads, .. } => { + storage_reads.fetch_add(1, Ordering::Relaxed); + } + } + } + + fn record_returned(&self, range: Option<&Range>, bytes: &Bytes) { + match &self.role { + ScanIoStoreRole::ObjectStore => self + .scan_io_metrics + .object_store_response_bytes_read + .add(bytes.len()), + ScanIoStoreRole::Metadata { + footer_payload_bytes, + file_size, + .. + } => { + self.scan_io_metrics.metadata_bytes.add(bytes.len()); + if range.is_some_and(|range| range.end == *file_size) && bytes.len() >= 8 { + if let Ok(footer) = FooterTail::try_from(&bytes[bytes.len() - 8..]) { + let _ = footer_payload_bytes.compare_exchange( + 0, + footer.metadata_length(), + Ordering::Relaxed, + Ordering::Relaxed, + ); + } + } + } } } } -impl Display for MetadataIoObjectStore { +impl Display for ScanIoObjectStore { fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { - write!(formatter, "metadata-io({})", self.inner) + write!(formatter, "scan-io({})", self.inner) } } #[async_trait] #[deny(clippy::missing_trait_methods)] -impl ObjectStore for MetadataIoObjectStore { +impl ObjectStore for ScanIoObjectStore { async fn put_opts( &self, location: &Path, @@ -429,14 +435,23 @@ impl ObjectStore for MetadataIoObjectStore { let attributes = result.attributes.clone(); let payload = if matches!(&result.payload, GetResultPayload::File(..)) { let bytes = result.bytes().await?; - self.scan_io_metrics.add_metadata_returned(bytes.len()); + self.record_returned(Some(&range), &bytes); GetResultPayload::Stream(futures::stream::once(async move { Ok(bytes) }).boxed()) } else { let scan_io_metrics = Arc::clone(&self.scan_io_metrics); + let metadata_read = matches!(self.role, ScanIoStoreRole::Metadata { .. }); GetResultPayload::Stream( result .into_stream() - .inspect_ok(move |bytes| scan_io_metrics.add_metadata_returned(bytes.len())) + .inspect_ok(move |bytes| { + if metadata_read { + scan_io_metrics.metadata_bytes.add(bytes.len()); + } else { + scan_io_metrics + .object_store_response_bytes_read + .add(bytes.len()); + } + }) .boxed(), ) }; @@ -454,11 +469,24 @@ impl ObjectStore for MetadataIoObjectStore { location: &Path, ranges: &[Range], ) -> ObjectStoreResult> { - self.record_request(ranges_bytes(ranges)); - let bytes = self.inner.get_ranges(location, ranges).await?; - self.scan_io_metrics - .add_metadata_returned(bytes.iter().map(Bytes::len).sum()); - Ok(bytes) + match &self.role { + ScanIoStoreRole::ObjectStore => { + coalesce_ranges( + ranges, + |range| self.get_range(location, range), + OBJECT_STORE_COALESCE_DEFAULT, + ) + .await + } + ScanIoStoreRole::Metadata { .. } => { + self.record_request(ranges_bytes(ranges)); + let bytes = self.inner.get_ranges(location, ranges).await?; + for (range, bytes) in ranges.iter().zip(bytes.iter()) { + self.record_returned(Some(range), bytes); + } + Ok(bytes) + } + } } fn delete_stream( diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index cfa03220c10..8c27698b4d7 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -49,7 +49,7 @@ use crate::execution::utils::SparkArrowConvert; use crate::jvm_bridge::JVMClasses; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_exec::init_datasource_exec; -use crate::parquet::parquet_support::prepare_object_store_with_configs; +use crate::parquet::parquet_support::{is_hdfs_scheme, prepare_object_store_with_configs}; use arrow::array::{Array, RecordBatch}; use datafusion::datasource::listing::PartitionedFile; use datafusion::execution::SendableRecordBatchStream; @@ -160,6 +160,8 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat let path: String = file_path.try_to_string(env).unwrap(); let object_store_config = get_object_store_options(env, object_store_options)?; + let is_hdfs_object_store = + url::Url::parse(&path).is_ok_and(|url| is_hdfs_scheme(&url, &object_store_config)); let (object_store_url, object_store_path) = prepare_object_store_with_configs( session_ctx.runtime_env(), path.clone(), @@ -213,6 +215,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat Some(data_schema), None, object_store_url, + is_hdfs_object_store, file_groups, None, data_filters, diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 29d4c97ebc7..055b216fc83 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -62,6 +62,7 @@ pub(crate) fn init_datasource_exec( data_schema: Option, partition_schema: Option, object_store_url: ObjectStoreUrl, + is_hdfs_object_store: bool, file_groups: Vec>, projection_vector: Option>, data_filters: Option>>, @@ -161,11 +162,7 @@ pub(crate) fn init_datasource_exec( let runtime_env = session_ctx.runtime_env(); let store = runtime_env.object_store(&object_store_url)?; let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); - let scan_io_source = if object_store_url == ObjectStoreUrl::local_filesystem() { - ScanIoSource::Local - } else { - ScanIoSource::ObjectStore - }; + let scan_io_source = scan_io_source(&object_store_url, is_hdfs_object_store); let reader_factory = Arc::new(EagerPageIndexReaderFactory::new( store, metadata_cache, @@ -223,6 +220,19 @@ pub(crate) fn init_datasource_exec( Ok(data_source_exec) } +fn scan_io_source(object_store_url: &ObjectStoreUrl, is_hdfs_object_store: bool) -> ScanIoSource { + let store_url: &url::Url = object_store_url.as_ref(); + if is_hdfs_object_store { + return ScanIoSource::OtherObjectStore; + } + + match store_url.scheme() { + "file" => ScanIoSource::Local, + "s3" | "gs" | "az" | "abfs" | "abfss" | "http" | "https" => ScanIoSource::ObjectStore, + _ => ScanIoSource::OtherObjectStore, + } +} + #[allow(clippy::too_many_arguments)] fn get_options( session_timezone: &str, @@ -300,6 +310,7 @@ mod tests { use arrow::array::Int32Array; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; + use bytes::Bytes; use datafusion::datasource::physical_plan::parquet::metadata::{ CachedParquetMetaData, DFParquetMetadata, }; @@ -311,6 +322,8 @@ mod tests { use datafusion::physical_plan::ExecutionPlan; use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename; use futures::StreamExt; + use object_store::memory::InMemory; + use object_store::{ObjectStore, ObjectStoreExt}; use parquet::arrow::ArrowWriter; use parquet::file::metadata::PageIndexPolicy; use parquet::file::properties::{EnabledStatistics, WriterProperties}; @@ -371,6 +384,7 @@ mod tests { Some(data_schema), None, ObjectStoreUrl::local_filesystem(), + false, vec![vec![partitioned_file]], projection, filters, @@ -456,6 +470,35 @@ mod tests { assert_eq!(global.coerce_int96_tz, Some("UTC".to_string())); } + #[test] + fn preserves_custom_hdfs_backend_range_reads_for_cloud_schemes() { + let options = HashMap::from([( + "fs.comet.libhdfs.schemes".to_string(), + "abfs,s3".to_string(), + )]); + + for scheme in ["abfs", "s3"] { + let url = ObjectStoreUrl::parse(format!("{scheme}://bucket")).unwrap(); + let original_url = url::Url::parse(format!("{scheme}://bucket").as_str()).unwrap(); + let is_hdfs_store = + crate::parquet::parquet_support::is_hdfs_scheme(&original_url, &options); + assert_eq!(scan_io_source(&url, false), ScanIoSource::ObjectStore); + assert_eq!( + scan_io_source(&url, is_hdfs_store), + ScanIoSource::OtherObjectStore + ); + } + + let s3a_original_url = url::Url::parse("s3a://bucket").unwrap(); + let normalized_s3_url = ObjectStoreUrl::parse("s3://bucket").unwrap(); + let is_hdfs_store = + crate::parquet::parquet_support::is_hdfs_scheme(&s3a_original_url, &options); + assert_eq!( + scan_io_source(&normalized_s3_url, is_hdfs_store), + ScanIoSource::ObjectStore + ); + } + // Regression test for #3978: DataFusion's opener requests `PageIndexPolicy::Skip` on the // initial metadata load and only loads the page index later, on demand, when row-group // pruning shows it is still needed (apache/datafusion#22857). That on-demand load bypasses @@ -499,6 +542,7 @@ mod tests { None, None, ObjectStoreUrl::local_filesystem(), + false, vec![vec![partitioned_file]], None, None, @@ -545,6 +589,13 @@ mod tests { #[tokio::test] async fn reports_cold_and_warm_metadata_io_without_data_reads() { let (filename, _schema) = write_scan_io_fixture(); + let file_bytes = std::fs::read(&filename).unwrap(); + let footer_bytes = usize::try_from(u32::from_le_bytes( + file_bytes[file_bytes.len() - 8..file_bytes.len() - 4] + .try_into() + .unwrap(), + )) + .unwrap(); let partitioned_file = PartitionedFile::from_path(filename).unwrap(); let session_ctx = Arc::new(SessionContext::new()); let runtime_env = session_ctx.runtime_env(); @@ -564,34 +615,24 @@ mod tests { .unwrap(); cold_reader.get_metadata(None).await.unwrap(); - let cold_metadata_requested = - reader_metric(&cold_metrics, "scan_io_metadata_bytes_requested"); - let cold_metadata_returned = - reader_metric(&cold_metrics, "scan_io_metadata_bytes_returned"); - assert!(cold_metadata_requested > 0); - assert_eq!(cold_metadata_returned, cold_metadata_requested); + let cold_metadata_bytes = reader_metric(&cold_metrics, "scan_io_metadata_bytes"); + assert!(cold_metadata_bytes > footer_bytes); + assert_eq!(reader_metric(&cold_metrics, "scan_io_data_bytes"), 0); + assert_eq!(reader_metric(&cold_metrics, "scan_io_footer_reads"), 1); assert_eq!( - reader_metric(&cold_metrics, "scan_io_data_bytes_requested"), - 0 + reader_metric(&cold_metrics, "scan_io_footer_bytes"), + footer_bytes ); assert_eq!( - reader_metric(&cold_metrics, "scan_io_bytes_requested"), - cold_metadata_requested - ); - assert_eq!( - reader_metric(&cold_metrics, "scan_io_local_bytes_requested"), - cold_metadata_requested - ); - assert_eq!( - reader_metric(&cold_metrics, "scan_io_local_bytes_returned"), - cold_metadata_returned + reader_metric(&cold_metrics, "scan_io_object_store_get_calls"), + 0 ); assert_eq!( - reader_metric(&cold_metrics, "scan_io_object_store_bytes_requested"), + reader_metric(&cold_metrics, "scan_io_object_store_get_requested_bytes"), 0 ); assert_eq!( - reader_metric(&cold_metrics, "scan_io_object_store_bytes_returned"), + reader_metric(&cold_metrics, "scan_io_object_store_response_bytes_read"), 0 ); assert_eq!( @@ -615,15 +656,9 @@ mod tests { .unwrap(); warm_reader.get_metadata(None).await.unwrap(); - assert_eq!( - reader_metric(&warm_metrics, "scan_io_metadata_bytes_requested"), - 0 - ); - assert_eq!( - reader_metric(&warm_metrics, "scan_io_metadata_bytes_returned"), - 0 - ); - assert_eq!(reader_metric(&warm_metrics, "scan_io_bytes_requested"), 0); + assert_eq!(reader_metric(&warm_metrics, "scan_io_metadata_bytes"), 0); + assert_eq!(reader_metric(&warm_metrics, "scan_io_footer_reads"), 0); + assert_eq!(reader_metric(&warm_metrics, "scan_io_footer_bytes"), 0); assert_eq!( reader_metric(&warm_metrics, "scan_io_metadata_cache_hits"), 1 @@ -669,13 +704,11 @@ mod tests { assert!(metadata.column_index().is_some()); assert!(metadata.offset_index().is_some()); - let requested = reader_metric(&metrics, "scan_io_metadata_bytes_requested"); - assert!(requested > 0); - assert_eq!( - reader_metric(&metrics, "scan_io_metadata_bytes_returned"), - requested - ); - assert_eq!(reader_metric(&metrics, "scan_io_data_bytes_requested"), 0); + let metadata_bytes = reader_metric(&metrics, "scan_io_metadata_bytes"); + assert!(metadata_bytes > 0); + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_footer_bytes"), 0); assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_hits"), 0); assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_misses"), 1); @@ -687,15 +720,10 @@ mod tests { .await .unwrap(); assert_eq!( - reader_metric(&metrics, "scan_io_metadata_bytes_requested"), - requested + page_index_length as usize - ); - assert_eq!( - reader_metric(&metrics, "scan_io_metadata_bytes_returned"), - requested + page_index_length as usize + reader_metric(&metrics, "scan_io_metadata_bytes"), + metadata_bytes + page_index_length as usize ); - assert_eq!(reader_metric(&metrics, "scan_io_data_bytes_requested"), 0); - assert_eq!(reader_metric(&metrics, "scan_io_data_bytes_returned"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes"), 0); } #[test] @@ -723,13 +751,158 @@ mod tests { .iter() .filter(|metric| metric.value().name().starts_with("scan_io_")) .collect::>(); - assert_eq!(scan_io_metrics.len(), 12); + assert_eq!(scan_io_metrics.len(), 9); for metric in scan_io_metrics { assert!(metric.labels().is_empty()); assert_eq!(metric.partition(), None); } } + #[tokio::test] + async fn reports_object_store_read_amplification_after_range_coalescing() { + let file_size = 524_416; + let location = object_store::path::Path::from("coalesced.parquet"); + let store: Arc = Arc::new(InMemory::new()); + store + .put(&location, Bytes::from(vec![0; file_size]).into()) + .await + .unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + let metadata_cache = session_ctx + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = EagerPageIndexReaderFactory::new( + store, + metadata_cache, + ScanIoSource::ObjectStore, + &metrics, + ); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::new(location.to_string(), file_size as u64), + None, + &metrics, + ) + .unwrap(); + let buffers = reader + .get_byte_ranges(vec![0..64, 524_352..524_416]) + .await + .unwrap(); + + assert_eq!(buffers.iter().map(Bytes::len).sum::(), 128); + assert_eq!(reader_metric(&metrics, "bytes_scanned"), 128); + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes"), 128); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_bytes"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_object_store_get_calls"), 1); + assert_eq!( + reader_metric(&metrics, "scan_io_object_store_get_requested_bytes"), + file_size + ); + assert_eq!( + reader_metric(&metrics, "scan_io_object_store_response_bytes_read"), + file_size + ); + } + + #[tokio::test] + async fn reports_object_store_footer_and_metadata_reads() { + let (filename, _schema) = write_scan_io_fixture(); + let file_bytes = std::fs::read(filename).unwrap(); + let file_size = file_bytes.len(); + let footer_bytes = usize::try_from(u32::from_le_bytes( + file_bytes[file_size - 8..file_size - 4].try_into().unwrap(), + )) + .unwrap(); + let location = object_store::path::Path::from("footer.parquet"); + let store: Arc = Arc::new(InMemory::new()); + store + .put(&location, Bytes::from(file_bytes).into()) + .await + .unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + let metadata_cache = session_ctx + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = EagerPageIndexReaderFactory::new( + store, + metadata_cache, + ScanIoSource::ObjectStore, + &metrics, + ); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::new(location.to_string(), file_size as u64), + Some(512 * 1024), + &metrics, + ) + .unwrap(); + reader.get_metadata(None).await.unwrap(); + + assert_eq!(reader_metric(&metrics, "scan_io_data_bytes"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_bytes"), file_size); + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 1); + assert_eq!( + reader_metric(&metrics, "scan_io_footer_bytes"), + footer_bytes + ); + assert_eq!(reader_metric(&metrics, "scan_io_object_store_get_calls"), 1); + assert_eq!( + reader_metric(&metrics, "scan_io_object_store_get_requested_bytes"), + file_size + ); + assert_eq!( + reader_metric(&metrics, "scan_io_object_store_response_bytes_read"), + file_size + ); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_misses"), 1); + } + + #[tokio::test] + async fn does_not_report_footer_metrics_when_metadata_decoding_fails() { + let (filename, _schema) = write_scan_io_fixture(); + let mut file_bytes = std::fs::read(&filename).unwrap(); + let file_size = file_bytes.len(); + let footer_bytes = usize::try_from(u32::from_le_bytes( + file_bytes[file_size - 8..file_size - 4].try_into().unwrap(), + )) + .unwrap(); + file_bytes[file_size - 8 - footer_bytes..file_size - 8].fill(0xff); + std::fs::write(&filename, file_bytes).unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + let runtime_env = session_ctx.runtime_env(); + let store = runtime_env + .object_store(ObjectStoreUrl::local_filesystem()) + .unwrap(); + let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = + EagerPageIndexReaderFactory::new(store, metadata_cache, ScanIoSource::Local, &metrics); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::from_path(filename).unwrap(), + Some(512 * 1024), + &metrics, + ) + .unwrap(); + + assert!(reader.get_metadata(None).await.is_err()); + assert!(reader_metric(&metrics, "scan_io_metadata_bytes") > 0); + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_footer_bytes"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_hits"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_misses"), 0); + } + #[tokio::test] async fn classifies_bloom_filter_reads_as_metadata() { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -774,13 +947,8 @@ mod tests { let bloom_bytes = scan_metric(&scan, "bytes_scanned"); assert!(bloom_bytes > 0); - assert_eq!(scan_metric(&scan, "scan_io_data_bytes_requested"), 0); - assert_eq!(scan_metric(&scan, "scan_io_data_bytes_returned"), 0); - assert!(scan_metric(&scan, "scan_io_metadata_bytes_requested") >= bloom_bytes); - assert_eq!( - scan_metric(&scan, "scan_io_metadata_bytes_requested"), - scan_metric(&scan, "scan_io_metadata_bytes_returned") - ); + assert_eq!(scan_metric(&scan, "scan_io_data_bytes"), 0); + assert!(scan_metric(&scan, "scan_io_metadata_bytes") >= bloom_bytes); } #[tokio::test] @@ -825,9 +993,9 @@ mod tests { ); drain_scan(&filtered_scan, &session_ctx).await; - let full_data = scan_metric(&full_scan, "scan_io_data_bytes_requested"); - let projected_data = scan_metric(&projected_scan, "scan_io_data_bytes_requested"); - let filtered_data = scan_metric(&filtered_scan, "scan_io_data_bytes_requested"); + let full_data = scan_metric(&full_scan, "scan_io_data_bytes"); + let projected_data = scan_metric(&projected_scan, "scan_io_data_bytes"); + let filtered_data = scan_metric(&filtered_scan, "scan_io_data_bytes"); assert!( projected_data < full_data, "projection should request fewer data bytes: projected={projected_data}, full={full_data}" @@ -838,31 +1006,19 @@ mod tests { ); for scan in [&full_scan, &projected_scan, &filtered_scan] { - let data_requested = scan_metric(scan, "scan_io_data_bytes_requested"); - let data_returned = scan_metric(scan, "scan_io_data_bytes_returned"); - let metadata_requested = scan_metric(scan, "scan_io_metadata_bytes_requested"); - let metadata_returned = scan_metric(scan, "scan_io_metadata_bytes_returned"); - assert_eq!(scan_metric(scan, "bytes_scanned"), data_requested); - assert_eq!(data_returned, data_requested); - assert_eq!(metadata_returned, metadata_requested); - assert_eq!( - scan_metric(scan, "scan_io_bytes_requested"), - data_requested + metadata_requested - ); assert_eq!( - scan_metric(scan, "scan_io_bytes_returned"), - data_returned + metadata_returned + scan_metric(scan, "bytes_scanned"), + scan_metric(scan, "scan_io_data_bytes") ); + assert_eq!(scan_metric(scan, "scan_io_object_store_get_calls"), 0); assert_eq!( - scan_metric(scan, "scan_io_local_bytes_requested"), - data_requested + metadata_requested + scan_metric(scan, "scan_io_object_store_get_requested_bytes"), + 0 ); assert_eq!( - scan_metric(scan, "scan_io_local_bytes_returned"), - data_returned + metadata_returned + scan_metric(scan, "scan_io_object_store_response_bytes_read"), + 0 ); - assert_eq!(scan_metric(scan, "scan_io_object_store_bytes_requested"), 0); - assert_eq!(scan_metric(scan, "scan_io_object_store_bytes_returned"), 0); } } } diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 5b22afa2609..585a894eeb5 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -469,9 +469,9 @@ fn create_hdfs_object_store( }) } -type ObjectStoreCache = RwLock>>; +type ObjectStoreCache = RwLock>>; -/// Process-wide cache of object stores, keyed by `(scheme://host:port, config_hash)`. +/// Process-wide cache of object stores, keyed by `(scheme://host:port, config_hash, hdfs_backend)`. /// /// ## Why static / process lifetime? /// @@ -486,8 +486,8 @@ type ObjectStoreCache = RwLock>>; /// /// ## Unbounded size /// -/// Cache entries are indexed by `(scheme://host:port, hash-of-configs)`. A typical Spark -/// job accesses a small, fixed set of buckets with a stable configuration, so the number of +/// Cache entries are indexed by `(scheme://host:port, hash-of-configs, hdfs_backend)`. A typical +/// Spark job accesses a small, fixed set of buckets with a stable configuration, so the number of /// distinct keys is O(buckets × credential-configs) and remains small throughout the job. /// Entries are cheap relative to the cost of creating a new object store (new HTTP /// connection pool + DNS resolution), and there is no meaningful benefit from eviction, so @@ -543,7 +543,7 @@ pub(crate) fn prepare_object_store_with_configs( ); let config_hash = hash_object_store_configs(object_store_configs); - let cache_key = (url_key.clone(), config_hash); + let cache_key = (url_key.clone(), config_hash, is_hdfs_scheme); // Check the cache first to reuse existing object store instances. // This enables HTTP connection pooling and avoids redundant DNS lookups. @@ -604,6 +604,47 @@ mod tests { #[cfg(not(feature = "hdfs-opendal"))] use std::collections::HashMap; + #[test] + fn cache_distinguishes_backend_routing_for_normalized_s3_aliases() { + let configs = std::collections::HashMap::from([( + "fs.comet.libhdfs.schemes".to_string(), + "s3".to_string(), + )]); + let cloud_url = url::Url::parse("s3a://scan-io-backend-routing").unwrap(); + let hdfs_url = url::Url::parse("s3://scan-io-backend-routing").unwrap(); + let config_hash = super::hash_object_store_configs(&configs); + let cloud_key = ( + "s3://scan-io-backend-routing".to_string(), + config_hash, + super::is_hdfs_scheme(&cloud_url, &configs), + ); + let hdfs_key = ( + "s3://scan-io-backend-routing".to_string(), + config_hash, + super::is_hdfs_scheme(&hdfs_url, &configs), + ); + let cloud_store: std::sync::Arc = + std::sync::Arc::new(object_store::memory::InMemory::new()); + let hdfs_store: std::sync::Arc = + std::sync::Arc::new(object_store::memory::InMemory::new()); + + let mut cache = super::object_store_cache().write().unwrap(); + cache.insert(cloud_key.clone(), std::sync::Arc::clone(&cloud_store)); + cache.insert(hdfs_key.clone(), std::sync::Arc::clone(&hdfs_store)); + + assert!(std::sync::Arc::ptr_eq( + cache.get(&cloud_key).unwrap(), + &cloud_store + )); + assert!(std::sync::Arc::ptr_eq( + cache.get(&hdfs_key).unwrap(), + &hdfs_store + )); + + cache.remove(&cloud_key); + cache.remove(&hdfs_key); + } + /// Parses the url, registers the object store, and returns a tuple of the object store url and object store path #[cfg(not(feature = "hdfs-opendal"))] pub(crate) fn prepare_object_store( diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 7e759bfca05..d26ca6fb072 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -282,38 +282,26 @@ object CometMetricNode { SQLMetrics.createMetric(sc, "Number of row groups matched by limit pruning (not pruned)"), "bytes_scanned" -> SQLMetrics.createSizeMetric(sc, "Number of bytes scanned"), - "scan_io_bytes_requested" -> + "scan_io_data_bytes" -> SQLMetrics.createSizeMetric( sc, - "Total data-page plus metadata byte ranges requested by native Parquet scan"), - "scan_io_bytes_returned" -> + "Projected Parquet data-page bytes returned to the reader"), + "scan_io_metadata_bytes" -> SQLMetrics.createSizeMetric( sc, - "Total data-page plus metadata bytes returned to native Parquet scan"), - "scan_io_data_bytes_requested" -> - SQLMetrics.createSizeMetric(sc, "Data-page byte ranges requested by native Parquet scan"), - "scan_io_data_bytes_returned" -> - SQLMetrics.createSizeMetric(sc, "Data-page bytes returned to native Parquet scan"), - "scan_io_metadata_bytes_requested" -> + "Footer-prefetch, page-index, and Bloom-filter bytes returned to the reader"), + "scan_io_footer_reads" -> + SQLMetrics.createMetric(sc, "Number of Parquet footer payloads read from storage"), + "scan_io_footer_bytes" -> + SQLMetrics.createSizeMetric(sc, "Serialized Parquet footer payload bytes read"), + "scan_io_object_store_get_calls" -> + SQLMetrics.createMetric(sc, "ObjectStore GET operations after range coalescing"), + "scan_io_object_store_get_requested_bytes" -> + SQLMetrics.createSizeMetric(sc, "ObjectStore GET range bytes after range coalescing"), + "scan_io_object_store_response_bytes_read" -> SQLMetrics.createSizeMetric( sc, - "Footer, page-index, and Bloom-filter byte ranges requested through the storage API"), - "scan_io_metadata_bytes_returned" -> - SQLMetrics.createSizeMetric( - sc, - "Footer, page-index, and Bloom-filter bytes returned through the storage API"), - "scan_io_object_store_bytes_requested" -> - SQLMetrics.createSizeMetric( - sc, - "Non-local ObjectStore API bytes requested, excluding metadata cache hits"), - "scan_io_object_store_bytes_returned" -> - SQLMetrics.createSizeMetric( - sc, - "Non-local ObjectStore API bytes returned, excluding metadata cache hits"), - "scan_io_local_bytes_requested" -> - SQLMetrics.createSizeMetric(sc, "Local file bytes requested by native Parquet scan"), - "scan_io_local_bytes_returned" -> - SQLMetrics.createSizeMetric(sc, "Local file bytes returned to native Parquet scan"), + "ObjectStore response bytes consumed after range coalescing"), "scan_io_metadata_cache_hits" -> SQLMetrics.createMetric(sc, "Metadata loads served without storage I/O"), "scan_io_metadata_cache_misses" -> diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 035a6d1af77..bd8ae540929 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -2418,16 +2418,13 @@ class CometExecSuite extends CometTestBase { assert(metrics.contains("time_elapsed_processing")) assert(metrics.contains("time_elapsed_scanning_until_data")) Seq( - "scan_io_bytes_requested", - "scan_io_bytes_returned", - "scan_io_data_bytes_requested", - "scan_io_data_bytes_returned", - "scan_io_metadata_bytes_requested", - "scan_io_metadata_bytes_returned", - "scan_io_object_store_bytes_requested", - "scan_io_object_store_bytes_returned", - "scan_io_local_bytes_requested", - "scan_io_local_bytes_returned", + "scan_io_data_bytes", + "scan_io_metadata_bytes", + "scan_io_footer_reads", + "scan_io_footer_bytes", + "scan_io_object_store_get_calls", + "scan_io_object_store_get_requested_bytes", + "scan_io_object_store_response_bytes_read", "scan_io_metadata_cache_hits", "scan_io_metadata_cache_misses").foreach { name => assert(metrics.contains(name), s"Missing $name. Available: ${metrics.keys}") @@ -2436,22 +2433,13 @@ class CometExecSuite extends CometTestBase { metrics("time_elapsed_scanning_total").value > 0, "time_elapsed_scanning_total should be > 0") assert(metrics("bytes_scanned").value > 0, "bytes_scanned should be > 0") - assert(metrics("scan_io_data_bytes_requested").value > 0) - assert(metrics("scan_io_data_bytes_returned").value > 0) - assert(metrics("scan_io_metadata_bytes_requested").value > 0) - assert(metrics("scan_io_metadata_bytes_returned").value > 0) - assert( - metrics("scan_io_bytes_requested").value == - metrics("scan_io_data_bytes_requested").value + - metrics("scan_io_metadata_bytes_requested").value) - assert( - metrics("scan_io_bytes_returned").value == - metrics("scan_io_data_bytes_returned").value + - metrics("scan_io_metadata_bytes_returned").value) - assert(metrics("scan_io_local_bytes_requested").value > 0) - assert(metrics("scan_io_local_bytes_returned").value > 0) - assert(metrics("scan_io_object_store_bytes_requested").value == 0) - assert(metrics("scan_io_object_store_bytes_returned").value == 0) + assert(metrics("scan_io_data_bytes").value > 0) + assert(metrics("scan_io_metadata_bytes").value > 0) + assert(metrics("scan_io_footer_reads").value > 0) + assert(metrics("scan_io_footer_bytes").value > 0) + assert(metrics("scan_io_object_store_get_calls").value == 0) + assert(metrics("scan_io_object_store_get_requested_bytes").value == 0) + assert(metrics("scan_io_object_store_response_bytes_read").value == 0) assert(metrics("scan_io_metadata_cache_misses").value > 0) assert(metrics("output_rows").value > 0, "output_rows should be > 0") assert(metrics("time_elapsed_opening").value > 0, "time_elapsed_opening should be > 0") From f6c6100d497398ee44748064f1ceaa31f13f2d70 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 22 Aug 2026 23:00:00 -0700 Subject: [PATCH 06/12] Fix native scan metric shutdown and object-store registry isolation --- native/core/src/execution/jni_api.rs | 37 +++++++++++++++++- native/core/src/parquet/parquet_exec.rs | 9 ++++- native/core/src/parquet/parquet_support.rs | 45 ++++++++++++++++++---- 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d80754736b7..048d7c63ff2 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -95,6 +95,7 @@ use std::time::{Duration, Instant}; use std::{sync::Arc, task::Poll}; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; +use tokio::task::JoinHandle; use crate::execution::memory_pools::{ create_memory_pool, handle_task_shared_pool_release, parse_memory_pool_config, MemoryPoolConfig, @@ -324,6 +325,7 @@ struct ExecutionContext { pub stream: Option, /// Receives batches from a spawned tokio task (async I/O path) pub batch_receiver: Option>>, + pub batch_producer: Option>, /// Native metrics pub metrics: Arc>>, // The interval in milliseconds to update metrics @@ -537,6 +539,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources, stream: None, batch_receiver: None, + batch_producer: None, metrics, metrics_update_interval, metrics_last_update_time: Instant::now(), @@ -835,7 +838,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // decreasing to 1 would serialize production and consumption. let (tx, rx) = mpsc::channel(2); let mut stream = stream; - get_runtime().spawn(async move { + let producer = get_runtime().spawn(async move { let result = std::panic::AssertUnwindSafe(async { while let Some(batch) = stream.next().await { if tx.send(batch).await.is_err() { @@ -862,6 +865,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } }); exec_context.batch_receiver = Some(rx); + exec_context.batch_producer = Some(producer); } else { exec_context.stream = Some(stream); } @@ -966,6 +970,11 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( try_unwrap_or_throw(&e, |env| unsafe { let execution_context = get_execution_context(exec_context); + execution_context.batch_receiver.take(); + if let Some(producer) = execution_context.batch_producer.take() { + stop_batch_producer(producer); + } + // Update metrics update_metrics(env, execution_context)?; @@ -989,6 +998,11 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( }) } +fn stop_batch_producer(producer: JoinHandle<()>) { + producer.abort(); + let _ = get_runtime().block_on(producer); +} + fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { if let Some(native_query) = &exec_context.root_op { let metrics = exec_context.metrics.as_obj(); @@ -1375,3 +1389,24 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( Ok(()) }) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn waits_for_background_batch_producer_shutdown() { + let (sender, receiver) = std::sync::mpsc::channel::<()>(); + let producer = get_runtime().spawn(async move { + let _sender = sender; + futures::future::pending::<()>().await; + }); + + stop_batch_producer(producer); + + assert_eq!( + receiver.try_recv(), + Err(std::sync::mpsc::TryRecvError::Disconnected) + ); + } +} diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 055b216fc83..e609efd5a66 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -228,7 +228,9 @@ fn scan_io_source(object_store_url: &ObjectStoreUrl, is_hdfs_object_store: bool) match store_url.scheme() { "file" => ScanIoSource::Local, - "s3" | "gs" | "az" | "abfs" | "abfss" | "http" | "https" => ScanIoSource::ObjectStore, + "s3" | "s3a" | "gs" | "az" | "abfs" | "abfss" | "http" | "https" => { + ScanIoSource::ObjectStore + } _ => ScanIoSource::OtherObjectStore, } } @@ -491,12 +493,17 @@ mod tests { let s3a_original_url = url::Url::parse("s3a://bucket").unwrap(); let normalized_s3_url = ObjectStoreUrl::parse("s3://bucket").unwrap(); + let preserved_s3a_url = ObjectStoreUrl::parse("s3a://bucket").unwrap(); let is_hdfs_store = crate::parquet::parquet_support::is_hdfs_scheme(&s3a_original_url, &options); assert_eq!( scan_io_source(&normalized_s3_url, is_hdfs_store), ScanIoSource::ObjectStore ); + assert_eq!( + scan_io_source(&preserved_s3a_url, is_hdfs_store), + ScanIoSource::ObjectStore + ); } // Regression test for #3978: DataFusion's opener requests `PageIndexPolicy::Skip` on the diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 585a894eeb5..265b378b78c 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -528,6 +528,7 @@ pub(crate) fn prepare_object_store_with_configs( ) -> Result<(ObjectStoreUrl, Path), ExecutionError> { let mut url = Url::parse(url.as_str()) .map_err(|e| ExecutionError::GeneralError(format!("Error parsing URL {url}: {e}")))?; + let original_url = url.clone(); let is_hdfs_scheme = is_hdfs_scheme(&url, object_store_configs); let mut scheme = url.scheme(); if !is_hdfs_scheme && scheme == "s3a" { @@ -581,8 +582,19 @@ pub(crate) fn prepare_object_store_with_configs( (store, path) }; - let object_store_url = ObjectStoreUrl::parse(url_key.clone())?; - runtime_env.register_object_store(&url, object_store); + let registration_url = if original_url.scheme() != url.scheme() + && self::is_hdfs_scheme(&url, object_store_configs) + { + &original_url + } else { + &url + }; + let object_store_url = ObjectStoreUrl::parse(format!( + "{}://{}", + registration_url.scheme(), + ®istration_url[url::Position::BeforeHost..url::Position::AfterPort], + ))?; + runtime_env.register_object_store(registration_url, object_store); Ok((object_store_url, object_store_path)) } @@ -628,19 +640,38 @@ mod tests { let hdfs_store: std::sync::Arc = std::sync::Arc::new(object_store::memory::InMemory::new()); - let mut cache = super::object_store_cache().write().unwrap(); - cache.insert(cloud_key.clone(), std::sync::Arc::clone(&cloud_store)); - cache.insert(hdfs_key.clone(), std::sync::Arc::clone(&hdfs_store)); + { + let mut cache = super::object_store_cache().write().unwrap(); + cache.insert(cloud_key.clone(), std::sync::Arc::clone(&cloud_store)); + cache.insert(hdfs_key.clone(), std::sync::Arc::clone(&hdfs_store)); + } + let runtime_env = + std::sync::Arc::new(datafusion::execution::runtime_env::RuntimeEnv::default()); + let (cloud_object_store_url, _) = super::prepare_object_store_with_configs( + std::sync::Arc::clone(&runtime_env), + cloud_url.to_string(), + &configs, + ) + .unwrap(); + let (hdfs_object_store_url, _) = super::prepare_object_store_with_configs( + std::sync::Arc::clone(&runtime_env), + hdfs_url.to_string(), + &configs, + ) + .unwrap(); + + assert_ne!(cloud_object_store_url, hdfs_object_store_url); assert!(std::sync::Arc::ptr_eq( - cache.get(&cloud_key).unwrap(), + &runtime_env.object_store(&cloud_object_store_url).unwrap(), &cloud_store )); assert!(std::sync::Arc::ptr_eq( - cache.get(&hdfs_key).unwrap(), + &runtime_env.object_store(&hdfs_object_store_url).unwrap(), &hdfs_store )); + let mut cache = super::object_store_cache().write().unwrap(); cache.remove(&cloud_key); cache.remove(&hdfs_key); } From 1d87afee03270b10399197267ed77e60c5b99455 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 23 Aug 2026 13:28:33 -0700 Subject: [PATCH 07/12] Fix native scan producer cleanup and object-store isolation --- native/core/src/execution/jni_api.rs | 31 ++++++++- native/core/src/parquet/parquet_exec.rs | 47 ++++++++++++- native/core/src/parquet/parquet_support.rs | 81 ++++++++++++++++++++-- 3 files changed, 150 insertions(+), 9 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 048d7c63ff2..4b1a9d18aaf 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1000,7 +1000,14 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( fn stop_batch_producer(producer: JoinHandle<()>) { producer.abort(); - let _ = get_runtime().block_on(producer); + let deadline = Instant::now() + Duration::from_millis(100); + while !producer.is_finished() { + if Instant::now() >= deadline { + return; + } + std::thread::sleep(Duration::from_millis(1)); + } + let _ = futures::executor::block_on(producer); } fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { @@ -1409,4 +1416,26 @@ mod tests { Err(std::sync::mpsc::TryRecvError::Disconnected) ); } + + #[test] + fn does_not_wait_indefinitely_for_blocked_batch_producer() { + let (started_sender, started_receiver) = std::sync::mpsc::channel(); + let (release_sender, release_receiver) = std::sync::mpsc::channel(); + let producer = get_runtime().spawn(async move { + started_sender.send(()).unwrap(); + release_receiver.recv().unwrap(); + }); + started_receiver.recv().unwrap(); + + let release_thread = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(500)); + release_sender.send(()).unwrap(); + }); + let started = Instant::now(); + stop_batch_producer(producer); + let elapsed = started.elapsed(); + release_thread.join().unwrap(); + + assert!(elapsed < Duration::from_millis(400)); + } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index e609efd5a66..b88419b7bf5 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -221,12 +221,11 @@ pub(crate) fn init_datasource_exec( } fn scan_io_source(object_store_url: &ObjectStoreUrl, is_hdfs_object_store: bool) -> ScanIoSource { - let store_url: &url::Url = object_store_url.as_ref(); if is_hdfs_object_store { return ScanIoSource::OtherObjectStore; } - match store_url.scheme() { + match physical_object_store_scheme(object_store_url) { "file" => ScanIoSource::Local, "s3" | "s3a" | "gs" | "az" | "abfs" | "abfss" | "http" | "https" => { ScanIoSource::ObjectStore @@ -235,6 +234,14 @@ fn scan_io_source(object_store_url: &ObjectStoreUrl, is_hdfs_object_store: bool) } } +fn physical_object_store_scheme(object_store_url: &ObjectStoreUrl) -> &str { + let store_url: &url::Url = object_store_url.as_ref(); + store_url + .scheme() + .split_once("+comet-") + .map_or_else(|| store_url.scheme(), |(scheme, _)| scheme) +} + #[allow(clippy::too_many_arguments)] fn get_options( session_timezone: &str, @@ -295,10 +302,15 @@ fn get_options( spark_parquet_options.allow_timestamp_ltz_to_ntz = allow_timestamp_ltz_to_ntz; if encryption_enabled { + let store_url: &url::Url = object_store_url.as_ref(); table_parquet_options.crypto.configure_factory( ENCRYPTION_FACTORY_ID, &CometEncryptionConfig { - uri_base: object_store_url.to_string(), + uri_base: format!( + "{}://{}/", + physical_object_store_scheme(object_store_url), + &store_url[url::Position::BeforeHost..url::Position::AfterPort], + ), }, ); } @@ -494,6 +506,8 @@ mod tests { let s3a_original_url = url::Url::parse("s3a://bucket").unwrap(); let normalized_s3_url = ObjectStoreUrl::parse("s3://bucket").unwrap(); let preserved_s3a_url = ObjectStoreUrl::parse("s3a://bucket").unwrap(); + let isolated_s3a_url = + ObjectStoreUrl::parse("s3a+comet-0123456789abcdef-native://bucket").unwrap(); let is_hdfs_store = crate::parquet::parquet_support::is_hdfs_scheme(&s3a_original_url, &options); assert_eq!( @@ -504,6 +518,33 @@ mod tests { scan_io_source(&preserved_s3a_url, is_hdfs_store), ScanIoSource::ObjectStore ); + assert_eq!( + scan_io_source(&isolated_s3a_url, is_hdfs_store), + ScanIoSource::ObjectStore + ); + } + + #[test] + fn preserves_physical_uri_for_isolated_encrypted_object_stores() { + let object_store_url = + ObjectStoreUrl::parse("s3a+comet-0123456789abcdef-native://bucket").unwrap(); + let (table_parquet_options, _) = get_options( + "UTC", + true, + false, + false, + false, + &object_store_url, + true, + &ParquetOptions::default(), + ); + let encryption_options: CometEncryptionConfig = table_parquet_options + .crypto + .factory_options + .to_extension_options() + .unwrap(); + + assert_eq!(encryption_options.uri_base, "s3a://bucket/"); } // Regression test for #3978: DataFusion's opener requests `PageIndexPolicy::Skip` on the diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 265b378b78c..2dc884bceb0 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -582,19 +582,28 @@ pub(crate) fn prepare_object_store_with_configs( (store, path) }; - let registration_url = if original_url.scheme() != url.scheme() - && self::is_hdfs_scheme(&url, object_store_configs) + let object_store_url = ObjectStoreUrl::parse(url_key.as_str())?; + let registration_url = if scheme != "file" + && runtime_env + .object_store(&object_store_url) + .is_ok_and(|existing| !Arc::ptr_eq(&existing, &object_store)) { - &original_url + let backend = if is_hdfs_scheme { "hdfs" } else { "native" }; + Url::parse(&format!( + "{}+comet-{config_hash:016x}-{backend}://{}", + original_url.scheme(), + &url[url::Position::BeforeHost..url::Position::AfterPort], + )) + .map_err(|e| ExecutionError::GeneralError(e.to_string()))? } else { - &url + url }; let object_store_url = ObjectStoreUrl::parse(format!( "{}://{}", registration_url.scheme(), ®istration_url[url::Position::BeforeHost..url::Position::AfterPort], ))?; - runtime_env.register_object_store(registration_url, object_store); + runtime_env.register_object_store(®istration_url, object_store); Ok((object_store_url, object_store_path)) } @@ -676,6 +685,68 @@ mod tests { cache.remove(&hdfs_key); } + #[test] + fn registry_distinguishes_backend_routing_across_configurations() { + let cloud_configs = std::collections::HashMap::from([( + "fs.comet.libhdfs.schemes".to_string(), + "s3".to_string(), + )]); + let hdfs_configs = std::collections::HashMap::from([( + "fs.comet.libhdfs.schemes".to_string(), + "s3a".to_string(), + )]); + let url = url::Url::parse("s3a://scan-io-mixed-backend-routing").unwrap(); + let cloud_key = ( + "s3://scan-io-mixed-backend-routing".to_string(), + super::hash_object_store_configs(&cloud_configs), + false, + ); + let hdfs_key = ( + "s3a://scan-io-mixed-backend-routing".to_string(), + super::hash_object_store_configs(&hdfs_configs), + true, + ); + let cloud_store: std::sync::Arc = + std::sync::Arc::new(object_store::memory::InMemory::new()); + let hdfs_store: std::sync::Arc = + std::sync::Arc::new(object_store::memory::InMemory::new()); + + { + let mut cache = super::object_store_cache().write().unwrap(); + cache.insert(cloud_key.clone(), std::sync::Arc::clone(&cloud_store)); + cache.insert(hdfs_key.clone(), std::sync::Arc::clone(&hdfs_store)); + } + + let runtime_env = + std::sync::Arc::new(datafusion::execution::runtime_env::RuntimeEnv::default()); + let (hdfs_object_store_url, _) = super::prepare_object_store_with_configs( + std::sync::Arc::clone(&runtime_env), + url.to_string(), + &hdfs_configs, + ) + .unwrap(); + let (cloud_object_store_url, _) = super::prepare_object_store_with_configs( + std::sync::Arc::clone(&runtime_env), + url.to_string(), + &cloud_configs, + ) + .unwrap(); + + assert_ne!(cloud_object_store_url, hdfs_object_store_url); + assert!(std::sync::Arc::ptr_eq( + &runtime_env.object_store(&cloud_object_store_url).unwrap(), + &cloud_store + )); + assert!(std::sync::Arc::ptr_eq( + &runtime_env.object_store(&hdfs_object_store_url).unwrap(), + &hdfs_store + )); + + let mut cache = super::object_store_cache().write().unwrap(); + cache.remove(&cloud_key); + cache.remove(&hdfs_key); + } + /// Parses the url, registers the object store, and returns a tuple of the object store url and object store path #[cfg(not(feature = "hdfs-opendal"))] pub(crate) fn prepare_object_store( From e8852255e42c0e91e8e480c5ea5adb8f5624166f Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 23 Aug 2026 15:57:30 -0700 Subject: [PATCH 08/12] Fix encrypted footer metric publication before key retrieval --- .../eager_page_index_reader_factory.rs | 26 +++-- native/core/src/parquet/parquet_exec.rs | 95 +++++++++++++++++++ 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index bb2a18fa9fd..92f8deb4cad 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -287,6 +287,7 @@ impl AsyncFileReader for EagerPageIndexReader { storage_reads: Arc::clone(&metadata_storage_reads), footer_payload_bytes: Arc::clone(&footer_payload_bytes), file_size: object_meta.size, + record_footer_immediately: !cache_enabled, }, }; @@ -306,7 +307,7 @@ impl AsyncFileReader for EagerPageIndexReader { if metadata.is_ok() { let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); - if footer_bytes > 0 { + if footer_bytes > 0 && cache_enabled { scan_io_metrics.footer_reads.add(1); scan_io_metrics.footer_bytes.add(footer_bytes); } @@ -330,6 +331,7 @@ enum ScanIoStoreRole { storage_reads: Arc, footer_payload_bytes: Arc, file_size: u64, + record_footer_immediately: bool, }, } @@ -367,17 +369,27 @@ impl ScanIoObjectStore { ScanIoStoreRole::Metadata { footer_payload_bytes, file_size, + record_footer_immediately, .. } => { self.scan_io_metrics.metadata_bytes.add(bytes.len()); if range.is_some_and(|range| range.end == *file_size) && bytes.len() >= 8 { if let Ok(footer) = FooterTail::try_from(&bytes[bytes.len() - 8..]) { - let _ = footer_payload_bytes.compare_exchange( - 0, - footer.metadata_length(), - Ordering::Relaxed, - Ordering::Relaxed, - ); + if footer_payload_bytes + .compare_exchange( + 0, + footer.metadata_length(), + Ordering::Relaxed, + Ordering::Relaxed, + ) + .is_ok() + && *record_footer_immediately + { + self.scan_io_metrics.footer_reads.add(1); + self.scan_io_metrics + .footer_bytes + .add(footer.metadata_length()); + } } } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index b88419b7bf5..a2533cbcb71 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -338,7 +338,10 @@ mod tests { use futures::StreamExt; use object_store::memory::InMemory; use object_store::{ObjectStore, ObjectStoreExt}; + use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::ArrowWriter; + use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever}; + use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::file::metadata::PageIndexPolicy; use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::fs::File; @@ -913,6 +916,98 @@ mod tests { assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_misses"), 1); } + #[tokio::test] + async fn reports_encrypted_footer_before_key_retrieval() { + struct ObservingKeyRetriever { + key: Vec, + metrics: Arc, + footer_bytes: usize, + } + + impl KeyRetriever for ObservingKeyRetriever { + fn retrieve_key(&self, _key_metadata: &[u8]) -> parquet::errors::Result> { + assert_eq!(reader_metric(&self.metrics, "scan_io_footer_reads"), 1); + assert_eq!( + reader_metric(&self.metrics, "scan_io_footer_bytes"), + self.footer_bytes + ); + Ok(self.key.clone()) + } + } + + let key = b"0123456789012345".to_vec(); + let encryption = FileEncryptionProperties::builder(key.clone()) + .with_footer_key_metadata(b"footer".to_vec()) + .build() + .unwrap(); + let properties = WriterProperties::builder() + .with_file_encryption_properties(encryption) + .build(); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let filename = get_temp_filename(); + let mut writer = + ArrowWriter::try_new(File::create(&filename).unwrap(), schema, Some(properties)) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let file_bytes = std::fs::read(&filename).unwrap(); + let file_size = file_bytes.len(); + let footer_bytes = usize::try_from(u32::from_le_bytes( + file_bytes[file_size - 8..file_size - 4].try_into().unwrap(), + )) + .unwrap(); + let location = object_store::path::Path::from("encrypted-footer.parquet"); + let store: Arc = Arc::new(InMemory::new()); + store + .put(&location, Bytes::from(file_bytes).into()) + .await + .unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + let metadata_cache = session_ctx + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let metrics = Arc::new(ExecutionPlanMetricsSet::new()); + let factory = EagerPageIndexReaderFactory::new( + store, + metadata_cache, + ScanIoSource::ObjectStore, + &metrics, + ); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::new(location.to_string(), file_size as u64), + Some(512 * 1024), + &metrics, + ) + .unwrap(); + let decryption = + FileDecryptionProperties::with_key_retriever(Arc::new(ObservingKeyRetriever { + key, + metrics: Arc::clone(&metrics), + footer_bytes, + })) + .build() + .unwrap(); + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption); + + reader.get_metadata(Some(&options)).await.unwrap(); + + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 1); + assert_eq!( + reader_metric(&metrics, "scan_io_footer_bytes"), + footer_bytes + ); + } + #[tokio::test] async fn does_not_report_footer_metrics_when_metadata_decoding_fails() { let (filename, _schema) = write_scan_io_fixture(); From 4733fc3a15f6cef2886bcbcda446a8c6d21d1872 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 23 Aug 2026 17:17:20 -0700 Subject: [PATCH 09/12] Count encrypted footers only after complete payload reads --- .../eager_page_index_reader_factory.rs | 46 ++++++++++------ native/core/src/parquet/parquet_exec.rs | 55 ++++++++++++++++++- 2 files changed, 82 insertions(+), 19 deletions(-) diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 92f8deb4cad..22cb1f194e4 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -72,7 +72,7 @@ use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; use parquet::file::metadata::{FooterTail, PageIndexPolicy, ParquetMetaData}; use std::fmt::{Debug, Display, Formatter}; use std::ops::Range; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -286,6 +286,7 @@ impl AsyncFileReader for EagerPageIndexReader { role: ScanIoStoreRole::Metadata { storage_reads: Arc::clone(&metadata_storage_reads), footer_payload_bytes: Arc::clone(&footer_payload_bytes), + footer_recorded: AtomicBool::new(false), file_size: object_meta.size, record_footer_immediately: !cache_enabled, }, @@ -330,6 +331,7 @@ enum ScanIoStoreRole { Metadata { storage_reads: Arc, footer_payload_bytes: Arc, + footer_recorded: AtomicBool, file_size: u64, record_footer_immediately: bool, }, @@ -368,6 +370,7 @@ impl ScanIoObjectStore { .add(bytes.len()), ScanIoStoreRole::Metadata { footer_payload_bytes, + footer_recorded, file_size, record_footer_immediately, .. @@ -375,21 +378,32 @@ impl ScanIoObjectStore { self.scan_io_metrics.metadata_bytes.add(bytes.len()); if range.is_some_and(|range| range.end == *file_size) && bytes.len() >= 8 { if let Ok(footer) = FooterTail::try_from(&bytes[bytes.len() - 8..]) { - if footer_payload_bytes - .compare_exchange( - 0, - footer.metadata_length(), - Ordering::Relaxed, - Ordering::Relaxed, - ) - .is_ok() - && *record_footer_immediately - { - self.scan_io_metrics.footer_reads.add(1); - self.scan_io_metrics - .footer_bytes - .add(footer.metadata_length()); - } + let _ = footer_payload_bytes.compare_exchange( + 0, + footer.metadata_length(), + Ordering::Relaxed, + Ordering::Relaxed, + ); + } + } + + if *record_footer_immediately { + let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); + let footer_end = file_size.saturating_sub(8); + if footer_bytes > 0 + && footer_end + .checked_sub(footer_bytes as u64) + .is_some_and(|footer_start| { + range.is_some_and(|range| { + range.start <= footer_start + && range.start.saturating_add(bytes.len() as u64) + >= footer_end + }) + }) + && !footer_recorded.swap(true, Ordering::Relaxed) + { + self.scan_io_metrics.footer_reads.add(1); + self.scan_io_metrics.footer_bytes.add(footer_bytes); } } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index a2533cbcb71..5f147f48522 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -337,14 +337,16 @@ mod tests { use datafusion_comet_spark_expr::test_common::file_util::get_temp_filename; use futures::StreamExt; use object_store::memory::InMemory; + use object_store::throttle::{ThrottleConfig, ThrottledStore}; use object_store::{ObjectStore, ObjectStoreExt}; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::ArrowWriter; use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever}; use parquet::encryption::encrypt::FileEncryptionProperties; - use parquet::file::metadata::PageIndexPolicy; + use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::fs::File; + use std::time::Duration; fn write_scan_io_fixture() -> (String, SchemaRef) { let schema = Arc::new(Schema::new(vec![ @@ -918,6 +920,15 @@ mod tests { #[tokio::test] async fn reports_encrypted_footer_before_key_retrieval() { + assert_encrypted_footer_before_key_retrieval(0).await; + } + + #[tokio::test] + async fn does_not_report_encrypted_footer_before_payload_is_read() { + assert_encrypted_footer_before_key_retrieval(512 * 1024).await; + } + + async fn assert_encrypted_footer_before_key_retrieval(footer_padding: usize) { struct ObservingKeyRetriever { key: Vec, metrics: Arc, @@ -942,6 +953,12 @@ mod tests { .unwrap(); let properties = WriterProperties::builder() .with_file_encryption_properties(encryption) + .set_key_value_metadata((footer_padding > 0).then(|| { + vec![KeyValue::new( + "padding".to_string(), + "x".repeat(footer_padding), + )] + })) .build(); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); let batch = RecordBatch::try_new( @@ -963,7 +980,18 @@ mod tests { )) .unwrap(); let location = object_store::path::Path::from("encrypted-footer.parquet"); - let store: Arc = Arc::new(InMemory::new()); + let store: Arc = if footer_padding > 0 { + assert!(footer_bytes > 512 * 1024); + Arc::new(ThrottledStore::new( + InMemory::new(), + ThrottleConfig { + wait_get_per_call: Duration::from_millis(10), + ..Default::default() + }, + )) + } else { + Arc::new(InMemory::new()) + }; store .put(&location, Bytes::from(file_bytes).into()) .await @@ -999,7 +1027,28 @@ mod tests { .unwrap(); let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption); - reader.get_metadata(Some(&options)).await.unwrap(); + if footer_padding > 0 { + let metadata = reader.get_metadata(Some(&options)); + tokio::pin!(metadata); + + tokio::select! { + result = &mut metadata => { + panic!("metadata load completed before the second footer read: {result:?}"); + } + () = async { + while reader_metric(&metrics, "scan_io_object_store_get_calls") < 2 { + tokio::task::yield_now().await; + } + } => { + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 0); + assert_eq!(reader_metric(&metrics, "scan_io_footer_bytes"), 0); + } + } + + metadata.await.unwrap(); + } else { + reader.get_metadata(Some(&options)).await.unwrap(); + } assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 1); assert_eq!( From 838434c89cc06d634237b7df7b564817e2e14059 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 23 Aug 2026 21:38:37 -0700 Subject: [PATCH 10/12] Fix bounded producer cleanup and plaintext footer accounting --- native/core/src/execution/jni_api.rs | 21 +++++- .../eager_page_index_reader_factory.rs | 40 +++++++++-- native/core/src/parquet/parquet_exec.rs | 67 +++++++++++++++++++ 3 files changed, 120 insertions(+), 8 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 4b1a9d18aaf..3c6f0d1e4f4 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1007,7 +1007,6 @@ fn stop_batch_producer(producer: JoinHandle<()>) { } std::thread::sleep(Duration::from_millis(1)); } - let _ = futures::executor::block_on(producer); } fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { @@ -1438,4 +1437,24 @@ mod tests { assert!(elapsed < Duration::from_millis(400)); } + + #[test] + fn stops_finished_batch_producer_with_exhausted_runtime_budget() { + let (sender, receiver) = std::sync::mpsc::channel(); + + get_runtime().spawn(async move { + let producer = get_runtime().spawn(async {}); + while !producer.is_finished() { + tokio::task::yield_now().await; + } + while tokio::task::coop::has_budget_remaining() { + tokio::task::coop::consume_budget().await; + } + + stop_batch_producer(producer); + sender.send(()).unwrap(); + }); + + assert_eq!(receiver.recv_timeout(Duration::from_millis(500)), Ok(())); + } } diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 22cb1f194e4..3f83210d929 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -309,8 +309,7 @@ impl AsyncFileReader for EagerPageIndexReader { if metadata.is_ok() { let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); if footer_bytes > 0 && cache_enabled { - scan_io_metrics.footer_reads.add(1); - scan_io_metrics.footer_bytes.add(footer_bytes); + metadata_store.record_footer(footer_bytes); } if cache_enabled { scan_io_metrics.record_metadata_cache_result( @@ -345,6 +344,18 @@ struct ScanIoObjectStore { } impl ScanIoObjectStore { + fn record_footer(&self, bytes: usize) { + if let ScanIoStoreRole::Metadata { + footer_recorded, .. + } = &self.role + { + if bytes > 0 && !footer_recorded.swap(true, Ordering::Relaxed) { + self.scan_io_metrics.footer_reads.add(1); + self.scan_io_metrics.footer_bytes.add(bytes); + } + } + } + fn record_request(&self, bytes: usize) { if bytes == 0 { return; @@ -370,7 +381,6 @@ impl ScanIoObjectStore { .add(bytes.len()), ScanIoStoreRole::Metadata { footer_payload_bytes, - footer_recorded, file_size, record_footer_immediately, .. @@ -400,10 +410,8 @@ impl ScanIoObjectStore { >= footer_end }) }) - && !footer_recorded.swap(true, Ordering::Relaxed) { - self.scan_io_metrics.footer_reads.add(1); - self.scan_io_metrics.footer_bytes.add(footer_bytes); + self.record_footer(footer_bytes); } } } @@ -504,7 +512,25 @@ impl ObjectStore for ScanIoObjectStore { ) .await } - ScanIoStoreRole::Metadata { .. } => { + ScanIoStoreRole::Metadata { + footer_payload_bytes, + file_size, + record_footer_immediately, + .. + } => { + let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); + if !record_footer_immediately + && footer_bytes > 0 + && !ranges.is_empty() + && file_size + .saturating_sub(8) + .checked_sub(footer_bytes as u64) + .is_some_and(|footer_start| { + ranges.iter().all(|range| range.end <= footer_start) + }) + { + self.record_footer(footer_bytes); + } self.record_request(ranges_bytes(ranges)); let bytes = self.inner.get_ranges(location, ranges).await?; for (range, bytes) in ranges.iter().zip(bytes.iter()) { diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 5f147f48522..4d5eb3a358e 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -918,6 +918,73 @@ mod tests { assert_eq!(reader_metric(&metrics, "scan_io_metadata_cache_misses"), 1); } + #[tokio::test] + async fn reports_plaintext_footer_before_page_index_read() { + let (filename, _schema) = write_scan_io_fixture(); + let file_bytes = std::fs::read(filename).unwrap(); + let file_size = file_bytes.len(); + let footer_bytes = usize::try_from(u32::from_le_bytes( + file_bytes[file_size - 8..file_size - 4].try_into().unwrap(), + )) + .unwrap(); + let location = object_store::path::Path::from("plaintext-footer.parquet"); + let store: Arc = Arc::new(ThrottledStore::new( + InMemory::new(), + ThrottleConfig { + wait_get_per_call: Duration::from_millis(10), + ..Default::default() + }, + )); + store + .put(&location, Bytes::from(file_bytes).into()) + .await + .unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + let metadata_cache = session_ctx + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = EagerPageIndexReaderFactory::new( + store, + metadata_cache, + ScanIoSource::ObjectStore, + &metrics, + ); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::new(location.to_string(), file_size as u64), + Some(footer_bytes + 8), + &metrics, + ) + .unwrap(); + let metadata = reader.get_metadata(None); + tokio::pin!(metadata); + + tokio::select! { + result = &mut metadata => { + panic!("metadata load completed before the page-index read: {result:?}"); + } + () = async { + while reader_metric(&metrics, "scan_io_object_store_get_calls") < 2 { + tokio::task::yield_now().await; + } + } => { + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 1); + assert_eq!(reader_metric(&metrics, "scan_io_footer_bytes"), footer_bytes); + } + } + + metadata.await.unwrap(); + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 1); + assert_eq!( + reader_metric(&metrics, "scan_io_footer_bytes"), + footer_bytes + ); + } + #[tokio::test] async fn reports_encrypted_footer_before_key_retrieval() { assert_encrypted_footer_before_key_retrieval(0).await; From e01bf0827a2477de48643b3d56181867a9a90d99 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 24 Aug 2026 10:04:48 -0700 Subject: [PATCH 11/12] Preserve footer metrics when prefetched page indexes fail --- .../eager_page_index_reader_factory.rs | 59 ++++++++++++++----- native/core/src/parquet/parquet_exec.rs | 58 +++++++++++++++++- 2 files changed, 101 insertions(+), 16 deletions(-) diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 3f83210d929..671e25c7677 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -69,11 +69,13 @@ use object_store::{ }; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; -use parquet::file::metadata::{FooterTail, PageIndexPolicy, ParquetMetaData}; +use parquet::file::metadata::{ + FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, +}; use std::fmt::{Debug, Display, Formatter}; use std::ops::Range; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum ScanIoSource { @@ -286,6 +288,7 @@ impl AsyncFileReader for EagerPageIndexReader { role: ScanIoStoreRole::Metadata { storage_reads: Arc::clone(&metadata_storage_reads), footer_payload_bytes: Arc::clone(&footer_payload_bytes), + footer_payload: OnceLock::new(), footer_recorded: AtomicBool::new(false), file_size: object_meta.size, record_footer_immediately: !cache_enabled, @@ -316,6 +319,8 @@ impl AsyncFileReader for EagerPageIndexReader { metadata_storage_reads.load(Ordering::Relaxed), ); } + } else if cache_enabled { + metadata_store.record_valid_footer(); } metadata @@ -330,6 +335,7 @@ enum ScanIoStoreRole { Metadata { storage_reads: Arc, footer_payload_bytes: Arc, + footer_payload: OnceLock, footer_recorded: AtomicBool, file_size: u64, record_footer_immediately: bool, @@ -344,6 +350,24 @@ struct ScanIoObjectStore { } impl ScanIoObjectStore { + fn record_valid_footer(&self) { + if let ScanIoStoreRole::Metadata { + footer_payload_bytes, + footer_payload, + footer_recorded, + .. + } = &self.role + { + if !footer_recorded.load(Ordering::Relaxed) + && footer_payload + .get() + .is_some_and(|payload| ParquetMetaDataReader::decode_metadata(payload).is_ok()) + { + self.record_footer(footer_payload_bytes.load(Ordering::Relaxed)); + } + } + } + fn record_footer(&self, bytes: usize) { if let ScanIoStoreRole::Metadata { footer_recorded, .. @@ -381,6 +405,7 @@ impl ScanIoObjectStore { .add(bytes.len()), ScanIoStoreRole::Metadata { footer_payload_bytes, + footer_payload, file_size, record_footer_immediately, .. @@ -397,21 +422,25 @@ impl ScanIoObjectStore { } } - if *record_footer_immediately { - let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); - let footer_end = file_size.saturating_sub(8); - if footer_bytes > 0 - && footer_end - .checked_sub(footer_bytes as u64) - .is_some_and(|footer_start| { - range.is_some_and(|range| { - range.start <= footer_start - && range.start.saturating_add(bytes.len() as u64) - >= footer_end - }) + let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); + let footer_end = file_size.saturating_sub(8); + if footer_bytes > 0 + && footer_end + .checked_sub(footer_bytes as u64) + .is_some_and(|footer_start| { + range.is_some_and(|range| { + range.start <= footer_start + && range.start.saturating_add(bytes.len() as u64) >= footer_end }) - { + }) + { + if *record_footer_immediately { self.record_footer(footer_bytes); + } else if let Some(range) = range { + let payload_start = + (footer_end - footer_bytes as u64 - range.start) as usize; + let _ = footer_payload + .set(bytes.slice(payload_start..payload_start + footer_bytes)); } } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 4d5eb3a358e..1599bf57617 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -343,7 +343,7 @@ mod tests { use parquet::arrow::ArrowWriter; use parquet::encryption::decrypt::{FileDecryptionProperties, KeyRetriever}; use parquet::encryption::encrypt::FileEncryptionProperties; - use parquet::file::metadata::{KeyValue, PageIndexPolicy}; + use parquet::file::metadata::{KeyValue, PageIndexPolicy, ParquetMetaDataReader}; use parquet::file::properties::{EnabledStatistics, WriterProperties}; use std::fs::File; use std::time::Duration; @@ -985,6 +985,62 @@ mod tests { ); } + #[tokio::test] + async fn reports_plaintext_footer_when_prefetched_page_index_is_invalid() { + let (filename, _schema) = write_scan_io_fixture(); + let mut file_bytes = std::fs::read(filename).unwrap(); + let file_size = file_bytes.len(); + let footer_bytes = usize::try_from(u32::from_le_bytes( + file_bytes[file_size - 8..file_size - 4].try_into().unwrap(), + )) + .unwrap(); + let footer_start = file_size - 8 - footer_bytes; + let footer = + ParquetMetaDataReader::decode_metadata(&file_bytes[footer_start..file_size - 8]) + .unwrap(); + let column = footer.row_group(0).column(0); + let index_start = usize::try_from(column.column_index_offset().unwrap()).unwrap(); + let index_bytes = usize::try_from(column.column_index_length().unwrap()).unwrap(); + file_bytes[index_start..index_start + index_bytes].fill(0xff); + + let location = object_store::path::Path::from("invalid-prefetched-page-index.parquet"); + let store: Arc = Arc::new(InMemory::new()); + store + .put(&location, Bytes::from(file_bytes).into()) + .await + .unwrap(); + + let session_ctx = Arc::new(SessionContext::new()); + let metadata_cache = session_ctx + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = EagerPageIndexReaderFactory::new( + store, + metadata_cache, + ScanIoSource::ObjectStore, + &metrics, + ); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::new(location.to_string(), file_size as u64), + Some(512 * 1024), + &metrics, + ) + .unwrap(); + + assert!(reader.get_metadata(None).await.is_err()); + assert_eq!(reader_metric(&metrics, "scan_io_metadata_bytes"), file_size); + assert_eq!(reader_metric(&metrics, "scan_io_object_store_get_calls"), 1); + assert_eq!(reader_metric(&metrics, "scan_io_footer_reads"), 1); + assert_eq!( + reader_metric(&metrics, "scan_io_footer_bytes"), + footer_bytes + ); + } + #[tokio::test] async fn reports_encrypted_footer_before_key_retrieval() { assert_encrypted_footer_before_key_retrieval(0).await; From 130ee02bb4d3bd127c154a20cca3070ea5979644 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 27 Aug 2026 17:57:38 +0000 Subject: [PATCH 12/12] Address scan I/O review and extract independent fixes --- docs/source/user-guide/latest/metrics.md | 48 +++- native/core/src/execution/jni_api.rs | 85 +----- .../src/execution/operators/parquet_writer.rs | 2 +- native/core/src/execution/planner.rs | 19 +- .../eager_page_index_reader_factory.rs | 183 +++++++++++++ native/core/src/parquet/mod.rs | 17 +- native/core/src/parquet/parquet_exec.rs | 127 +++------ native/core/src/parquet/parquet_support.rs | 247 +++++++----------- .../sql/benchmark/CometReadBenchmark.scala | 86 ++++++ 9 files changed, 455 insertions(+), 359 deletions(-) diff --git a/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index a8b2d6bdb8a..c3fb35d8e67 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -79,11 +79,50 @@ Here is a guide to some of the native metrics. | `spilled_bytes` | Actual bytes written to native shuffle spill files on disk. | | `memory_spilled_bytes` | Uncompressed Arrow backing-buffer and partition-index memory spilled. | +### Native Parquet scans + +Native Parquet scans expose these counters in the Spark SQL metric map as well as native +execution metrics. Counters accumulate per scan operator; they do not instrument individual rows. + +| Metric | Description | +| ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `scan_io_data_bytes` | Bytes returned to the Parquet reader for projected data-page ranges. | +| `scan_io_metadata_bytes` | Bytes returned for footer prefetches, page indexes, and Bloom filters. A footer prefetch can also contain unused data bytes. | +| `scan_io_footer_reads` | Storage reads of complete serialized footer payloads, counted once per metadata open. Plaintext payloads must decode successfully; encrypted payloads are counted before key retrieval/decryption, even if those later fail. | +| `scan_io_footer_bytes` | Serialized footer payload bytes, excluding the final eight-byte trailer. These bytes are already included in `scan_io_metadata_bytes`. | +| `scan_io_object_store_get_calls` | Nonempty GET operations at the native remote `ObjectStore` API, after range coalescing. Does not count transport retries or HEAD requests. | +| `scan_io_object_store_get_requested_bytes` | Requested coalesced range bytes at that API. A failed or partly consumed request can contribute requested bytes without the same number of response bytes. | +| `scan_io_object_store_response_bytes_read` | Response bytes actually consumed at that API, including bytes fetched between coalesced ranges. Not HTTP wire bytes. | +| `scan_io_metadata_cache_hits` | Successful, cache-eligible metadata opens requiring no storage reads. | +| `scan_io_metadata_cache_misses` | Successful, cache-eligible metadata opens requiring storage reads. Failed opens and encrypted opens, which bypass this shared cache, increment neither cache counter. | + +Reader-level and object-store bytes are two views of the same reads; do not add them together. +Likewise, footer bytes are a subset of metadata bytes, not a third reader-level category. A warm +metadata-cache hit contributes no new metadata or footer I/O. A valid plaintext footer followed by +a page-index failure still contributes footer bytes; an invalid plaintext footer does not. + +Remote counters follow the backend selected during object-store construction, including native +S3 (`s3`/`s3a`), GCS (`gs`), Azure (`az`, `adl`, `azure`, `abfs`, `abfss`), and HTTP(S) stores. +Local files, in-memory stores, and HDFS/custom backends (including cloud-looking schemes selected +through `fs.comet.libhdfs.schemes`) retain reader-level counters but have zero remote counters. +The native cloud wrapper observes default range coalescing; custom `get_ranges` implementations +require an explicit accounting contract before being composed with that wrapper. + +For data-bearing scans, comparing remote response bytes with reader data bytes can reveal +coalescing and metadata overhead. For a metadata-only scan, data bytes are zero: report the +metadata and remote totals instead of dividing by zero. Cancellation can leave late asynchronous +work outside the final metric snapshot; these counters are not a guarantee of complete network +traffic accounting after cancellation. + ## Task-Level Input Metrics on Spark 4.1+ -Comet's native scans set `inputMetrics.bytesRead` to the actual file IO performed by the -DataFusion parquet reader (`bytes_scanned`). This is the truthful number you would see at the -filesystem layer. +Comet's native scans populate `inputMetrics.bytesRead` from the existing `bytes_scanned` +counter. It counts requested data/Bloom-filter ranges through the Parquet reader's byte-read +methods, not all filesystem I/O. Footer and page-index reads through metadata loading bypass +this counter, and range coalescing can fetch more bytes than the logical ranges request. The +additional scan I/O metrics above expose those differences without changing `bytes_scanned`. +The native `scan_efficiency_ratio` still uses `bytes_scanned` as its numerator and has the same +blind spots; it is not the remote read-amplification ratio described above. Spark 4.1 changed its own parquet reader to pre-open the `SeekableInputStream` and read the file footer outside the `FileScanRDD.compute()` thread. Spark's `inputMetrics.bytesRead` is updated @@ -100,4 +139,5 @@ unaffected and remains exactly equal between Comet and Spark. If you compare Comet's `bytesRead` against vanilla Spark's on Spark 4.1+ (via the Spark UI or the REST API), expect Comet's number to be substantially larger for small files, and closer to -Spark's for large files. Comet's value reflects what the storage layer actually delivered. +Spark's for large files in that workload. Neither metric should be interpreted as complete +filesystem or network traffic accounting. diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 3c6f0d1e4f4..d80754736b7 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -95,7 +95,6 @@ use std::time::{Duration, Instant}; use std::{sync::Arc, task::Poll}; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; -use tokio::task::JoinHandle; use crate::execution::memory_pools::{ create_memory_pool, handle_task_shared_pool_release, parse_memory_pool_config, MemoryPoolConfig, @@ -325,7 +324,6 @@ struct ExecutionContext { pub stream: Option, /// Receives batches from a spawned tokio task (async I/O path) pub batch_receiver: Option>>, - pub batch_producer: Option>, /// Native metrics pub metrics: Arc>>, // The interval in milliseconds to update metrics @@ -539,7 +537,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( input_sources, stream: None, batch_receiver: None, - batch_producer: None, metrics, metrics_update_interval, metrics_last_update_time: Instant::now(), @@ -838,7 +835,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // decreasing to 1 would serialize production and consumption. let (tx, rx) = mpsc::channel(2); let mut stream = stream; - let producer = get_runtime().spawn(async move { + get_runtime().spawn(async move { let result = std::panic::AssertUnwindSafe(async { while let Some(batch) = stream.next().await { if tx.send(batch).await.is_err() { @@ -865,7 +862,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } }); exec_context.batch_receiver = Some(rx); - exec_context.batch_producer = Some(producer); } else { exec_context.stream = Some(stream); } @@ -970,11 +966,6 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( try_unwrap_or_throw(&e, |env| unsafe { let execution_context = get_execution_context(exec_context); - execution_context.batch_receiver.take(); - if let Some(producer) = execution_context.batch_producer.take() { - stop_batch_producer(producer); - } - // Update metrics update_metrics(env, execution_context)?; @@ -998,17 +989,6 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( }) } -fn stop_batch_producer(producer: JoinHandle<()>) { - producer.abort(); - let deadline = Instant::now() + Duration::from_millis(100); - while !producer.is_finished() { - if Instant::now() >= deadline { - return; - } - std::thread::sleep(Duration::from_millis(1)); - } -} - fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> { if let Some(native_query) = &exec_context.root_op { let metrics = exec_context.metrics.as_obj(); @@ -1395,66 +1375,3 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( Ok(()) }) } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn waits_for_background_batch_producer_shutdown() { - let (sender, receiver) = std::sync::mpsc::channel::<()>(); - let producer = get_runtime().spawn(async move { - let _sender = sender; - futures::future::pending::<()>().await; - }); - - stop_batch_producer(producer); - - assert_eq!( - receiver.try_recv(), - Err(std::sync::mpsc::TryRecvError::Disconnected) - ); - } - - #[test] - fn does_not_wait_indefinitely_for_blocked_batch_producer() { - let (started_sender, started_receiver) = std::sync::mpsc::channel(); - let (release_sender, release_receiver) = std::sync::mpsc::channel(); - let producer = get_runtime().spawn(async move { - started_sender.send(()).unwrap(); - release_receiver.recv().unwrap(); - }); - started_receiver.recv().unwrap(); - - let release_thread = std::thread::spawn(move || { - std::thread::sleep(Duration::from_millis(500)); - release_sender.send(()).unwrap(); - }); - let started = Instant::now(); - stop_batch_producer(producer); - let elapsed = started.elapsed(); - release_thread.join().unwrap(); - - assert!(elapsed < Duration::from_millis(400)); - } - - #[test] - fn stops_finished_batch_producer_with_exhausted_runtime_budget() { - let (sender, receiver) = std::sync::mpsc::channel(); - - get_runtime().spawn(async move { - let producer = get_runtime().spawn(async {}); - while !producer.is_finished() { - tokio::task::yield_now().await; - } - while tokio::task::coop::has_budget_remaining() { - tokio::task::coop::consume_budget().await; - } - - stop_batch_producer(producer); - sender.send(()).unwrap(); - }); - - assert_eq!(receiver.recv_timeout(Duration::from_millis(500)), Ok(())); - } -} diff --git a/native/core/src/execution/operators/parquet_writer.rs b/native/core/src/execution/operators/parquet_writer.rs index d6e81b85e11..94e37d4b854 100644 --- a/native/core/src/execution/operators/parquet_writer.rs +++ b/native/core/src/execution/operators/parquet_writer.rs @@ -311,7 +311,7 @@ impl ParquetWriterExec { #[cfg(feature = "hdfs-opendal")] { // Use prepare_object_store_with_configs to create and register the object store - let (_object_store_url, object_store_path) = prepare_object_store_with_configs( + let (_object_store_url, object_store_path, _) = prepare_object_store_with_configs( _runtime_env, output_file_path.to_string(), object_store_options, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 76e5f725114..ab9421f4e74 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -94,7 +94,7 @@ use iceberg::expr::Bind; use crate::execution::operators::ExecutionError::GeneralError; use crate::execution::shuffle::{CometPartitioning, CompressionCodec}; use crate::execution::spark_plan::SparkPlan; -use crate::parquet::parquet_support::{is_hdfs_scheme, prepare_object_store_with_configs}; +use crate::parquet::parquet_support::prepare_object_store_with_configs; use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, @@ -1631,13 +1631,12 @@ impl PhysicalPlanner { .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let is_hdfs_object_store = url::Url::parse(&one_file) - .is_ok_and(|url| is_hdfs_scheme(&url, &object_store_options)); - let (object_store_url, _) = prepare_object_store_with_configs( - self.session_ctx.runtime_env(), - one_file, - &object_store_options, - )?; + let (object_store_url, _, object_store_backend) = + prepare_object_store_with_configs( + self.session_ctx.runtime_env(), + one_file, + &object_store_options, + )?; // Get files for this partition let files = self.get_partitioned_files(partition_files)?; @@ -1648,7 +1647,7 @@ impl PhysicalPlanner { Some(data_schema), Some(partition_schema), object_store_url, - is_hdfs_object_store, + object_store_backend, file_groups, Some(projection_vector), Some(data_filters?), @@ -1686,7 +1685,7 @@ impl PhysicalPlanner { .and_then(|f| f.partitioned_file.first()) .map(|f| f.file_path.clone()) .ok_or(GeneralError("Failed to locate file".to_string()))?; - let (object_store_url, _) = prepare_object_store_with_configs( + let (object_store_url, _, _) = prepare_object_store_with_configs( self.session_ctx.runtime_env(), one_file, &object_store_options, diff --git a/native/core/src/parquet/eager_page_index_reader_factory.rs b/native/core/src/parquet/eager_page_index_reader_factory.rs index 671e25c7677..bdadc72f983 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -225,8 +225,14 @@ struct EagerPageIndexReader { } impl AsyncFileReader for EagerPageIndexReader { + // Pinned parquet 58.4 uses this single-range entry point for Bloom filters, and the + // multi-range entry point below for data pages. Footer/page-index loading is instrumented + // separately in get_metadata. The scan tests pin this method contract; revisit it when + // upgrading parquet rather than assuming an offset alone identifies metadata. fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { let requested = range_bytes(&range); + // Preserve the existing requested-range metric. Metadata fetched through get_metadata + // bypasses it, and object-store coalescing can fetch more bytes than these ranges. self.file_metrics.bytes_scanned.add(requested); let scan_io_metrics = Arc::clone(&self.scan_io_metrics); let future = self.inner.get_bytes(range); @@ -397,6 +403,13 @@ impl ScanIoObjectStore { } } + // Footer accounting follows the DataFusion 54.1/parquet 58.4 metadata push decoder: a tail + // read supplies the payload length, then one or more reads supply the complete payload. + // Plaintext payloads are retained here and counted only after metadata succeeds, or when + // get_ranges observes the subsequent page-index request wholly below the footer. Thus a + // later index failure does not erase a decoded footer. Encrypted payloads are counted as + // soon as complete, before key retrieval/decryption; their counter measures payload I/O, + // not successful authentication. Recheck this protocol when upgrading the decoder. fn record_returned(&self, range: Option<&Range>, bytes: &Bytes) { match &self.role { ScanIoStoreRole::ObjectStore => self @@ -435,6 +448,9 @@ impl ScanIoObjectStore { }) { if *record_footer_immediately { + // Encrypted opens bypass the shared cache and retrieve the key only after + // the complete encrypted payload is read. Count completed payload I/O, + // even if key retrieval, authentication, or metadata decoding later fails. self.record_footer(footer_bytes); } else if let Some(range) = range { let payload_start = @@ -534,6 +550,11 @@ impl ObjectStore for ScanIoObjectStore { ) -> ObjectStoreResult> { match &self.role { ScanIoStoreRole::ObjectStore => { + // Supported native cloud stores use object_store's default coalescing. Apply it + // here so get_opts observes the coalesced requests, not just logical ranges. + // This deliberately bypasses an inner get_ranges override: a cache or custom + // backend must define that observation boundary before being composed here. + // Local/custom HDFS backends are not wrapped in this role and keep delegation. coalesce_ranges( ranges, |range| self.get_range(location, range), @@ -558,6 +579,11 @@ impl ObjectStore for ScanIoObjectStore { ranges.iter().all(|range| range.end <= footer_start) }) { + // In parquet 58.4, plaintext page-index requests below the footer begin only + // after its metadata payload has decoded successfully. Record that footer + // before awaiting the indexes, so a later index-read failure does not erase + // a successful footer read. Encrypted opens use the complete-payload path + // above; cached plaintext opens are recorded only after metadata succeeds. self.record_footer(footer_bytes); } self.record_request(ranges_bytes(ranges)); @@ -627,3 +653,160 @@ impl Drop for EagerPageIndexReader { .set_total(self.partitioned_file.object_meta.size as usize); } } + +#[cfg(test)] +mod tests { + use super::*; + use object_store::memory::InMemory; + + #[derive(Debug)] + struct RecordingRangeStore { + inner: InMemory, + range_calls: AtomicUsize, + get_calls: AtomicUsize, + } + + impl Display for RecordingRangeStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "recording-range-store") + } + } + + #[async_trait] + impl ObjectStore for RecordingRangeStore { + async fn put_opts( + &self, + p: &Path, + v: PutPayload, + o: PutOptions, + ) -> ObjectStoreResult { + self.inner.put_opts(p, v, o).await + } + + async fn put_multipart_opts( + &self, + p: &Path, + o: PutMultipartOptions, + ) -> ObjectStoreResult> { + self.inner.put_multipart_opts(p, o).await + } + + async fn get_opts(&self, p: &Path, o: GetOptions) -> ObjectStoreResult { + self.get_calls.fetch_add(1, Ordering::Relaxed); + self.inner.get_opts(p, o).await + } + + async fn get_ranges( + &self, + p: &Path, + ranges: &[Range], + ) -> ObjectStoreResult> { + self.range_calls.fetch_add(1, Ordering::Relaxed); + self.inner.get_ranges(p, ranges).await + } + + fn delete_stream( + &self, + paths: futures::stream::BoxStream<'static, ObjectStoreResult>, + ) -> futures::stream::BoxStream<'static, ObjectStoreResult> { + self.inner.delete_stream(paths) + } + + fn list( + &self, + p: Option<&Path>, + ) -> futures::stream::BoxStream<'static, ObjectStoreResult> { + self.inner.list(p) + } + + async fn list_with_delimiter(&self, p: Option<&Path>) -> ObjectStoreResult { + self.inner.list_with_delimiter(p).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> ObjectStoreResult<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + options: RenameOptions, + ) -> ObjectStoreResult<()> { + self.inner.rename_opts(from, to, options).await + } + } + + #[tokio::test] + async fn preserves_custom_range_delegation_for_local_and_other_backends() { + assert_range_read_contract(ScanIoSource::Local).await; + // This is also the classification returned for custom libhdfs schemes, including s3. + assert_range_read_contract(ScanIoSource::OtherObjectStore).await; + } + + #[tokio::test] + async fn remote_metrics_observe_default_coalescing_instead_of_inner_override() { + assert_range_read_contract(ScanIoSource::ObjectStore).await; + } + + async fn assert_range_read_contract(source: ScanIoSource) { + let store = Arc::new(RecordingRangeStore { + inner: InMemory::new(), + range_calls: AtomicUsize::new(0), + get_calls: AtomicUsize::new(0), + }); + let location = Path::from("ranges.parquet"); + store + .put(&location, Bytes::from_static(b"0123456789").into()) + .await + .unwrap(); + let runtime = datafusion::execution::runtime_env::RuntimeEnv::default(); + let metrics = ExecutionPlanMetricsSet::new(); + let factory = EagerPageIndexReaderFactory::new( + Arc::clone(&store) as Arc, + runtime.cache_manager.get_file_metadata_cache(), + source, + &metrics, + ); + let mut reader = factory + .create_reader( + 0, + PartitionedFile::new(location.to_string(), 10), + None, + &metrics, + ) + .unwrap(); + let result = reader.get_byte_ranges(vec![0..2, 4..6]).await.unwrap(); + assert_eq!( + result, + vec![Bytes::from_static(b"01"), Bytes::from_static(b"45")] + ); + let remote = source == ScanIoSource::ObjectStore; + assert_eq!( + store.range_calls.load(Ordering::Relaxed), + usize::from(!remote) + ); + assert_eq!(store.get_calls.load(Ordering::Relaxed), usize::from(remote)); + assert_eq!( + metrics + .clone_inner() + .sum_by_name("scan_io_data_bytes") + .unwrap() + .as_usize(), + 4 + ); + assert_eq!( + metrics + .clone_inner() + .sum_by_name("scan_io_object_store_response_bytes_read") + .unwrap() + .as_usize(), + if remote { 6 } else { 0 } + ); + } +} diff --git a/native/core/src/parquet/mod.rs b/native/core/src/parquet/mod.rs index 8c27698b4d7..6a06a34cd7f 100644 --- a/native/core/src/parquet/mod.rs +++ b/native/core/src/parquet/mod.rs @@ -49,7 +49,7 @@ use crate::execution::utils::SparkArrowConvert; use crate::jvm_bridge::JVMClasses; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_exec::init_datasource_exec; -use crate::parquet::parquet_support::{is_hdfs_scheme, prepare_object_store_with_configs}; +use crate::parquet::parquet_support::prepare_object_store_with_configs; use arrow::array::{Array, RecordBatch}; use datafusion::datasource::listing::PartitionedFile; use datafusion::execution::SendableRecordBatchStream; @@ -160,13 +160,12 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat let path: String = file_path.try_to_string(env).unwrap(); let object_store_config = get_object_store_options(env, object_store_options)?; - let is_hdfs_object_store = - url::Url::parse(&path).is_ok_and(|url| is_hdfs_scheme(&url, &object_store_config)); - let (object_store_url, object_store_path) = prepare_object_store_with_configs( - session_ctx.runtime_env(), - path.clone(), - &object_store_config, - )?; + let (object_store_url, object_store_path, object_store_backend) = + prepare_object_store_with_configs( + session_ctx.runtime_env(), + path.clone(), + &object_store_config, + )?; let required_schema_buffer = env.convert_byte_array(&required_schema)?; let required_schema = Arc::new(deserialize_schema(&required_schema_buffer)?); @@ -215,7 +214,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat Some(data_schema), None, object_store_url, - is_hdfs_object_store, + object_store_backend, file_groups, None, data_filters, diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 1599bf57617..d99000ff726 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -18,6 +18,7 @@ use crate::execution::operators::ExecutionError; use crate::parquet::eager_page_index_reader_factory::{EagerPageIndexReaderFactory, ScanIoSource}; use crate::parquet::encryption_support::{CometEncryptionConfig, ENCRYPTION_FACTORY_ID}; +use crate::parquet::parquet_support::ObjectStoreBackend; use crate::parquet::parquet_support::SparkParquetOptions; use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory; use arrow::datatypes::{Field, SchemaRef}; @@ -62,7 +63,7 @@ pub(crate) fn init_datasource_exec( data_schema: Option, partition_schema: Option, object_store_url: ObjectStoreUrl, - is_hdfs_object_store: bool, + object_store_backend: ObjectStoreBackend, file_groups: Vec>, projection_vector: Option>, data_filters: Option>>, @@ -159,10 +160,14 @@ pub(crate) fn init_datasource_exec( // cached with the footer, at the cost of losing the skip's benefit when it would have // applied. Filed upstream as apache/datafusion#23978; revert this once that's fixed. // + // Preserve bytes_scanned's existing requested data/Bloom-filter range accounting. Footer + // and page-index reads through get_metadata bypass it, and coalescing may fetch extra bytes. + // The scan I/O counters expose those gaps without changing task inputMetrics.bytesRead or + // scan_efficiency_ratio, which continue to derive from bytes_scanned. let runtime_env = session_ctx.runtime_env(); let store = runtime_env.object_store(&object_store_url)?; let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache(); - let scan_io_source = scan_io_source(&object_store_url, is_hdfs_object_store); + let scan_io_source = scan_io_source(object_store_backend); let reader_factory = Arc::new(EagerPageIndexReaderFactory::new( store, metadata_cache, @@ -220,26 +225,12 @@ pub(crate) fn init_datasource_exec( Ok(data_source_exec) } -fn scan_io_source(object_store_url: &ObjectStoreUrl, is_hdfs_object_store: bool) -> ScanIoSource { - if is_hdfs_object_store { - return ScanIoSource::OtherObjectStore; +fn scan_io_source(backend: ObjectStoreBackend) -> ScanIoSource { + match backend { + ObjectStoreBackend::Local => ScanIoSource::Local, + ObjectStoreBackend::Remote => ScanIoSource::ObjectStore, + ObjectStoreBackend::Other => ScanIoSource::OtherObjectStore, } - - match physical_object_store_scheme(object_store_url) { - "file" => ScanIoSource::Local, - "s3" | "s3a" | "gs" | "az" | "abfs" | "abfss" | "http" | "https" => { - ScanIoSource::ObjectStore - } - _ => ScanIoSource::OtherObjectStore, - } -} - -fn physical_object_store_scheme(object_store_url: &ObjectStoreUrl) -> &str { - let store_url: &url::Url = object_store_url.as_ref(); - store_url - .scheme() - .split_once("+comet-") - .map_or_else(|| store_url.scheme(), |(scheme, _)| scheme) } #[allow(clippy::too_many_arguments)] @@ -302,15 +293,10 @@ fn get_options( spark_parquet_options.allow_timestamp_ltz_to_ntz = allow_timestamp_ltz_to_ntz; if encryption_enabled { - let store_url: &url::Url = object_store_url.as_ref(); table_parquet_options.crypto.configure_factory( ENCRYPTION_FACTORY_ID, &CometEncryptionConfig { - uri_base: format!( - "{}://{}/", - physical_object_store_scheme(object_store_url), - &store_url[url::Position::BeforeHost..url::Position::AfterPort], - ), + uri_base: object_store_url.to_string(), }, ); } @@ -403,7 +389,7 @@ mod tests { Some(data_schema), None, ObjectStoreUrl::local_filesystem(), - false, + ObjectStoreBackend::Local, vec![vec![partitioned_file]], projection, filters, @@ -489,69 +475,6 @@ mod tests { assert_eq!(global.coerce_int96_tz, Some("UTC".to_string())); } - #[test] - fn preserves_custom_hdfs_backend_range_reads_for_cloud_schemes() { - let options = HashMap::from([( - "fs.comet.libhdfs.schemes".to_string(), - "abfs,s3".to_string(), - )]); - - for scheme in ["abfs", "s3"] { - let url = ObjectStoreUrl::parse(format!("{scheme}://bucket")).unwrap(); - let original_url = url::Url::parse(format!("{scheme}://bucket").as_str()).unwrap(); - let is_hdfs_store = - crate::parquet::parquet_support::is_hdfs_scheme(&original_url, &options); - assert_eq!(scan_io_source(&url, false), ScanIoSource::ObjectStore); - assert_eq!( - scan_io_source(&url, is_hdfs_store), - ScanIoSource::OtherObjectStore - ); - } - - let s3a_original_url = url::Url::parse("s3a://bucket").unwrap(); - let normalized_s3_url = ObjectStoreUrl::parse("s3://bucket").unwrap(); - let preserved_s3a_url = ObjectStoreUrl::parse("s3a://bucket").unwrap(); - let isolated_s3a_url = - ObjectStoreUrl::parse("s3a+comet-0123456789abcdef-native://bucket").unwrap(); - let is_hdfs_store = - crate::parquet::parquet_support::is_hdfs_scheme(&s3a_original_url, &options); - assert_eq!( - scan_io_source(&normalized_s3_url, is_hdfs_store), - ScanIoSource::ObjectStore - ); - assert_eq!( - scan_io_source(&preserved_s3a_url, is_hdfs_store), - ScanIoSource::ObjectStore - ); - assert_eq!( - scan_io_source(&isolated_s3a_url, is_hdfs_store), - ScanIoSource::ObjectStore - ); - } - - #[test] - fn preserves_physical_uri_for_isolated_encrypted_object_stores() { - let object_store_url = - ObjectStoreUrl::parse("s3a+comet-0123456789abcdef-native://bucket").unwrap(); - let (table_parquet_options, _) = get_options( - "UTC", - true, - false, - false, - false, - &object_store_url, - true, - &ParquetOptions::default(), - ); - let encryption_options: CometEncryptionConfig = table_parquet_options - .crypto - .factory_options - .to_extension_options() - .unwrap(); - - assert_eq!(encryption_options.uri_base, "s3a://bucket/"); - } - // Regression test for #3978: DataFusion's opener requests `PageIndexPolicy::Skip` on the // initial metadata load and only loads the page index later, on demand, when row-group // pruning shows it is still needed (apache/datafusion#22857). That on-demand load bypasses @@ -595,7 +518,7 @@ mod tests { None, None, ObjectStoreUrl::local_filesystem(), - false, + ObjectStoreBackend::Local, vec![vec![partitioned_file]], None, None, @@ -1043,15 +966,20 @@ mod tests { #[tokio::test] async fn reports_encrypted_footer_before_key_retrieval() { - assert_encrypted_footer_before_key_retrieval(0).await; + assert_encrypted_footer_before_key_retrieval(0, false).await; } #[tokio::test] async fn does_not_report_encrypted_footer_before_payload_is_read() { - assert_encrypted_footer_before_key_retrieval(512 * 1024).await; + assert_encrypted_footer_before_key_retrieval(512 * 1024, false).await; } - async fn assert_encrypted_footer_before_key_retrieval(footer_padding: usize) { + #[tokio::test] + async fn counts_complete_encrypted_footer_even_when_authentication_fails() { + assert_encrypted_footer_before_key_retrieval(0, true).await; + } + + async fn assert_encrypted_footer_before_key_retrieval(footer_padding: usize, corrupt: bool) { struct ObservingKeyRetriever { key: Vec, metrics: Arc, @@ -1096,13 +1024,18 @@ mod tests { writer.write(&batch).unwrap(); writer.close().unwrap(); - let file_bytes = std::fs::read(&filename).unwrap(); + let mut file_bytes = std::fs::read(&filename).unwrap(); let file_size = file_bytes.len(); let footer_bytes = usize::try_from(u32::from_le_bytes( file_bytes[file_size - 8..file_size - 4].try_into().unwrap(), )) .unwrap(); let location = object_store::path::Path::from("encrypted-footer.parquet"); + if corrupt { + // Corrupt the encrypted payload's authentication tag, preserving its length/trailer + // and key metadata. Complete I/O is counted before decryption is attempted. + file_bytes[file_size - 9] ^= 1; + } let store: Arc = if footer_padding > 0 { assert!(footer_bytes > 512 * 1024); Arc::new(ThrottledStore::new( @@ -1169,6 +1102,8 @@ mod tests { } metadata.await.unwrap(); + } else if corrupt { + assert!(reader.get_metadata(Some(&options)).await.is_err()); } else { reader.get_metadata(Some(&options)).await.unwrap(); } diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index 2dc884bceb0..d4f4cce9ab7 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -37,7 +37,7 @@ use datafusion::physical_plan::ColumnarValue; use datafusion_comet_spark_expr::EvalMode; use log::debug; use object_store::path::Path; -use object_store::{parse_url, ObjectStore}; +use object_store::{parse_url, ObjectStore, ObjectStoreScheme}; use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use std::collections::HashMap; use std::sync::OnceLock; @@ -469,9 +469,9 @@ fn create_hdfs_object_store( }) } -type ObjectStoreCache = RwLock>>; +type ObjectStoreCache = RwLock>>; -/// Process-wide cache of object stores, keyed by `(scheme://host:port, config_hash, hdfs_backend)`. +/// Process-wide cache of object stores, keyed by `(scheme://host:port, config_hash)`. /// /// ## Why static / process lifetime? /// @@ -486,8 +486,8 @@ type ObjectStoreCache = RwLock /// /// ## Unbounded size /// -/// Cache entries are indexed by `(scheme://host:port, hash-of-configs, hdfs_backend)`. A typical -/// Spark job accesses a small, fixed set of buckets with a stable configuration, so the number of +/// Cache entries are indexed by `(scheme://host:port, hash-of-configs)`. A typical Spark +/// job accesses a small, fixed set of buckets with a stable configuration, so the number of /// distinct keys is O(buckets × credential-configs) and remains small throughout the job. /// Entries are cheap relative to the cost of creating a new object store (new HTTP /// connection pool + DNS resolution), and there is no meaningful benefit from eviction, so @@ -519,17 +519,43 @@ fn hash_object_store_configs(configs: &HashMap) -> u64 { hasher.finish() } -/// Parses the url, registers the object store with configurations, and returns a tuple of the object store url -/// and object store path +/// The selected backend, independent of the URL used to register it in DataFusion. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ObjectStoreBackend { + Local, + Remote, + Other, +} + +fn object_store_backend(url: &Url, is_hdfs: bool) -> Result { + if is_hdfs { + // Custom libhdfs schemes may look like cloud URLs but retain their own range-read API. + return Ok(ObjectStoreBackend::Other); + } + let (scheme, _) = + ObjectStoreScheme::parse(url).map_err(|e| ExecutionError::GeneralError(e.to_string()))?; + Ok(match scheme { + ObjectStoreScheme::Local => ObjectStoreBackend::Local, + ObjectStoreScheme::AmazonS3 + | ObjectStoreScheme::GoogleCloudStorage + | ObjectStoreScheme::MicrosoftAzure + | ObjectStoreScheme::Http => ObjectStoreBackend::Remote, + // Memory and future backends are not implicitly classified as remote network traffic. + _ => ObjectStoreBackend::Other, + }) +} + +/// Parses the URL, selects and registers the backend, and returns its registry URL, path, and +/// classification. Callers must use this classification rather than infer it from a scheme alias. pub(crate) fn prepare_object_store_with_configs( runtime_env: Arc, url: String, object_store_configs: &HashMap, -) -> Result<(ObjectStoreUrl, Path), ExecutionError> { +) -> Result<(ObjectStoreUrl, Path, ObjectStoreBackend), ExecutionError> { let mut url = Url::parse(url.as_str()) .map_err(|e| ExecutionError::GeneralError(format!("Error parsing URL {url}: {e}")))?; - let original_url = url.clone(); let is_hdfs_scheme = is_hdfs_scheme(&url, object_store_configs); + let backend = object_store_backend(&url, is_hdfs_scheme)?; let mut scheme = url.scheme(); if !is_hdfs_scheme && scheme == "s3a" { scheme = "s3"; @@ -544,7 +570,7 @@ pub(crate) fn prepare_object_store_with_configs( ); let config_hash = hash_object_store_configs(object_store_configs); - let cache_key = (url_key.clone(), config_hash, is_hdfs_scheme); + let cache_key = (url_key.clone(), config_hash); // Check the cache first to reuse existing object store instances. // This enables HTTP connection pooling and avoids redundant DNS lookups. @@ -582,33 +608,65 @@ pub(crate) fn prepare_object_store_with_configs( (store, path) }; - let object_store_url = ObjectStoreUrl::parse(url_key.as_str())?; - let registration_url = if scheme != "file" - && runtime_env - .object_store(&object_store_url) - .is_ok_and(|existing| !Arc::ptr_eq(&existing, &object_store)) - { - let backend = if is_hdfs_scheme { "hdfs" } else { "native" }; - Url::parse(&format!( - "{}+comet-{config_hash:016x}-{backend}://{}", - original_url.scheme(), - &url[url::Position::BeforeHost..url::Position::AfterPort], - )) - .map_err(|e| ExecutionError::GeneralError(e.to_string()))? - } else { - url - }; - let object_store_url = ObjectStoreUrl::parse(format!( - "{}://{}", - registration_url.scheme(), - ®istration_url[url::Position::BeforeHost..url::Position::AfterPort], - ))?; - runtime_env.register_object_store(®istration_url, object_store); - Ok((object_store_url, object_store_path)) + let object_store_url = ObjectStoreUrl::parse(url_key.clone())?; + runtime_env.register_object_store(&url, object_store); + Ok((object_store_url, object_store_path, backend)) } #[cfg(test)] mod tests { + #[test] + fn classifies_the_selected_backend_using_object_store_parser() { + use super::{is_hdfs_scheme, object_store_backend, ObjectStoreBackend}; + let configs = std::collections::HashMap::from([( + "fs.comet.libhdfs.schemes".to_string(), + "s3,abfs".to_string(), + )]); + for address in [ + "s3://bucket/path", + "s3a://bucket/path", + "gs://bucket/path", + "az://container/path", + "adl://container/path", + "azure://container/path", + "abfs://container/path", + "abfss://container/path", + "http://example.com/path", + "https://example.com/path", + "https://account.blob.core.windows.net/container/path", + ] { + let url = url::Url::parse(address).unwrap(); + assert_eq!( + object_store_backend(&url, false).unwrap(), + ObjectStoreBackend::Remote, + "{address}" + ); + if is_hdfs_scheme(&url, &configs) { + assert_eq!( + object_store_backend(&url, true).unwrap(), + ObjectStoreBackend::Other + ); + } + } + assert_eq!( + object_store_backend(&url::Url::parse("file:///tmp/a").unwrap(), false).unwrap(), + ObjectStoreBackend::Local + ); + assert_eq!( + object_store_backend(&url::Url::parse("memory:///a").unwrap(), false).unwrap(), + ObjectStoreBackend::Other + ); + // These spellings are not accepted native backends in pinned object_store 0.13.2. + for scheme in ["gcs", "wasb", "wasbs", "s3n"] { + let url = url::Url::parse(&format!("{scheme}://bucket/path")).unwrap(); + assert!(object_store_backend(&url, false).is_err()); + assert_eq!( + object_store_backend(&url, true).unwrap(), + ObjectStoreBackend::Other + ); + } + } + #[cfg(not(feature = "hdfs-opendal"))] use datafusion::execution::object_store::ObjectStoreUrl; #[cfg(not(feature = "hdfs-opendal"))] @@ -625,128 +683,6 @@ mod tests { #[cfg(not(feature = "hdfs-opendal"))] use std::collections::HashMap; - #[test] - fn cache_distinguishes_backend_routing_for_normalized_s3_aliases() { - let configs = std::collections::HashMap::from([( - "fs.comet.libhdfs.schemes".to_string(), - "s3".to_string(), - )]); - let cloud_url = url::Url::parse("s3a://scan-io-backend-routing").unwrap(); - let hdfs_url = url::Url::parse("s3://scan-io-backend-routing").unwrap(); - let config_hash = super::hash_object_store_configs(&configs); - let cloud_key = ( - "s3://scan-io-backend-routing".to_string(), - config_hash, - super::is_hdfs_scheme(&cloud_url, &configs), - ); - let hdfs_key = ( - "s3://scan-io-backend-routing".to_string(), - config_hash, - super::is_hdfs_scheme(&hdfs_url, &configs), - ); - let cloud_store: std::sync::Arc = - std::sync::Arc::new(object_store::memory::InMemory::new()); - let hdfs_store: std::sync::Arc = - std::sync::Arc::new(object_store::memory::InMemory::new()); - - { - let mut cache = super::object_store_cache().write().unwrap(); - cache.insert(cloud_key.clone(), std::sync::Arc::clone(&cloud_store)); - cache.insert(hdfs_key.clone(), std::sync::Arc::clone(&hdfs_store)); - } - - let runtime_env = - std::sync::Arc::new(datafusion::execution::runtime_env::RuntimeEnv::default()); - let (cloud_object_store_url, _) = super::prepare_object_store_with_configs( - std::sync::Arc::clone(&runtime_env), - cloud_url.to_string(), - &configs, - ) - .unwrap(); - let (hdfs_object_store_url, _) = super::prepare_object_store_with_configs( - std::sync::Arc::clone(&runtime_env), - hdfs_url.to_string(), - &configs, - ) - .unwrap(); - - assert_ne!(cloud_object_store_url, hdfs_object_store_url); - assert!(std::sync::Arc::ptr_eq( - &runtime_env.object_store(&cloud_object_store_url).unwrap(), - &cloud_store - )); - assert!(std::sync::Arc::ptr_eq( - &runtime_env.object_store(&hdfs_object_store_url).unwrap(), - &hdfs_store - )); - - let mut cache = super::object_store_cache().write().unwrap(); - cache.remove(&cloud_key); - cache.remove(&hdfs_key); - } - - #[test] - fn registry_distinguishes_backend_routing_across_configurations() { - let cloud_configs = std::collections::HashMap::from([( - "fs.comet.libhdfs.schemes".to_string(), - "s3".to_string(), - )]); - let hdfs_configs = std::collections::HashMap::from([( - "fs.comet.libhdfs.schemes".to_string(), - "s3a".to_string(), - )]); - let url = url::Url::parse("s3a://scan-io-mixed-backend-routing").unwrap(); - let cloud_key = ( - "s3://scan-io-mixed-backend-routing".to_string(), - super::hash_object_store_configs(&cloud_configs), - false, - ); - let hdfs_key = ( - "s3a://scan-io-mixed-backend-routing".to_string(), - super::hash_object_store_configs(&hdfs_configs), - true, - ); - let cloud_store: std::sync::Arc = - std::sync::Arc::new(object_store::memory::InMemory::new()); - let hdfs_store: std::sync::Arc = - std::sync::Arc::new(object_store::memory::InMemory::new()); - - { - let mut cache = super::object_store_cache().write().unwrap(); - cache.insert(cloud_key.clone(), std::sync::Arc::clone(&cloud_store)); - cache.insert(hdfs_key.clone(), std::sync::Arc::clone(&hdfs_store)); - } - - let runtime_env = - std::sync::Arc::new(datafusion::execution::runtime_env::RuntimeEnv::default()); - let (hdfs_object_store_url, _) = super::prepare_object_store_with_configs( - std::sync::Arc::clone(&runtime_env), - url.to_string(), - &hdfs_configs, - ) - .unwrap(); - let (cloud_object_store_url, _) = super::prepare_object_store_with_configs( - std::sync::Arc::clone(&runtime_env), - url.to_string(), - &cloud_configs, - ) - .unwrap(); - - assert_ne!(cloud_object_store_url, hdfs_object_store_url); - assert!(std::sync::Arc::ptr_eq( - &runtime_env.object_store(&cloud_object_store_url).unwrap(), - &cloud_store - )); - assert!(std::sync::Arc::ptr_eq( - &runtime_env.object_store(&hdfs_object_store_url).unwrap(), - &hdfs_store - )); - - let mut cache = super::object_store_cache().write().unwrap(); - cache.remove(&cloud_key); - cache.remove(&hdfs_key); - } - /// Parses the url, registers the object store, and returns a tuple of the object store url and object store path #[cfg(not(feature = "hdfs-opendal"))] pub(crate) fn prepare_object_store( @@ -755,6 +691,7 @@ mod tests { ) -> Result<(ObjectStoreUrl, Path), ExecutionError> { use crate::parquet::parquet_support::prepare_object_store_with_configs; prepare_object_store_with_configs(runtime_env, url, &HashMap::new()) + .map(|(url, path, _)| (url, path)) } #[cfg(not(feature = "hdfs-opendal"))] diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometReadBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometReadBenchmark.scala index 30090d1d126..1055240cd73 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometReadBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometReadBenchmark.scala @@ -34,6 +34,7 @@ import org.apache.spark.TestUtils import org.apache.spark.benchmark.Benchmark import org.apache.spark.sql.{DataFrame, SparkSession} import org.apache.spark.sql.execution.datasources.parquet.VectorizedParquetRecordReader +import org.apache.spark.sql.execution.metric.SQLMetrics import org.apache.spark.sql.types._ import org.apache.spark.sql.vectorized.ColumnVector @@ -47,6 +48,87 @@ import org.apache.comet.{CometConf, WithHdfsCluster} */ class CometReadBaseBenchmark extends CometBenchmarkBase { + /** + * Measure the nine scan I/O accumulators separately from storage work. Run this benchmark with + * `--scan-metric-overhead`; add `--reverse-cases` to reverse the real-job comparison. The first + * two cases isolate driver creation and task-copy/merge costs. The last runs 10,000 Spark tasks + * with zero or nine extra SQL accumulators, including task serialization and scheduler updates. + * This is not an end-to-end scan benchmark: it excludes native counters, JNI traversal, the SQL + * UI's per-node rendering, and storage I/O. + */ + def scanMetricAccumulatorBenchmark(reverseCases: Boolean): Unit = { + val names = Seq( + "scan_io_data_bytes", + "scan_io_metadata_bytes", + "scan_io_footer_reads", + "scan_io_footer_bytes", + "scan_io_object_store_get_calls", + "scan_io_object_store_get_requested_bytes", + "scan_io_object_store_response_bytes_read", + "scan_io_metadata_cache_hits", + "scan_io_metadata_cache_misses") + def createMetrics() = names.map { name => + if (name.endsWith("bytes") || name.endsWith("bytes_read")) { + SQLMetrics.createSizeMetric(spark.sparkContext, name) + } else { + SQLMetrics.createMetric(spark.sparkContext, name) + } + } + + val operators = 1024 + val creation = + new Benchmark("Scan I/O SQL metrics: creation", operators, minNumIters = 5, output = output) + creation.addCase("nine metrics per operator") { _ => + var count = 0 + while (count < operators) { + assert(createMetrics().size == 9) + count += 1 + } + } + creation.run() + + val tasks = 10000 + val driverMetrics = createMetrics() + val updates = new Benchmark( + "Scan I/O SQL metrics: task snapshots", + tasks, + minNumIters = 5, + output = output) + updates.addCase("copy, update and merge nine metrics") { _ => + driverMetrics.foreach(_.reset()) + var task = 0 + while (task < tasks) { + driverMetrics.foreach { driver => + val local = driver.copyAndReset() + local.add(64L) + driver.merge(local) + } + task += 1 + } + assert(driverMetrics.forall(_.value == tasks * 64L)) + } + updates.run() + + val partitions = spark.sparkContext.parallelize(0 until tasks, tasks) + val jobs = new Benchmark( + "Scan I/O SQL metrics: 10000-task Spark job", + tasks, + minNumIters = 3, + output = output) + val cases = if (reverseCases) Seq(9, 0) else Seq(0, 9) + cases.foreach { count => + val metrics = if (count == 0) Seq.empty else createMetrics() + jobs.addCase(s"$count extra SQL accumulators") { _ => + metrics.foreach(_.reset()) + partitions.foreachPartition { _ => + metrics.foreach(_.add(64L)) + } + assert(metrics.forall(_.value == tasks * 64L)) + } + } + jobs.run() + } + def numericScanBenchmark(values: Int, dataType: DataType): Unit = { val sqlBenchmark = new Benchmark(s"SQL Single ${dataType.sql} Column Scan", values, output = output) @@ -316,6 +398,10 @@ class CometReadBaseBenchmark extends CometBenchmarkBase { } override def runCometBenchmark(mainArgs: Array[String]): Unit = { + if (mainArgs.contains("--scan-metric-overhead")) { + scanMetricAccumulatorBenchmark(mainArgs.contains("--reverse-cases")) + return + } runBenchmarkWithTable("Parquet Reader", 1024 * 1024 * 15) { v => Seq( BooleanType,