From 53acf989526d05c64bce16225f9f00f4054793d6 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 10 Oct 2026 14:57:37 +0300 Subject: [PATCH] Bound ListingTable metadata requests by configured concurrency --- datafusion/catalog-listing/src/table.rs | 65 ++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 8 deletions(-) diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index e0c349e04ea67..f65b8db1d37d4 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -57,8 +57,23 @@ use datafusion_physical_plan::empty::EmptyExec; use futures::{Stream, StreamExt, TryStreamExt, future, stream}; use object_store::ObjectStore; use std::collections::{HashMap, HashSet}; +use std::future::Future; use std::sync::Arc; +async fn bounded_try_join_all( + futures: I, + concurrency: usize, +) -> Result, E> +where + I: IntoIterator, + F: Future>, +{ + stream::iter(futures) + .buffered(concurrency) + .try_collect() + .await +} + /// Result of a file listing operation from [`ListingTable::list_files_for_scan`]. #[derive(Debug)] pub struct ListFilesResult { @@ -806,9 +821,12 @@ impl ListingTable { listing_time_filters: &'a [Expr], file_limit: Option, ) -> datafusion_common::Result<(FileGroup, bool)> { - // list files (with partitions) - let file_list = future::try_join_all(self.table_paths.iter().map(|table_path| { - pruned_partition_list( + // Bound metadata requests for explicit file lists as well as directories. + let meta_fetch_concurrency = + ctx.config_options().execution.meta_fetch_concurrency.get(); + let mut requests = Vec::with_capacity(self.table_paths.len()); + for table_path in &self.table_paths { + requests.push(pruned_partition_list( ctx, store.as_ref(), table_path, @@ -816,11 +834,9 @@ impl ListingTable { &self.options.file_extension, self.options.auto_file_suffixes.as_deref(), &self.options.table_partition_cols, - ) - })) - .await?; - let meta_fetch_concurrency = - ctx.config_options().execution.meta_fetch_concurrency.get(); + )); + } + let file_list = bounded_try_join_all(requests, meta_fetch_concurrency).await?; // Table paths can overlap, for example when one path is a directory and // another names a file inside it. A ListingTable uses one object store, // so the object path uniquely identifies a file within this scan. @@ -1132,6 +1148,39 @@ mod tests { use arrow::compute::SortOptions; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::Poll; + + #[test] + fn test_bounded_metadata_requests_keep_order() { + let active = Arc::new(AtomicUsize::new(0)); + let maximum = Arc::new(AtomicUsize::new(0)); + let requests = (0..64).map(|index| { + let active = Arc::clone(&active); + let maximum = Arc::clone(&maximum); + async move { + let current = active.fetch_add(1, Ordering::SeqCst) + 1; + maximum.fetch_max(current, Ordering::SeqCst); + let mut first_poll = true; + future::poll_fn(|cx| { + if first_poll { + first_poll = false; + cx.waker().wake_by_ref(); + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; + active.fetch_sub(1, Ordering::SeqCst); + Ok::<_, ()>(index) + } + }); + let result = futures::executor::block_on(bounded_try_join_all(requests, 8)); + assert_eq!(result, Ok((0..64).collect())); + assert_eq!(maximum.load(Ordering::SeqCst), 8); + assert_eq!(active.load(Ordering::SeqCst), 0); + } /// Helper to create a PhysicalSortExpr fn sort_expr(