viirya commented on code in PR #7094:
URL: https://github.com/apache/datafusion-comet/pull/7094#discussion_r4316733828
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -750,6 +788,164 @@ 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<dyn ObjectStore>,
+ scan_io_metrics: Arc<ScanIoMetrics>,
+ metadata_cache: Arc<FileMetadataCache>,
+ object_meta: &ObjectMeta,
+ metadata_size_hint: Option<usize>,
+ file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
+ page_index_policy: Option<PageIndexPolicy>,
+) -> ParquetResult<(Arc<ParquetMetaData>, 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<Option<(Arc<ParquetMetaData>, usize)>>;
+
+/// The scan-wide state of `with_next_file_metadata_prefetch`.
+#[derive(Debug)]
+struct NextFileMetadataPrefetch {
+ store: Arc<dyn ObjectStore>,
+ scan_io_metrics: Arc<ScanIoMetrics>,
+ metadata_cache: Arc<FileMetadataCache>,
+ /// 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<FileRange>), (ObjectMeta,
Option<usize>)>,
+ /// 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<Option<(Path, MetadataLoad)>>,
+}
+
+impl NextFileMetadataPrefetch {
+ fn new(
+ files: &[PartitionedFile],
+ store: Arc<dyn ObjectStore>,
+ scan_io_metrics: Arc<ScanIoMetrics>,
+ metadata_cache: Arc<FileMetadataCache>,
+ ) -> Self {
+ let mut next_files = HashMap::new();
+ for (position, file) in files.iter().enumerate() {
+ let location = &file.object_meta.location;
+ let next = files[position + 1..]
Review Comment:
+1, I hit the same thing reading it. Within a run of one path, each range
scans the rest of the run before giving up, so a run of n ranges costs about
n²/2 path comparisons. Spark 3.5's `FilePartition.getFilePartitions` keeps the
split ranges when `maxPartitionNum` makes it repack, so a long run in one task
is reachable, even if it takes an unusual config. The fix can stay small. The
successor of `files[i]` is `files[i + 1]` when their paths differ, and
otherwise it's the successor of `files[i + 1]`. So one reverse pass over the
indices gives the same map in linear time, including for a path that shows up
again later in the group, like `[a, b, a, c]`. If you fill the map during that
reverse pass, use `insert` rather than `or_insert_with`, so the earliest
occurrence of a duplicate `(path, range)` key still wins as it does today. The
existing adjacent-range test already covers what the map should contain.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
