diff --git a/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index c249b8972dd..0fcacd0ba8f 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -230,12 +230,17 @@ execution metrics. Counters accumulate per scan operator; they do not instrument | `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. | +| `scan_io_metadata_prefetches` | Opens that used a load which read storage ahead of them, started when the previous file opened (`spark.comet.parquet.prefetchNextFileMetadata.enabled`). Each also counts as a cache miss. | | `row_groups_pruned_dynamic_filter` | Row groups skipped by a runtime predicate whose value changes during execution, such as a local TopK threshold. | 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. +metadata-cache hit contributes no new metadata or footer I/O. A prefetched load's footer and +metadata bytes count when the prefetch reads them, so a prefetch that no open uses, because the +scan stopped early or pruned the file before opening it, adds bytes with no matching cache miss and +no `scan_io_metadata_prefetches`. +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. diff --git a/docs/source/user-guide/latest/parquet.md b/docs/source/user-guide/latest/parquet.md index 4d6a8dab0f8..1c96006a145 100644 --- a/docs/source/user-guide/latest/parquet.md +++ b/docs/source/user-guide/latest/parquet.md @@ -215,6 +215,18 @@ more splits overall, so some idle tasks may remain — tune the value against yo See [#3817](https://github.com/apache/datafusion-comet/issues/3817#issuecomment-4193279630) for a worked example and further discussion. +### Prefetching the Next File's Metadata + +A native Parquet scan opens a task's files one after another, and each open waits for a round trip to read +the file's footer and page index before its data fetch can start. For tasks that read many small files from +high-latency storage such as S3, that wait is paid once per file. By default, Comet loads the next file's +footer and page index in the background while the current file is read. Set +`spark.comet.parquet.prefetchNextFileMetadata.enabled` to `false` to turn this off. + +The cost is one extra footer and page index read for a file that is never opened, such as when a scan stops +early or the file is pruned before it opens. Scans of encrypted Parquet files never prefetch. The +`scan_io_metadata_prefetches` metric counts the opens that used a load which read storage ahead of them. + ## Known Issues This section lists open issues in which Comet's Parquet reads and writes give a different result diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 930dcd5bfc7..265338a2283 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -137,10 +137,11 @@ use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, - COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, - COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, + COMET_STRICT_TIMESTAMP_LABELS, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; +use crate::parquet::parquet_exec::CometParquetScanOptions; use crate::parquet::parquet_support::CometObjectStoreRegistry; use datafusion_comet_proto::spark_operator::operator::OpStruct; use log::{info, warn}; @@ -1029,6 +1030,13 @@ fn prepare_datafusion_session_context( session_config.set_str("datafusion.execution.parquet.reorder_filters", "true"); } + // Native Parquet scan options that DataFusion has no setting for travel as a session config + // extension, which every scan in the plan reads. + session_config = session_config.with_extension(Arc::new(CometParquetScanOptions { + prefetch_next_file_metadata: spark_config + .get_bool(COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED), + })); + // Translate Comet's spill codec into DataFusion's spill compression, which DataFusion // otherwise leaves uncompressed. Set before the pass-through below for the same reason, // so an explicit `spark.comet.datafusion.execution.spill_compression` still wins. diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index c27d4d11d14..511b5733388 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -24,6 +24,8 @@ pub(crate) const COMET_MAX_TEMP_DIRECTORY_SIZE: &str = "spark.comet.maxTempDirec pub(crate) const COMET_DEBUG_MEMORY: &str = "spark.comet.debug.memory"; pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; +pub(crate) const COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED: &str = + "spark.comet.parquet.prefetchNextFileMetadata.enabled"; pub(crate) const COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: &str = "spark.comet.exec.aggregate.skipPartial.enabled"; pub(crate) const COMET_EXEC_SPILL_COMPRESSION_CODEC: &str = 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 25023f4580c..8ebe639bc87 100644 --- a/native/core/src/parquet/eager_page_index_reader_factory.rs +++ b/native/core/src/parquet/eager_page_index_reader_factory.rs @@ -76,6 +76,21 @@ //! row filters enabled, a deferred conversion failure is raised on the first data-page request. //! Footer, Bloom-filter and page-index reads still proceed, so format pruning can discard a file //! without decoding it, but row selection cannot hide the failure. +//! +//! With `with_next_file_metadata_prefetch`, readers also load the next file's metadata before +//! DataFusion opens it. DataFusion's `FileStream` keeps one I/O outstanding and plans a file +//! only after the previous file's reader has returned its last batch (apache/datafusion#20916 +//! removed its next-file prefetch), so without this every file boundary waits for a footer and +//! page-index round trip. When a file opens, its reader spawns a task that loads the metadata of +//! the next file in the scan's file group through the cache, and that file's open waits for the +//! task instead of reading the footer itself. This assumes DataFusion opens the files in the +//! group's order, which holds unless the source reorders them (`ParquetSource::reorder_files`); +//! a reordered scan would only waste its prefetches. +//! +//! This is a stopgap for the overlap DataFusion dropped. It can go once `FileStream` opens the next +//! file while the current reader runs again; apache/datafusion#5209 (buffer more than one file +//! open) and apache/datafusion#18470 (prefetch the next row group) are the open issues that would +//! bring it back. use crate::parquet::datetime_rebase::stamp_int96_leaves; use crate::parquet::parquet_support::SparkParquetOptions; @@ -85,6 +100,7 @@ use crate::parquet::schema_adapter::{ use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; use async_trait::async_trait; use bytes::Bytes; +use datafusion::common::runtime::SpawnedTask; use datafusion::common::{DataFusionError, Result as DFResult}; use datafusion::datasource::physical_plan::parquet::metadata::{ CachedParquetMetaData, DFParquetMetadata, @@ -97,7 +113,7 @@ use datafusion::physical_plan::metrics::{ Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricType, }; use datafusion_comet_common::SparkError; -use datafusion_datasource::PartitionedFile; +use datafusion_datasource::{FileRange, PartitionedFile}; use futures::future::BoxFuture; use futures::{FutureExt, StreamExt, TryStreamExt}; use object_store::path::Path; @@ -111,6 +127,7 @@ use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::{encode_arrow_schema, parquet_to_arrow_schema, ARROW_SCHEMA_META_KEY}; use parquet::basic::{ConvertedType, LogicalType}; +use parquet::encryption::decrypt::FileDecryptionProperties; use parquet::errors::{ParquetError, Result as ParquetResult}; use parquet::file::metadata::{FileMetaData, KeyValue, ParquetMetaDataBuilder}; use parquet::file::metadata::{ @@ -145,6 +162,7 @@ struct ScanIoMetrics { object_store_response_bytes_read: Count, metadata_cache_hits: Count, metadata_cache_misses: Count, + metadata_prefetches: Count, } impl ScanIoMetrics { @@ -167,6 +185,7 @@ impl ScanIoMetrics { ), metadata_cache_hits: count_counter(metrics, "scan_io_metadata_cache_hits"), metadata_cache_misses: count_counter(metrics, "scan_io_metadata_cache_misses"), + metadata_prefetches: count_counter(metrics, "scan_io_metadata_prefetches"), } } @@ -215,6 +234,7 @@ pub struct EagerPageIndexReaderFactory { stamp_int96_leaves: bool, conversion_check: Option<(SchemaRef, SparkParquetOptions, bool)>, deferred_rejections: Arc>>, + next_file_prefetch: Option>, } impl EagerPageIndexReaderFactory { @@ -247,6 +267,7 @@ impl EagerPageIndexReaderFactory { stamp_int96_leaves: false, conversion_check: None, deferred_rejections: Arc::new(Mutex::new(HashMap::new())), + next_file_prefetch: None, } } @@ -280,6 +301,26 @@ impl EagerPageIndexReaderFactory { self.stamp_int96_leaves = enabled; self } + + /// With `Some(files)`, load the footer and page index of the file DataFusion opens next while + /// the current file is read (see the module docs). `files` is the scan's one file group, in + /// the order DataFusion opens it. `None`, the default, leaves the prefetch off. A prefetch + /// carries no decryption properties, so an open that has them neither waits for nor starts + /// one. + pub(crate) fn with_next_file_metadata_prefetch( + mut self, + files: Option<&[PartitionedFile]>, + ) -> Self { + self.next_file_prefetch = files.map(|files| { + Arc::new(NextFileMetadataPrefetch::new( + files, + Arc::clone(&self.store), + Arc::clone(&self.scan_io_metrics), + Arc::clone(&self.metadata_cache), + )) + }); + self + } } /// True when `node` or any node under it carries a field id, the way Spark's @@ -332,6 +373,7 @@ impl ParquetFileReaderFactory for EagerPageIndexReaderFactory { conversion_check: self.conversion_check.clone(), deferred_rejections: Arc::clone(&self.deferred_rejections), deferred_rejection, + next_file_prefetch: self.next_file_prefetch.clone(), })) } } @@ -352,6 +394,7 @@ struct EagerPageIndexReader { conversion_check: Option<(SchemaRef, SparkParquetOptions, bool)>, deferred_rejections: Arc>>, deferred_rejection: Option, + next_file_prefetch: Option>, } // Arrow infers ENUM as Binary, losing the distinction from raw binary that Spark needs. @@ -606,11 +649,18 @@ impl AsyncFileReader for EagerPageIndexReader { let spark_variant_schema = self.spark_variant_schema; let require_field_ids = self.require_field_ids; let stamp_enabled = self.stamp_int96_leaves; + let file_decryption_properties = options + .and_then(|o| o.file_decryption_properties()) + .map(Arc::clone); + let cache_enabled = file_decryption_properties.is_none(); + // Take the load of this file's metadata that the previous file's open started, if any, + // and start the next file's. + let pending_load = self + .next_file_prefetch + .as_ref() + .filter(|_| cache_enabled) + .and_then(|prefetch| prefetch.open(&self.partitioned_file, metadata_size_hint)); 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(); if spark_variant_schema && file_decryption_properties.is_some() { return Err(ParquetError::General( "Projected Variant with Parquet encryption requires Spark fallback".to_string(), @@ -621,50 +671,38 @@ impl AsyncFileReader for EagerPageIndexReader { } else { options.map(|o| o.column_index_policy()) }; - let metadata_storage_reads = Arc::new(AtomicUsize::new(0)); - let footer_payload_bytes = Arc::new(AtomicUsize::new(0)); - let metadata_store = ScanIoObjectStore { - inner: store, - scan_io_metrics: Arc::clone(&scan_io_metrics), - 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, - }, + // Wait for the prefetch rather than reading the footer a second time. A prefetch + // that failed or was cancelled yields nothing, and this open then loads the metadata + // itself, so any error it reports is the open's own. + let prefetched = match pending_load { + Some(load) => load.await.ok().flatten(), + None => None, }; - - let metadata = DFParquetMetadata::new(&metadata_store, &object_meta) - .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(Arc::clone(&metadata_cache))) - .with_metadata_size_hint(metadata_size_hint) - .with_page_index_policy(page_index_policy) - .fetch_metadata() - .await - .map_err(|e| { - parquet::errors::ParquetError::General(format!( - "Failed to fetch metadata for file {}: {e}", - object_meta.location, - )) - }); - - if metadata.is_ok() { - let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); - if footer_bytes > 0 && cache_enabled { - metadata_store.record_footer(footer_bytes); + let (metadata, storage_reads) = match prefetched { + Some((metadata, storage_reads)) => { + // The load read from storage ahead of this open, so the open records a cache + // miss, as it would have without the prefetch, and the load counts as overlapped. + if storage_reads > 0 { + scan_io_metrics.metadata_prefetches.add(1); + } + (metadata, storage_reads) } - if cache_enabled { - scan_io_metrics.record_metadata_cache_result( - metadata_storage_reads.load(Ordering::Relaxed), - ); + None => { + load_metadata( + store, + Arc::clone(&scan_io_metrics), + Arc::clone(&metadata_cache), + &object_meta, + metadata_size_hint, + file_decryption_properties, + page_index_policy, + ) + .await? } - } else if cache_enabled { - metadata_store.record_valid_footer(); + }; + if cache_enabled { + scan_io_metrics.record_metadata_cache_result(storage_reads); } - - let metadata = metadata?; // Spark's missing field id check, over the raw Parquet schema as in // `ParquetReadSupport`. The JNI layer unwraps the `External` error, so the JVM sees // the exception Spark raises. @@ -750,6 +788,172 @@ impl AsyncFileReader for EagerPageIndexReader { } } +/// Loads a file's metadata through `metadata_cache`, with the page index as `page_index_policy` +/// asks, counting the metadata and footer bytes it reads from storage. Returns the metadata +/// together with the number of storage reads it took, which is zero when the cache already held +/// the metadata. +async fn load_metadata( + store: Arc, + scan_io_metrics: Arc, + metadata_cache: Arc, + object_meta: &ObjectMeta, + metadata_size_hint: Option, + file_decryption_properties: Option>, + page_index_policy: Option, +) -> ParquetResult<(Arc, usize)> { + let cache_enabled = file_decryption_properties.is_none(); + let storage_reads = Arc::new(AtomicUsize::new(0)); + let footer_payload_bytes = Arc::new(AtomicUsize::new(0)); + let metadata_store = ScanIoObjectStore { + inner: store, + scan_io_metrics, + role: ScanIoStoreRole::Metadata { + storage_reads: Arc::clone(&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, + }, + }; + + 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) + .with_page_index_policy(page_index_policy) + .fetch_metadata() + .await + .map_err(|e| { + ParquetError::General(format!( + "Failed to fetch metadata for file {}: {e}", + object_meta.location, + )) + }); + + if metadata.is_ok() { + let footer_bytes = footer_payload_bytes.load(Ordering::Relaxed); + if footer_bytes > 0 && cache_enabled { + metadata_store.record_footer(footer_bytes); + } + } else if cache_enabled { + metadata_store.record_valid_footer(); + } + Ok((metadata?, storage_reads.load(Ordering::Relaxed))) +} + +/// A load of a file's metadata started before DataFusion opens the file. Dropping it aborts the +/// load. +/// Its result carries the number of storage reads it took, for the open that uses it to record. +type MetadataLoad = SpawnedTask, usize)>>; + +/// The scan-wide state of `with_next_file_metadata_prefetch`. +#[derive(Debug)] +struct NextFileMetadataPrefetch { + store: Arc, + scan_io_metrics: Arc, + metadata_cache: Arc, + /// For each file of the group, keyed by path and byte range, the next file with another path + /// and that file's metadata size hint. Further ranges of a file's own path are skipped, + /// because opening the file already loads their metadata. + next_files: HashMap<(Path, Option), (ObjectMeta, Option)>, + /// The load in flight, with the path it loads. Each open replaces it unless it is for the + /// next file, so a load for a file DataFusion never opens, such as one pruned before its + /// open, is dropped at the next open. + in_flight: Mutex>, +} + +impl NextFileMetadataPrefetch { + fn new( + files: &[PartitionedFile], + store: Arc, + scan_io_metrics: Arc, + metadata_cache: Arc, + ) -> Self { + // One pass from the back, so a long run of ranges of one path costs linear time. The next + // file of `files[i]` is `files[i + 1]` when its path differs and `files[i + 1]`'s own next + // file otherwise. Inserting while walking backwards lets the earliest of any duplicate + // key win. + let mut next_files = HashMap::new(); + let mut next: Option<&PartitionedFile> = None; + for (position, file) in files.iter().enumerate().rev() { + let location = &file.object_meta.location; + if let Some(following) = files.get(position + 1) { + if &following.object_meta.location != location { + next = Some(following); + } + } + if let Some(next) = next { + next_files.insert( + (location.clone(), file.range.clone()), + (next.object_meta.clone(), next.metadata_size_hint), + ); + } + } + Self { + store, + scan_io_metrics, + metadata_cache, + next_files, + in_flight: Mutex::new(None), + } + } + + /// Called as `file` opens: takes the load of its metadata that the previous open started, if + /// any, and starts loading the next file's. + fn open( + &self, + file: &PartitionedFile, + metadata_size_hint: Option, + ) -> Option { + let mut in_flight = self.in_flight.lock().unwrap(); + let pending = in_flight + .take_if(|(location, _)| *location == file.object_meta.location) + .map(|(_, load)| load); + let key = (file.object_meta.location.clone(), file.range.clone()); + match self.next_files.get(&key) { + // Another range of an earlier file already started this load. + Some((next, _)) + if in_flight + .as_ref() + .is_some_and(|(location, _)| *location == next.location) => {} + next => { + *in_flight = + next.and_then(|(next, hint)| self.start(next, hint.or(metadata_size_hint))); + } + } + pending + } + + /// Starts loading `object_meta`'s metadata on the current Tokio runtime, if there is one. + fn start( + &self, + object_meta: &ObjectMeta, + metadata_size_hint: Option, + ) -> Option<(Path, MetadataLoad)> { + tokio::runtime::Handle::try_current().ok()?; + let store = Arc::clone(&self.store); + let scan_io_metrics = Arc::clone(&self.scan_io_metrics); + let metadata_cache = Arc::clone(&self.metadata_cache); + let object_meta = object_meta.clone(); + let location = object_meta.location.clone(); + let load = SpawnedTask::spawn(async move { + load_metadata( + store, + scan_io_metrics, + metadata_cache, + &object_meta, + metadata_size_hint, + None, + Some(PageIndexPolicy::Optional), + ) + .await + .ok() + }); + Some((location, load)) + } +} + #[derive(Debug)] enum ScanIoStoreRole { ObjectStore, @@ -1094,21 +1298,82 @@ mod tests { }, }; + /// An `InMemory` store that counts the gets and range calls of each path. Reads of the path + /// set by `with_gate` wait until `release` is called, and each held read keeps a receiver of + /// `gate`. #[derive(Debug)] - struct RecordingRangeStore { + struct RecordingStore { inner: InMemory, - range_calls: AtomicUsize, - get_calls: AtomicUsize, + gets: Mutex>, + range_calls: Mutex>, + gated: Option, + gate: tokio::sync::watch::Sender, } - impl Display for RecordingRangeStore { + impl RecordingStore { + fn new() -> Self { + Self { + inner: InMemory::new(), + gets: Mutex::new(HashMap::new()), + range_calls: Mutex::new(HashMap::new()), + gated: None, + gate: tokio::sync::watch::channel(false).0, + } + } + + fn with_gate(mut self, gated: &str) -> Self { + self.gated = Some(Path::from(gated)); + self + } + + fn release(&self) { + self.gate.send_replace(true); + } + + /// The gets and range calls of `location`. + fn reads(&self, location: &str) -> usize { + let location = Path::from(location); + [&self.gets, &self.range_calls] + .iter() + .map(|calls| calls.lock().unwrap().get(&location).copied().unwrap_or(0)) + .sum() + } + + fn get_calls(&self) -> usize { + self.gets.lock().unwrap().values().sum() + } + + fn range_calls(&self) -> usize { + self.range_calls.lock().unwrap().values().sum() + } + + /// Waits until a read of the gated path is held at the gate. + async fn wait_for_held_read(&self) { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + while self.gate.receiver_count() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the gated path was never read"); + } + + async fn hold_if_gated(&self, location: &Path) { + if self.gated.as_ref() == Some(location) { + let mut gate = self.gate.subscribe(); + gate.wait_for(|released| *released).await.unwrap(); + } + } + } + + impl Display for RecordingStore { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "recording-range-store") + write!(f, "recording-store") } } #[async_trait] - impl ObjectStore for RecordingRangeStore { + impl ObjectStore for RecordingStore { async fn put_opts( &self, p: &Path, @@ -1127,7 +1392,8 @@ mod tests { } async fn get_opts(&self, p: &Path, o: GetOptions) -> ObjectStoreResult { - self.get_calls.fetch_add(1, Ordering::Relaxed); + *self.gets.lock().unwrap().entry(p.clone()).or_default() += 1; + self.hold_if_gated(p).await; self.inner.get_opts(p, o).await } @@ -1136,7 +1402,13 @@ mod tests { p: &Path, ranges: &[Range], ) -> ObjectStoreResult> { - self.range_calls.fetch_add(1, Ordering::Relaxed); + *self + .range_calls + .lock() + .unwrap() + .entry(p.clone()) + .or_default() += 1; + self.hold_if_gated(p).await; self.inner.get_ranges(p, ranges).await } @@ -1190,11 +1462,7 @@ mod tests { } 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 store = Arc::new(RecordingStore::new()); let location = Path::from("ranges.parquet"); store .put(&location, Bytes::from_static(b"0123456789").into()) @@ -1222,11 +1490,8 @@ mod tests { 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!(store.range_calls(), usize::from(!remote)); + assert_eq!(store.get_calls(), usize::from(remote)); assert_eq!( metrics .clone_inner() @@ -1496,4 +1761,245 @@ mod tests { &with_spark_arrow_schema(Arc::clone(&rewritten)).unwrap() )); } + + async fn put_int_files(store: &InMemory, locations: &[&str]) -> Vec { + let mut files = vec![]; + for location in locations { + files.push(put_int_file(store, location, None).await); + } + files + } + + fn prefetching_factory( + store: Arc, + files: &[PartitionedFile], + metrics: &ExecutionPlanMetricsSet, + ) -> EagerPageIndexReaderFactory { + let runtime = datafusion::execution::runtime_env::RuntimeEnv::default(); + EagerPageIndexReaderFactory::new( + store, + runtime.cache_manager.get_file_metadata_cache(), + ScanIoSource::Local, + metrics, + ) + .with_next_file_metadata_prefetch(Some(files)) + } + + /// Opens `file` the way DataFusion's opener does: a new reader, then its metadata. + async fn open( + factory: &EagerPageIndexReaderFactory, + file: &PartitionedFile, + metrics: &ExecutionPlanMetricsSet, + ) -> ParquetResult> { + let mut reader = factory + .create_reader(0, file.clone(), None, metrics) + .unwrap(); + reader.get_metadata(None).await + } + + fn metric(metrics: &ExecutionPlanMetricsSet, name: &str) -> usize { + metrics + .clone_inner() + .sum_by_name(name) + .unwrap_or_else(|| panic!("missing metric {name}")) + .as_usize() + } + + fn in_flight_is_empty(factory: &EagerPageIndexReaderFactory) -> bool { + let prefetch = factory.next_file_prefetch.as_ref().unwrap(); + let in_flight = prefetch.in_flight.lock().unwrap(); + in_flight.is_none() + } + + /// Opening a file loads the next file's metadata in the background. Every file is read from + /// storage, so each open is a cache miss, and the two loads that ran ahead of their open are + /// counted as prefetches. Opening the last file starts nothing. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn opening_a_file_loads_the_next_files_metadata_ahead() { + let store = Arc::new(InMemory::new()); + let files = put_int_files(&store, &["a.parquet", "b.parquet", "c.parquet"]).await; + let metrics = ExecutionPlanMetricsSet::new(); + let factory = prefetching_factory(store, &files, &metrics); + + for file in &files { + let metadata = open(&factory, file, &metrics).await.unwrap(); + assert_eq!(metadata.file_metadata().num_rows(), 2); + assert!(metadata.column_index().is_some() && metadata.offset_index().is_some()); + } + + assert_eq!(metric(&metrics, "scan_io_metadata_cache_misses"), 3); + assert_eq!(metric(&metrics, "scan_io_metadata_cache_hits"), 0); + assert_eq!(metric(&metrics, "scan_io_metadata_prefetches"), 2); + assert_eq!(metric(&metrics, "scan_io_footer_reads"), 3); + assert!(in_flight_is_empty(&factory)); + } + + /// A file DataFusion never opens, such as one pruned before its open, doesn't keep the load + /// started for it: the next open reads its own metadata, and the load it then starts replaces + /// the stale one. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn a_load_for_a_file_never_opened_is_replaced_by_the_next_one() { + let store = Arc::new(InMemory::new()); + let files = put_int_files( + &store, + &["a.parquet", "b.parquet", "c.parquet", "d.parquet"], + ) + .await; + let metrics = ExecutionPlanMetricsSet::new(); + let factory = prefetching_factory(store, &files, &metrics); + + // b is never opened. + for file in [&files[0], &files[2], &files[3]] { + open(&factory, file, &metrics).await.unwrap(); + } + + assert_eq!(metric(&metrics, "scan_io_metadata_cache_misses"), 3); + assert_eq!(metric(&metrics, "scan_io_metadata_cache_hits"), 0); + assert_eq!(metric(&metrics, "scan_io_metadata_prefetches"), 1); + assert!(in_flight_is_empty(&factory)); + } + + /// A file whose metadata is still being read ahead when DataFusion opens it waits for that + /// read instead of starting its own. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn an_open_waits_for_the_load_of_its_metadata_in_flight() { + let store = Arc::new(RecordingStore::new().with_gate("b.parquet")); + let files = put_int_files(&store.inner, &["a.parquet", "b.parquet"]).await; + let metrics = ExecutionPlanMetricsSet::new(); + let factory = prefetching_factory(Arc::clone(&store) as _, &files, &metrics); + open(&factory, &files[0], &metrics).await.unwrap(); + store.wait_for_held_read().await; + + let mut metadata = Box::pin(open(&factory, &files[1], &metrics)); + assert!((&mut metadata).now_or_never().is_none()); + assert_eq!(store.reads("b.parquet"), 1, "the open read b itself"); + store.release(); + assert_eq!(metadata.await.unwrap().file_metadata().num_rows(), 2); + + assert_eq!(metric(&metrics, "scan_io_metadata_cache_misses"), 2); + assert_eq!(metric(&metrics, "scan_io_metadata_cache_hits"), 0); + assert_eq!(metric(&metrics, "scan_io_metadata_prefetches"), 1); + assert_eq!(metric(&metrics, "scan_io_footer_reads"), 2); + } + + /// The next file is the next one with a different path. Further ranges of a file share its + /// metadata, so they share its next file too, and opening them doesn't start a second load. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn the_next_file_skips_further_ranges_of_the_same_path() { + let store = Arc::new(RecordingStore::new()); + let a = put_int_file(&store.inner, "a.parquet", None).await; + let b = put_int_file(&store.inner, "b.parquet", None).await; + let files = [ + a.clone().with_range(0, 50), + a.with_range(50, 100), + b.clone(), + ]; + let metrics = ExecutionPlanMetricsSet::new(); + let factory = prefetching_factory(Arc::clone(&store) as _, &files, &metrics); + + for file in &files { + open(&factory, file, &metrics).await.unwrap(); + } + + // b is read once, by the load that the first range of a started. A second range of a + // with no next file would drop that load, and b's open would read it again. + let single_load = Arc::new(RecordingStore::new()); + let alone = put_int_file(&single_load.inner, "b.parquet", None).await; + let alone_metrics = ExecutionPlanMetricsSet::new(); + let alone_factory = prefetching_factory( + Arc::clone(&single_load) as _, + std::slice::from_ref(&alone), + &alone_metrics, + ); + open(&alone_factory, &alone, &alone_metrics).await.unwrap(); + assert_eq!(store.reads("b.parquet"), single_load.reads("b.parquet")); + assert_eq!(metric(&metrics, "scan_io_metadata_prefetches"), 1); + assert!(in_flight_is_empty(&factory)); + } + + /// The map of next files comes from one pass over the group, so a path that comes back later + /// keeps the next file of its earliest range, and a long run of ranges of one path builds + /// without a search per range. + #[tokio::test] + async fn the_next_file_map_handles_repeated_paths_and_long_runs() { + let file = |name: &str| PartitionedFile::new(name.to_string(), 100); + let mut files = vec![ + file("a.parquet"), + file("b.parquet"), + file("a.parquet"), + file("c.parquet"), + ]; + // 40,000 ranges of one path: quadratic construction takes seconds, a linear pass doesn't. + files.extend((0..40_000).map(|i| file("d.parquet").with_range(i, i + 1))); + let metrics = ExecutionPlanMetricsSet::new(); + let started = std::time::Instant::now(); + let factory = prefetching_factory(Arc::new(InMemory::new()), &files, &metrics); + assert!(started.elapsed() < std::time::Duration::from_secs(1)); + + let prefetch = factory.next_file_prefetch.as_ref().unwrap(); + let next = |file: &PartitionedFile| { + prefetch + .next_files + .get(&(file.object_meta.location.clone(), file.range.clone())) + .map(|(next, _)| next.location.to_string()) + }; + // The first a wins over the second, which would give c. + assert_eq!(next(&files[0]).as_deref(), Some("b.parquet")); + assert_eq!(next(&files[1]).as_deref(), Some("a.parquet")); + assert_eq!(next(&files[3]).as_deref(), Some("d.parquet")); + // d has no file after its ranges. + assert_eq!(next(&files[4]), None); + assert_eq!(next(&files[files.len() - 1]), None); + assert_eq!(prefetch.next_files.len(), 3); + } + + /// A prefetch that fails yields nothing, so the file's open loads the metadata itself and + /// reports the error. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn a_failed_prefetch_leaves_the_open_to_report_its_own_error() { + let store = Arc::new(RecordingStore::new()); + let a = put_int_file(&store.inner, "a.parquet", None).await; + let content = Bytes::from_static(b"not a parquet file"); + let bad = PartitionedFile::new("bad.parquet".to_string(), content.len() as u64); + store + .put(&Path::from("bad.parquet"), content.into()) + .await + .unwrap(); + let metrics = ExecutionPlanMetricsSet::new(); + let files = [a, bad]; + let factory = prefetching_factory(Arc::clone(&store) as _, &files, &metrics); + + open(&factory, &files[0], &metrics).await.unwrap(); + // The prefetch's failure doesn't surface until the open of `bad`, which reads it again. + let error = open(&factory, &files[1], &metrics).await.unwrap_err(); + assert!( + error + .to_string() + .contains("Failed to fetch metadata for file bad.parquet"), + "unexpected error: {error}" + ); + assert!( + store.reads("bad.parquet") >= 2, + "both the prefetch and the open should read bad.parquet", + ); + assert_eq!(metric(&metrics, "scan_io_metadata_prefetches"), 0); + } + + /// Dropping the scan cancels a metadata load that is still in flight. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn dropping_the_scan_aborts_a_metadata_load_in_flight() { + let store = Arc::new(RecordingStore::new().with_gate("b.parquet")); + let files = put_int_files(&store.inner, &["a.parquet", "b.parquet"]).await; + let metrics = ExecutionPlanMetricsSet::new(); + let factory = prefetching_factory(Arc::clone(&store) as _, &files, &metrics); + open(&factory, &files[0], &metrics).await.unwrap(); + store.wait_for_held_read().await; + + drop(factory); + // The held read drops its gate receiver only when its task is aborted. + tokio::time::timeout(std::time::Duration::from_secs(10), store.gate.closed()) + .await + .expect("the metadata load in flight was not aborted"); + assert_eq!(metric(&metrics, "scan_io_metadata_prefetches"), 0); + } } diff --git a/native/core/src/parquet/parquet_exec.rs b/native/core/src/parquet/parquet_exec.rs index 660d1073892..aa4bbb020a1 100644 --- a/native/core/src/parquet/parquet_exec.rs +++ b/native/core/src/parquet/parquet_exec.rs @@ -52,6 +52,17 @@ mod variant_tests; /// scan would. pub(crate) const METADATA_SIZE_HINT: usize = 512 * 1024; +/// Native Parquet scan options that are not DataFusion options. +/// `prepare_datafusion_session_context` attaches them to each plan's `SessionConfig` from the +/// plan's Comet configs, and every scan the plan builds reads them there. A session without them +/// scans with every option off. +#[derive(Debug)] +pub(crate) struct CometParquetScanOptions { + /// Load the next file's footer and page index while the current file is read + /// (`spark.comet.parquet.prefetchNextFileMetadata.enabled`). + pub prefetch_next_file_metadata: bool, +} + /// Initializes a DataSourceExec plan with a ParquetSource for Comet's native Parquet scan. /// /// `required_schema`: Schema to be projected by the scan. @@ -233,6 +244,19 @@ pub(crate) fn init_datasource_exec( // its footer metadata (see `datetime_rebase.rs`), which is how the expression adapter // attributes timestamp columns to Spark's INT64 vs INT96 rebase specs. let scan_io_source = scan_io_source(object_store_backend); + // Comet's scans read one file group, the only shape the prefetch handles. An encrypted scan + // never prefetches, since a prefetch carries no decryption properties. + let prefetch_files = match file_groups.as_slice() { + [files] + if !encryption_enabled + && session_config + .get_extension::() + .is_some_and(|options| options.prefetch_next_file_metadata) => + { + Some(files.as_slice()) + } + _ => None, + }; let reader_factory = Arc::new( EagerPageIndexReaderFactory::new( store, @@ -250,7 +274,8 @@ pub(crate) fn init_datasource_exec( && data_filters .as_ref() .is_some_and(|filters| !filters.is_empty()), - ), + ) + .with_next_file_metadata_prefetch(prefetch_files), ); parquet_source = parquet_source.with_parquet_file_reader_factory(reader_factory); @@ -1491,7 +1516,7 @@ mod tests { .iter() .filter(|metric| metric.value().name().starts_with("scan_io_")) .collect::>(); - assert_eq!(scan_io_metrics.len(), 11); + assert_eq!(scan_io_metrics.len(), 12); for metric in scan_io_metrics { assert!(metric.labels().is_empty()); assert_eq!(metric.partition(), None); diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6ffcf4913a7..260dccadfc2 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -200,6 +200,18 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.parquet.prefetchNextFileMetadata.enabled") + .category(CATEGORY_PARQUET) + .doc( + "When enabled, a native Parquet scan that reads several files in one task loads the " + + "next file's footer and page index in the background while it reads the current " + + "file, so opening the next file doesn't wait for a storage round trip. This helps " + + "most when tasks read many small files from high-latency storage such as S3. Scans " + + "of encrypted Parquet files never load ahead.") + .booleanConf + .createWithDefault(true) + val COMET_CONVERT_FROM_PARQUET_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.convert.parquet.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index a4f9d0db7ea..5790554a6bd 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -898,6 +898,7 @@ object CometExecIterator extends Logging { CometConf.COMET_EXPLAIN_NATIVE_ENABLED, CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, CometConf.COMET_NATIVE_UDF_ENABLED, + CometConf.COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, CometConf.COMET_STRICT_TIMESTAMP_LABELS, CometConf.COMET_TRACING_ENABLED).foreach { entry => 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 f390d4f51f1..095d6854578 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 @@ -451,6 +451,10 @@ object CometMetricNode { SQLMetrics.createMetric(sc, "Metadata loads served without storage I/O"), "scan_io_metadata_cache_misses" -> SQLMetrics.createMetric(sc, "Metadata loads requiring storage I/O"), + "scan_io_metadata_prefetches" -> + SQLMetrics.createMetric( + sc, + "Metadata loads read from storage ahead of the file's open and used by it"), "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 4ae13863fb5..346b0b8cd91 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -3453,7 +3453,8 @@ class CometExecSuite extends CometTestBase { "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 => + "scan_io_metadata_cache_misses", + "scan_io_metadata_prefetches").foreach { name => assert(metrics.contains(name), s"Missing $name. Available: ${metrics.keys}") } assert( @@ -3483,6 +3484,32 @@ class CometExecSuite extends CometTestBase { } } + test("native scan loads the next file's metadata while it reads the current file") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark.range(0, 4000, 1, 4).toDF("id").write.parquet(path) + Seq(true, false).foreach { prefetch => + withSQLConf( + CometConf.COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED.key -> prefetch.toString, + // Pack the four files into one task. + SQLConf.FILES_MIN_PARTITION_NUM.key -> "1", + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "1") { + val (_, cometPlan) = checkSparkAnswerAndOperator(spark.read.parquet(path)) + val scan = find(cometPlan)(_.isInstanceOf[CometNativeScanExec]) + assert(scan.isDefined, "Expected to find a CometNativeScanExec") + assert(scan.get.executeColumnar().getNumPartitions == 1) + val metrics = scan.get.metrics + // Every file is read from storage either way. With the prefetch, each open after the + // first finds its metadata already loaded. + assert(metrics("scan_io_metadata_cache_misses").value == 4) + assert(metrics("scan_io_metadata_cache_hits").value == 0) + assert(metrics("scan_io_metadata_prefetches").value == (if (prefetch) 3 else 0)) + assert(metrics("scan_io_footer_reads").value == 4) + } + } + } + } + test("Comet native metrics: project and filter") { withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "true") { withParquetTable((0 until 5).map(i => (i, i + 1)), "tbl") {