andygrove commented on code in PR #7094:
URL: https://github.com/apache/datafusion-comet/pull/7094#discussion_r4326867276
##########
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;
Review Comment:
I'd keep it a constant. `CometParquetScanOptions` holds per-plan settings
that come from Comet configs, and this one isn't configurable. It's also shared
with `delta_dv.rs`, whose cache-populating footer fetch has to issue the same
read the scan would, and a fixed constant guarantees that. If we want it
tunable later, it could move there with a config behind it.
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -1496,4 +1700,367 @@ 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 {
Review Comment:
Merged them into one `RecordingStore` that counts gets and range calls per
path, with an optional gated path set by `with_gate`. `get_calls` and
`range_calls` are now sums over paths. The local reader loads metadata through
`get_ranges`, so range calls are counted and gated per path too. `open` is
renamed `release`. (c73321f34)
##########
native/core/src/parquet/eager_page_index_reader_factory.rs:
##########
@@ -1496,4 +1700,367 @@ 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. 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(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"), 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(GatedStore::new("unused.parquet"));
+ 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();
+ }
+
+ // The ranges of a share one load. The second range finds it already
in the cache.
+ let next = |file: &PartitionedFile| {
+ let prefetch = factory.next_file_prefetch.as_ref().unwrap();
+ 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.parquet"));
+ assert_eq!(next(&files[1]).as_deref(), Some("b.parquet"));
+ assert_eq!(next(&files[2]), None);
Review Comment:
Removed the `next` closure and the `next_files` assertions. The read count
of `b` and `scan_io_metadata_prefetches == 1` cover a wrong successor, and
`the_next_file_map_handles_repeated_paths_and_long_runs` keeps the map shape.
(c73321f34)
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -1326,6 +1351,59 @@ mod tests {
);
}
+ /// A scan of several files loads metadata ahead only when the session sets
+ /// `prefetch_next_file_metadata`, which the JVM side does by default.
With it, each later
+ /// file's footer and page index are loaded while the file before it is
read, and the batches
+ /// match a scan without it.
+ #[tokio::test]
+ async fn prefetching_the_next_files_metadata_keeps_the_scan_result() {
Review Comment:
Removed. `init_test_scan` takes a single `PartitionedFile` again, so the
five call-site edits are gone too. (c73321f34)
##########
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 setting unset, which prefetches by default.
+ Seq(Some(true), Some(false), None).foreach { setting =>
Review Comment:
Fixed, it runs `Seq(true, false)`. (c73321f34)
--
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]
