viirya commented on code in PR #7094:
URL: https://github.com/apache/datafusion-comet/pull/7094#discussion_r4304251145
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -750,6 +776,167 @@ 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.
+type MetadataLoad = SpawnedTask<Option<Arc<ParquetMetaData>>>;
+
+/// 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..]
+ .iter()
+ .find(|next| &next.object_meta.location != location);
+ if let Some(next) = next {
+ next_files
+ .entry((location.clone(), file.range.clone()))
+ .or_insert_with(|| (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<usize>,
+ ) -> Option<MetadataLoad> {
+ 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<usize>,
+ ) -> 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 {
+ let (metadata, storage_reads) = load_metadata(
+ store,
+ Arc::clone(&scan_io_metrics),
+ metadata_cache,
+ &object_meta,
+ metadata_size_hint,
+ None,
+ Some(PageIndexPolicy::Optional),
+ )
+ .await
+ .ok()?;
+ if storage_reads > 0 {
+ scan_io_metrics.metadata_prefetches.add(1);
Review Comment:
Related to the comment on line 692: if the consuming open records the
storage reads itself, this counter could be incremented where the prefetched
load is actually used. Then it counts overlapped loads, and unused loads only
show up as extra metadata bytes, as the doc already says.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -621,50 +666,31 @@ 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) => (metadata, 0),
+ 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?
}
- if cache_enabled {
- scan_io_metrics.record_metadata_cache_result(
- metadata_storage_reads.load(Ordering::Relaxed),
- );
- }
- } else if cache_enabled {
- metadata_store.record_valid_footer();
+ };
+ if cache_enabled {
+ scan_io_metrics.record_metadata_cache_result(storage_reads);
Review Comment:
With the prefetch on, a scan of N cold files reports 1
`scan_io_metadata_cache_misses` and N-1 `scan_io_metadata_cache_hits`, even
though every one of those files was read from storage. Someone using these
counters to judge the metadata cache, for example after #7082, would think the
cache served most opens. You can't get the real hit count back by subtracting
`scan_io_metadata_prefetches` either, because that counter also includes loads
no open used. Could the prefetch task return its `storage_reads` along with the
metadata, so the open that uses it records a miss? Then hits still mean "served
from the cache", and `scan_io_metadata_prefetches` tells how many misses were
overlapped with the previous file. The metrics doc row and the
`CometMetricNode` description would change to match.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -1496,4 +1683,319 @@ mod tests {
&with_spark_arrow_schema(Arc::clone(&rewritten)).unwrap()
));
}
+
+ /// An `InMemory` store whose reads of one path wait until `open` is
called, counting the
+ /// reads of each path. Each held read keeps a receiver of `gate`.
+ #[derive(Debug)]
+ struct GatedStore {
+ inner: InMemory,
+ gated: Path,
+ gate: tokio::sync::watch::Sender<bool>,
+ reads: Mutex<HashMap<Path, usize>>,
+ }
+
+ impl GatedStore {
+ fn new(gated: &str) -> Self {
+ Self {
+ inner: InMemory::new(),
+ gated: Path::from(gated),
+ gate: tokio::sync::watch::channel(false).0,
+ reads: Mutex::new(HashMap::new()),
+ }
+ }
+
+ fn open(&self) {
+ self.gate.send_replace(true);
+ }
+
+ fn reads(&self, location: &str) -> usize {
+ let reads = self.reads.lock().unwrap();
+ reads.get(&Path::from(location)).copied().unwrap_or(0)
+ }
+
+ /// 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");
+ }
+ }
+
+ impl Display for GatedStore {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "gated-store")
+ }
+ }
+
+ #[async_trait]
+ impl ObjectStore for GatedStore {
+ async fn put_opts(
+ &self,
+ p: &Path,
+ v: PutPayload,
+ o: PutOptions,
+ ) -> ObjectStoreResult<PutResult> {
+ self.inner.put_opts(p, v, o).await
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ p: &Path,
+ o: PutMultipartOptions,
+ ) -> ObjectStoreResult<Box<dyn MultipartUpload>> {
+ self.inner.put_multipart_opts(p, o).await
+ }
+
+ async fn get_opts(&self, p: &Path, o: GetOptions) ->
ObjectStoreResult<GetResult> {
+ *self.reads.lock().unwrap().entry(p.clone()).or_default() += 1;
+ if p == &self.gated {
+ let mut gate = self.gate.subscribe();
+ gate.wait_for(|open| *open).await.unwrap();
+ }
+ self.inner.get_opts(p, o).await
+ }
+
+ fn delete_stream(
+ &self,
+ paths: futures::stream::BoxStream<'static,
ObjectStoreResult<Path>>,
+ ) -> futures::stream::BoxStream<'static, ObjectStoreResult<Path>> {
+ self.inner.delete_stream(paths)
+ }
+
+ fn list(
+ &self,
+ p: Option<&Path>,
+ ) -> futures::stream::BoxStream<'static,
ObjectStoreResult<ObjectMeta>> {
+ self.inner.list(p)
+ }
+
+ async fn list_with_delimiter(&self, p: Option<&Path>) ->
ObjectStoreResult<ListResult> {
+ 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
+ }
+ }
+
+ async fn put_int_files(store: &InMemory, locations: &[&str]) ->
Vec<PartitionedFile> {
+ let mut files = vec![];
+ for location in locations {
+ files.push(put_int_file(store, location, None).await);
+ }
+ files
+ }
+
+ fn prefetching_factory(
+ store: Arc<dyn ObjectStore>,
+ 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<Arc<ParquetMetaData>> {
+ 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, so
only the first open
+ /// reads storage. 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"), 1);
+ assert_eq!(metric(&metrics, "scan_io_metadata_cache_hits"), 2);
+ 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"), 2);
+ assert_eq!(metric(&metrics, "scan_io_metadata_cache_hits"), 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(GatedStore::new("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.open();
+ assert_eq!(metadata.await.unwrap().file_metadata().num_rows(), 2);
+
+ assert_eq!(metric(&metrics, "scan_io_metadata_cache_misses"), 1);
+ assert_eq!(metric(&metrics, "scan_io_metadata_cache_hits"), 1);
+ 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.
+ #[test]
+ fn the_next_file_skips_further_ranges_of_the_same_path() {
+ let file = |path: &str, range: Option<(i64, i64)>| {
+ let file = PartitionedFile::new(path.to_string(), 100);
+ match range {
+ Some((start, end)) => file.with_range(start, end),
+ None => file,
+ }
+ };
+ let files = vec![
+ file("a", Some((0, 50))),
+ file("a", Some((50, 100))),
+ file("b", None),
+ ];
+ let runtime =
datafusion::execution::runtime_env::RuntimeEnv::default();
+ let prefetch = NextFileMetadataPrefetch::new(
+ &files,
+ Arc::new(InMemory::new()),
+ Arc::new(ScanIoMetrics::new(&ExecutionPlanMetricsSet::new())),
+ runtime.cache_manager.get_file_metadata_cache(),
+ );
+ let next = |file: &PartitionedFile| {
+ prefetch
+ .next_files
+ .get(&(file.object_meta.location.clone(), file.range.clone()))
+ .map(|(next, _)| next.location.to_string())
+ };
+
+ assert_eq!(next(&files[0]).as_deref(), Some("b"));
+ assert_eq!(next(&files[1]).as_deref(), Some("b"));
+ assert_eq!(next(&files[2]), None);
+ assert_eq!(next(&file("a", Some((0, 10)))), None);
+ }
+
+ /// 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(GatedStore::new("unused.parquet"));
+ 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}"
+ );
+ let reads_per_load = store.reads("bad.parquet") / 2;
Review Comment:
`reads_per_load = reads / 2` followed by `reads == 2 * reads_per_load` only
checks that the count is even and nonzero. Would
`assert!(store.reads("bad.parquet") >= 2, ...)` say more plainly that both the
prefetch and the open read the file? Another option is to measure one load's
reads against a fresh store and compare with that.
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -233,6 +244,17 @@ 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.
+ let prefetch_files = match file_groups.as_slice() {
Review Comment:
An encrypted scan still builds `NextFileMetadataPrefetch` here. The rule
that encrypted scans never prefetch then depends on
`CometEncryptionFactory::get_file_decryption_properties` returning `Some` for
every file. If it ever returned `None` for a plaintext file in an encrypted
table, that open would start a load of the next, encrypted file with no keys.
`init_datasource_exec` already knows `encryption_enabled`, so could
`prefetch_files` also require `!encryption_enabled`? Then the rule is visible
where the scan is set up, and the `cache_enabled` filter in `get_metadata`
stays as a backstop.
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -1326,6 +1349,64 @@ mod tests {
);
}
+ /// The prefetch is opt-in: a scan of several files loads metadata ahead
only when the session
+ /// sets `prefetch_next_file_metadata`. With it, only the first file's
open reads metadata from
+ /// storage, since each later file's footer and page index were loaded
while the file before
+ /// it was read, and the batches match a scan without it.
+ #[tokio::test]
+ async fn prefetching_the_next_files_metadata_keeps_the_scan_result() {
+ let fixtures: Vec<_> = (0..3).map(|_|
write_scan_io_fixture()).collect();
+ let schema = Arc::clone(&fixtures[0].1);
+ let files: Vec<_> = fixtures
+ .into_iter()
+ .map(|(filename, _)| PartitionedFile::from_path(filename).unwrap())
+ .collect();
+
+ let mut results = vec![];
+ // `None` is a session without the options, which every Comet plan
attaches.
Review Comment:
`None` and `Some(false)` run the same code path here, because `is_some_and`
treats both as off. Could we keep only `Some(false)` and `Some(true)`? The
comment also reads as if a plan might not have the options, when every Comet
plan attaches them.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -3483,6 +3484,36 @@ 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)
+ // `None` leaves the config unset: the prefetch is opt-in.
+ Seq(Some(true), Some(false), None).foreach { prefetch =>
Review Comment:
The same goes for this test. `CometExecIterator` always sends the config's
value, and the default is `false`, so the `None` and `Some(false)` runs are the
same on the native side. Dropping one of them saves a full query run.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -750,6 +776,167 @@ 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.
+type MetadataLoad = SpawnedTask<Option<Arc<ParquetMetaData>>>;
+
+/// 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..]
+ .iter()
+ .find(|next| &next.object_meta.location != location);
+ if let Some(next) = next {
+ next_files
+ .entry((location.clone(), file.range.clone()))
+ .or_insert_with(|| (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<usize>,
+ ) -> Option<MetadataLoad> {
+ 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.
Review Comment:
No test reaches this arm.
`the_next_file_skips_further_ranges_of_the_same_path` only checks the
`next_files` map. Spark can put adjacent ranges of one file in the same task
when `spark.sql.files.maxPartitionNum` repacks partitions, so this path is
reachable. Could that test open `[a 0..50, a 50..100, b]` through the factory
instead, and assert that `b` is read once and `scan_io_metadata_prefetches` is
1? That would cover both the map and this branch.
--
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]
