andygrove commented on code in PR #7094:
URL: https://github.com/apache/datafusion-comet/pull/7094#discussion_r4319738593


##########
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:
   Good catch, and thanks to both of you for the reproduction and the suggested 
fix. The successor map is now built in one reverse pass: the next file of 
`files[i]` is `files[i + 1]` when the paths differ, and otherwise the next file 
of `files[i + 1]`. It inserts while walking backwards, so the earliest of a 
duplicate `(path, range)` key still wins, including for `[a, b, a, c]`. A new 
test builds the map for that group plus 40,000 ranges of one path and checks 
the successors. It fails on the old loop (6.7 s in a debug build) and passes on 
the new one. (4e0521200b)



-- 
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]

Reply via email to