andygrove commented on code in PR #7094:
URL: https://github.com/apache/datafusion-comet/pull/7094#discussion_r4308773814
##########
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:
Agreed, a counter that reports N-1 hits on a cold scan would mislead anyone
judging the cache after #7082. The prefetch task now returns its storage reads
with the metadata, and the open that uses it records a miss.
`scan_io_metadata_cache_hits` means "served from the cache" again. I updated
the `metrics.md` row, the `CometMetricNode` description, and the expected
counts in the Rust and Scala tests. (6be7a15a49)
##########
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:
Done as you suggested. `scan_io_metadata_prefetches` is now incremented in
the open that uses the load, only when the load read storage. It counts
overlapped loads, and a load nobody used shows up only in
`scan_io_metadata_bytes`. The `metrics.md` row says so. (6be7a15a49)
##########
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:
Good point. The rule shouldn't depend on `get_file_decryption_properties`
returning `Some` for every file. `prefetch_files` now also requires
`!encryption_enabled`, so an encrypted scan never builds the prefetch. The
`cache_enabled` filter in `get_metadata` stays as a backstop. (6be7a15a49)
##########
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:
You're right, the old test only covered the map. It now opens `[a 0..50, a
50..100, b]` through the factory. It asserts that `b` is read as many times as
a standalone load of `b`, so no second load starts, and that
`scan_io_metadata_prefetches` is 1. (6be7a15a49)
##########
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:
Replaced with `assert!(store.reads("bad.parquet") >= 2, ...)`, with a
message saying that both the prefetch and the open read the file. (6be7a15a49)
##########
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:
Fixed. The test runs only `false` and `true`, and the comment about a
session without the options is gone. The production check still uses
`is_some_and`, but I can make a missing extension an error if you'd prefer.
(6be7a15a49)
--
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]
