From fe0b78cb998ccd70121a55988c30512e0adfd67a Mon Sep 17 00:00:00 2001 From: Sushil Singh Date: Wed, 6 May 2026 10:59:34 -0400 Subject: [PATCH 1/2] perf(core): prune partition subtrees during descent and parallelize listing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `FileLister::list_relevant_partition_paths` previously walked every partition subtree to leaf depth, then filtered the results — and the recursive descent in `get_leaf_dirs` ran sequentially, paying full LIST latency at every level. On remote object stores with selective partition filters this is the dominant cost. Two changes per the re-scoped scope in issue #205: 1. `get_leaf_dirs` (`crates/core/src/storage/mod.rs`) now takes an optional per-child predicate and a parallelism bound. Subtrees the predicate rejects are skipped without listing. Sibling subdirs at each level are walked concurrently via `buffer_unordered`. 2. `PartitionPruner::should_include_prefix` evaluates filters whose field is already present in a partial partition path. Filters for deeper fields are deferred. The single `_hoodie_partition_path` field used by timestamp-based key generators conservatively admits at intermediate levels — `should_include` still runs at the leaf. `FileLister::list_relevant_partition_paths` plumbs both: the predicate filters out lake-format metadata dirs (`.hoodie`, `_delta_log`, `metadata`) and calls `should_include_prefix` for partition pruning. Parallelism uses the existing `hoodie.plan.listing.parallelism` config (default 10). Scope is the no-MDT listing path only; the metadata-table FILES path is unchanged. Closes #205. Tests: - `get_leaf_dirs_descends_in_parallel_and_returns_all_leaves` - `get_leaf_dirs_predicate_prunes_subtree` - `test_partition_pruner_should_include_prefix_*` (6 cases: empty prefix, deferred filter, evaluable filter prunes, leaf parity with `should_include`, single `_hoodie_partition_path` admits, and too-many-parts fails open) - `list_partition_paths_prunes_subtrees_when_top_level_filter_rejects` - `list_partition_paths_prunes_at_deeper_level_when_only_inner_field_filtered` --- crates/core/src/storage/mod.rs | 139 ++++++++++++++++---- crates/core/src/table/listing.rs | 95 +++++++++++--- crates/core/src/table/partition.rs | 202 +++++++++++++++++++++++++++++ 3 files changed, 399 insertions(+), 37 deletions(-) diff --git a/crates/core/src/storage/mod.rs b/crates/core/src/storage/mod.rs index e7e7cad1..e69fe24f 100644 --- a/crates/core/src/storage/mod.rs +++ b/crates/core/src/storage/mod.rs @@ -25,10 +25,11 @@ use std::sync::Arc; use arrow::compute::concat_batches; use arrow::record_batch::RecordBatch; use arrow_schema::SchemaRef; -use async_recursion::async_recursion; use bytes::Bytes; use futures::StreamExt; -use futures::stream::BoxStream; +use futures::TryStreamExt; +use futures::future::BoxFuture; +use futures::stream::{self, BoxStream}; use object_store::path::Path as ObjPath; use object_store::{ObjectStore, parse_url_opts}; use parquet::arrow::async_reader::ParquetObjectReader; @@ -427,36 +428,90 @@ impl Storage { } } +/// Predicate evaluated for each child directory during recursive descent. +/// +/// Returning `false` prunes the entire subtree rooted at that child — no +/// further `list_dirs` calls happen under it. The path is relative to the +/// storage base (the same form returned by [`get_leaf_dirs`]). +pub type LeafDirPredicate<'a> = &'a (dyn Fn(&str) -> bool + Send + Sync); + /// Get relative paths of leaf directories under a given directory. /// +/// Walks subdirectories in parallel up to `parallelism` concurrent +/// `list_dirs` calls. An optional `predicate` prunes subtrees that cannot +/// match (e.g., partition pruning during file listing); subtrees rejected +/// by the predicate are skipped without listing their contents. +/// +/// # Parameters +/// - `subdir`: starting directory relative to the storage base. `None` walks +/// from the root. +/// - `predicate`: if `Some`, called with each child path; subtrees where the +/// predicate returns `false` are pruned. +/// - `parallelism`: maximum concurrent `list_dirs` calls. Values `< 1` are +/// clamped to 1 (sequential descent). +/// /// **Example** /// - /usr/hudi/table_name /// - /usr/hudi/table_name/.hoodie /// - /usr/hudi/table_name/dt=2024/month=01/day=01 /// - /usr/hudi/table_name/dt=2025/month=02 /// -/// the result is \[".hoodie", "dt=2024/mont=01/day=01", "dt=2025/month=02"\] -#[async_recursion] -pub async fn get_leaf_dirs(storage: &Storage, subdir: Option<&str>) -> Result> { - let mut leaf_dirs = Vec::new(); - let child_dirs = storage.list_dirs(subdir).await?; - if child_dirs.is_empty() { - leaf_dirs.push(subdir.unwrap_or_default().to_owned()); - } else { +/// the result is \[".hoodie", "dt=2024/month=01/day=01", "dt=2025/month=02"\] +pub fn get_leaf_dirs<'a>( + storage: &'a Storage, + subdir: Option<&'a str>, + predicate: Option>, + parallelism: usize, +) -> BoxFuture<'a, Result>> { + get_leaf_dirs_inner( + storage, + subdir.map(String::from), + predicate, + parallelism.max(1), + ) +} + +fn get_leaf_dirs_inner<'a>( + storage: &'a Storage, + subdir: Option, + predicate: Option>, + parallelism: usize, +) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let child_dirs = storage.list_dirs(subdir.as_deref()).await?; + if child_dirs.is_empty() { + return Ok(vec![subdir.unwrap_or_default()]); + } + + let mut next_subdirs: Vec = Vec::with_capacity(child_dirs.len()); for child_dir in child_dirs { - let mut next_subdir = PathBuf::new(); - if let Some(curr) = subdir { - next_subdir.push(curr); + let mut path = PathBuf::new(); + if let Some(curr) = subdir.as_deref() { + path.push(curr); } - next_subdir.push(child_dir); - let next_subdir = next_subdir + path.push(&child_dir); + let next = path .to_str() - .ok_or_else(|| InvalidPath(format!("Failed to convert path: {next_subdir:?}")))?; - let curr_leaf_dir = get_leaf_dirs(storage, Some(next_subdir)).await?; - leaf_dirs.extend(curr_leaf_dir); + .ok_or_else(|| InvalidPath(format!("Failed to convert path: {path:?}")))? + .to_owned(); + if predicate.map(|p| p(&next)).unwrap_or(true) { + next_subdirs.push(next); + } } - } - Ok(leaf_dirs) + + if next_subdirs.is_empty() { + return Ok(Vec::new()); + } + + stream::iter(next_subdirs) + .map(move |path| get_leaf_dirs_inner(storage, Some(path), predicate, parallelism)) + .buffer_unordered(parallelism) + .try_fold(Vec::new(), |mut acc, mut sub| async move { + acc.append(&mut sub); + Ok::<_, StorageError>(acc) + }) + .await + }) } #[cfg(test)] @@ -596,7 +651,8 @@ mod tests { ) .unwrap(); let storage = Storage::new_with_base_url(base_url).unwrap(); - let leaf_dirs = get_leaf_dirs(&storage, None).await.unwrap(); + let mut leaf_dirs = get_leaf_dirs(&storage, None, None, 1).await.unwrap(); + leaf_dirs.sort(); assert_eq!( leaf_dirs, vec![".hoodie", "part1", "part2/part22", "part3/part32/part33"] @@ -609,7 +665,7 @@ mod tests { Url::from_directory_path(canonicalize(Path::new("tests/data/leaf_dir")).unwrap()) .unwrap(); let storage = Storage::new_with_base_url(base_url).unwrap(); - let leaf_dirs = get_leaf_dirs(&storage, None).await.unwrap(); + let leaf_dirs = get_leaf_dirs(&storage, None, None, 1).await.unwrap(); assert_eq!( leaf_dirs, vec![""], @@ -617,6 +673,45 @@ mod tests { ); } + #[tokio::test] + async fn get_leaf_dirs_descends_in_parallel_and_returns_all_leaves() { + let base_url = Url::from_directory_path( + canonicalize(Path::new("tests/data/timeline/commits_stub")).unwrap(), + ) + .unwrap(); + let storage = Storage::new_with_base_url(base_url).unwrap(); + let mut leaf_dirs = get_leaf_dirs(&storage, None, None, 8).await.unwrap(); + leaf_dirs.sort(); + assert_eq!( + leaf_dirs, + vec![".hoodie", "part1", "part2/part22", "part3/part32/part33"], + "Parallel descent must yield the same set of leaves as sequential." + ); + } + + #[tokio::test] + async fn get_leaf_dirs_predicate_prunes_subtree() { + let base_url = Url::from_directory_path( + canonicalize(Path::new("tests/data/timeline/commits_stub")).unwrap(), + ) + .unwrap(); + let storage = Storage::new_with_base_url(base_url).unwrap(); + let predicate = |path: &str| -> bool { + let top = path.split('/').next().unwrap_or(""); + top != ".hoodie" && top != "part3" + }; + let predicate_ref: LeafDirPredicate<'_> = &predicate; + let mut leaf_dirs = get_leaf_dirs(&storage, None, Some(predicate_ref), 4) + .await + .unwrap(); + leaf_dirs.sort(); + assert_eq!( + leaf_dirs, + vec!["part1", "part2/part22"], + "Subtrees rejected by the predicate must not appear in the output." + ); + } + #[tokio::test] async fn storage_get_file_info() { let base_url = diff --git a/crates/core/src/table/listing.rs b/crates/core/src/table/listing.rs index 1a2916c5..302f2374 100644 --- a/crates/core/src/table/listing.rs +++ b/crates/core/src/table/listing.rs @@ -26,7 +26,7 @@ use crate::file_group::base_file::BaseFile; use crate::file_group::log_file::LogFile; use crate::metadata::LAKE_FORMAT_METADATA_DIRS; use crate::statistics::estimator::FileStatsEstimator; -use crate::storage::{Storage, get_leaf_dirs}; +use crate::storage::{LeafDirPredicate, Storage, get_leaf_dirs}; use crate::table::partition::{ EMPTY_PARTITION_PATH, PARTITION_METAFIELD_PREFIX, PartitionPruner, is_table_partitioned, }; @@ -168,26 +168,36 @@ impl FileLister { return Ok(vec![EMPTY_PARTITION_PATH.to_string()]); } - let top_level_dirs: Vec = self - .storage - .list_dirs(None) - .await? - .into_iter() - .filter(|dir| !LAKE_FORMAT_METADATA_DIRS.contains(&dir.as_str())) - .collect(); + let parallelism: usize = self.hudi_configs.get_or_default(ListingParallelism).into(); + let pruner = &self.partition_pruner; + let pruner_is_empty = pruner.is_empty(); + let predicate = move |path: &str| -> bool { + // Skip table-format metadata dirs (.hoodie, _delta_log, metadata) + // anywhere in the descent — they live at the table root in + // practice, so a top-level segment match is sufficient. + let top = path.split('/').next().unwrap_or(""); + if LAKE_FORMAT_METADATA_DIRS.contains(&top) { + return false; + } + if pruner_is_empty { + return true; + } + pruner.should_include_prefix(path) + }; + let predicate_ref: LeafDirPredicate<'_> = &predicate; - let mut partition_paths = Vec::new(); - for dir in top_level_dirs { - partition_paths.extend(get_leaf_dirs(&self.storage, Some(&dir)).await?); - } + let partition_paths = + get_leaf_dirs(&self.storage, None, Some(predicate_ref), parallelism).await?; - if partition_paths.is_empty() || self.partition_pruner.is_empty() { + if pruner_is_empty { return Ok(partition_paths); } - + // Final leaf filter for cases the prefix predicate must conservatively + // admit (single _hoodie_partition_path field with timestamp-based + // keygen): the filter compares the full leaf path string. Ok(partition_paths .into_iter() - .filter(|path_str| self.partition_pruner.should_include(path_str)) + .filter(|p| pruner.should_include(p)) .collect()) } @@ -240,6 +250,7 @@ impl FileLister { #[cfg(test)] mod test { use super::*; + use crate::expr::filter::Filter; use crate::table::Table; use hudi_test::SampleTable; use std::collections::HashSet; @@ -281,4 +292,58 @@ mod test { ]) ) } + + #[tokio::test] + async fn list_partition_paths_prunes_subtrees_when_top_level_filter_rejects() { + let base_url = SampleTable::V6ComplexkeygenHivestyle.url_to_cow(); + let hudi_table = Table::new(base_url.path()).await.unwrap(); + let fs_view = &hudi_table.file_system_view; + let partition_schema = hudi_table.get_partition_schema().await.unwrap(); + let pruner = PartitionPruner::new( + &[Filter::try_from(("byteField", "<", "20")).unwrap()], + &partition_schema, + hudi_table.hudi_configs.as_ref(), + ) + .unwrap(); + let lister = FileLister::new( + fs_view.hudi_configs.clone(), + fs_view.storage.clone(), + pruner, + ); + let partition_paths = lister.list_relevant_partition_paths().await.unwrap(); + let partition_path_set: HashSet<&str> = + HashSet::from_iter(partition_paths.iter().map(|p| p.as_str())); + // byteField=20 and byteField=30 should be pruned at level 1 — their + // shortField subdirs must never be listed. + assert_eq!( + partition_path_set, + HashSet::from_iter(vec!["byteField=10/shortField=300"]) + ) + } + + #[tokio::test] + async fn list_partition_paths_prunes_at_deeper_level_when_only_inner_field_filtered() { + let base_url = SampleTable::V6ComplexkeygenHivestyle.url_to_cow(); + let hudi_table = Table::new(base_url.path()).await.unwrap(); + let fs_view = &hudi_table.file_system_view; + let partition_schema = hudi_table.get_partition_schema().await.unwrap(); + let pruner = PartitionPruner::new( + &[Filter::try_from(("shortField", "=", "300")).unwrap()], + &partition_schema, + hudi_table.hudi_configs.as_ref(), + ) + .unwrap(); + let lister = FileLister::new( + fs_view.hudi_configs.clone(), + fs_view.storage.clone(), + pruner, + ); + let partition_paths = lister.list_relevant_partition_paths().await.unwrap(); + let partition_path_set: HashSet<&str> = + HashSet::from_iter(partition_paths.iter().map(|p| p.as_str())); + assert_eq!( + partition_path_set, + HashSet::from_iter(vec!["byteField=10/shortField=300"]) + ) + } } diff --git a/crates/core/src/table/partition.rs b/crates/core/src/table/partition.rs index 6a4a1f75..bf6bc2f2 100644 --- a/crates/core/src/table/partition.rs +++ b/crates/core/src/table/partition.rs @@ -174,6 +174,98 @@ impl PartitionPruner { }) } + /// Returns `true` if a partition path *prefix* could lead to an included + /// partition. + /// + /// Used during recursive directory listing to short-circuit subtrees whose + /// partition columns have already been ruled out: only filters whose field + /// is present in the prefix are evaluated; filters for deeper fields are + /// deferred to lower levels. + /// + /// Returns `true` for the empty prefix (root) and on parsing errors — + /// callers must still apply [`Self::should_include`] at the leaf to filter + /// cases this method conservatively admits (notably the single + /// `_hoodie_partition_path` field used by timestamp-based key generators, + /// where intermediate-level string comparison would be incorrect). + pub fn should_include_prefix(&self, partition_path_prefix: &str) -> bool { + if partition_path_prefix.is_empty() { + return true; + } + let segments = match self.parse_prefix_segments(partition_path_prefix) { + Ok(s) => s, + Err(_) => return true, + }; + + self.and_filters.iter().all(|filter| { + match segments.get(filter.field.name()) { + Some(segment_value) => match filter.apply_comparison(segment_value) { + Ok(scalar) => scalar.value(0), + Err(_) => true, + }, + None => true, // Filter's field not yet in the prefix — defer. + } + }) + } + + fn parse_prefix_segments( + &self, + partition_path_prefix: &str, + ) -> Result>> { + let partition_path = if self.is_url_encoded { + percent_encoding::percent_decode(partition_path_prefix.as_bytes()) + .decode_utf8()? + .into_owned() + } else { + partition_path_prefix.to_string() + }; + + // Single _hoodie_partition_path field (timestamp-based key generators): + // the filter compares the full path string; we cannot meaningfully + // truncate it at intermediate levels without knowing the leaf depth, so + // include the subtree and let `should_include` handle the leaf. + if self.schema.fields().len() == 1 + && self.schema.field(0).name() == MetaField::PartitionPath.as_ref() + { + return Ok(HashMap::new()); + } + + let parts: Vec<&str> = partition_path.split('/').collect(); + + if parts.len() > self.schema.fields().len() { + return Err(InvalidPartitionPath(format!( + "Partition path prefix has {} part(s) but schema has only {} field(s)", + parts.len(), + self.schema.fields().len() + ))); + } + + self.schema + .fields() + .iter() + .take(parts.len()) + .zip(parts) + .map(|(field, part)| { + let value = if self.is_hive_style { + let (name, value) = part.split_once('=').ok_or(InvalidPartitionPath( + format!("Partition path should be hive-style but got {part}"), + ))?; + if name != field.name() { + return Err(InvalidPartitionPath(format!( + "Partition path should contain {} but got {}", + field.name(), + name + ))); + } + value + } else { + part + }; + let scalar = SchemableFilter::cast_value(&[value], field.data_type())?; + Ok((field.name().to_string(), scalar)) + }) + .collect() + } + /// Transforms user filters on data columns to filters on partition path columns /// based on the configured key generator. fn transform_filters_for_keygen( @@ -401,6 +493,116 @@ mod tests { assert!(!pruner.should_include("date=2023-02-01/category=B/count=10")); } + #[test] + fn test_partition_pruner_should_include_prefix_admits_root() { + let schema = create_test_schema(); + let configs = create_hudi_configs(true, false); + let filter = Filter::try_from(("category", "=", "A")).unwrap(); + let pruner = PartitionPruner::new(&[filter], &schema, &configs).unwrap(); + assert!( + pruner.should_include_prefix(""), + "An empty prefix must be included so root descent proceeds." + ); + } + + #[test] + fn test_partition_pruner_should_include_prefix_defers_when_field_not_yet_in_prefix() { + let schema = create_test_schema(); + let configs = create_hudi_configs(true, false); + let filter = Filter::try_from(("count", "<=", "10")).unwrap(); + let pruner = PartitionPruner::new(&[filter], &schema, &configs).unwrap(); + assert!( + pruner.should_include_prefix("date=2023-02-01"), + "Filter on `count` cannot evaluate at level 1; must include and defer." + ); + assert!( + pruner.should_include_prefix("date=2023-02-01/category=A"), + "Filter on `count` cannot evaluate at level 2; must include and defer." + ); + } + + #[test] + fn test_partition_pruner_should_include_prefix_prunes_when_filter_evaluable() { + let schema = create_test_schema(); + let configs = create_hudi_configs(true, false); + let filter_eq_a = Filter::try_from(("category", "=", "A")).unwrap(); + let pruner = PartitionPruner::new(&[filter_eq_a], &schema, &configs).unwrap(); + assert!(pruner.should_include_prefix("date=2023-02-01/category=A")); + assert!(!pruner.should_include_prefix("date=2023-02-01/category=B")); + } + + #[test] + fn test_partition_pruner_should_include_prefix_at_leaf_matches_should_include() { + let schema = create_test_schema(); + let configs = create_hudi_configs(true, false); + let filter_gt_date = Filter::try_from(("date", ">", "2023-01-01")).unwrap(); + let filter_eq_a = Filter::try_from(("category", "=", "A")).unwrap(); + let filter_lte_100 = Filter::try_from(("count", "<=", "100")).unwrap(); + let pruner = PartitionPruner::new( + &[filter_gt_date, filter_eq_a, filter_lte_100], + &schema, + &configs, + ) + .unwrap(); + + let leaf_in = "date=2023-02-01/category=A/count=10"; + let leaf_out = "date=2022-12-31/category=A/count=10"; + assert_eq!( + pruner.should_include_prefix(leaf_in), + pruner.should_include(leaf_in) + ); + assert_eq!( + pruner.should_include_prefix(leaf_out), + pruner.should_include(leaf_out) + ); + } + + #[test] + fn test_partition_pruner_should_include_prefix_admits_single_partition_path_field() { + // Single _hoodie_partition_path (timestamp-based keygen): prefix must + // fail-open so leaf-level `should_include` makes the final decision. + let configs = HudiConfigs::new([ + ("hoodie.table.partition.fields", "ts"), + ( + "hoodie.table.keygenerator.class", + "org.apache.hudi.keygen.TimestampBasedKeyGenerator", + ), + ("hoodie.keygen.timebased.timestamp.type", "DATE_STRING"), + ( + "hoodie.keygen.timebased.input.dateformat", + "yyyy-MM-dd'T'HH:mm:ssZ", + ), + ("hoodie.keygen.timebased.output.dateformat", "yyyy/MM/dd"), + ("hoodie.datasource.write.hive_style_partitioning", "true"), + ]); + let partition_schema = Schema::new(vec![Field::new( + MetaField::PartitionPath.as_ref(), + DataType::Utf8, + false, + )]); + let user_filter = Filter { + field: "ts".to_string(), + operator: ExprOperator::Gte, + values: vec!["2024-01-15T00:00:00Z".to_string()], + }; + let pruner = PartitionPruner::new(&[user_filter], &partition_schema, &configs).unwrap(); + assert!(pruner.should_include_prefix("year=2023")); + assert!(pruner.should_include_prefix("year=2023/month=12")); + assert!(pruner.should_include_prefix("year=2024/month=01")); + } + + #[test] + fn test_partition_pruner_should_include_prefix_too_many_parts_fails_open() { + let schema = create_test_schema(); + let configs = create_hudi_configs(true, false); + let pruner = PartitionPruner::new(&[], &schema, &configs).unwrap(); + // 4 parts vs 3-field schema — parse error, fail-open. + assert!( + pruner.should_include_prefix("date=2023-02-01/category=A/count=10/extra=1"), + "Over-deep prefix must fail-open rather than incorrectly prune." + ); + } + #[test] fn test_partition_pruner_parse_segments() { let schema = create_test_schema(); From 44e2a5fb547254f058fb0d8b61b49d6eaca0521e Mon Sep 17 00:00:00 2001 From: Sushil Singh Date: Wed, 6 May 2026 13:30:51 -0400 Subject: [PATCH 2/2] fix(core): cap get_leaf_dirs concurrency globally with a semaphore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `buffer_unordered(parallelism)` only bounds fan-out per recursion level. A recursive descent over a D-level partition tree could therefore reach up to P^D concurrent `list_dirs` calls — for the default `parallelism=10` and a 3-level hive-style table, that is 1000 simultaneous LIST requests; for a date/hour/minute/second layout, 10000. Replace the per-level cap with a `tokio::sync::Semaphore` shared across the whole descent. Each recursive call acquires a permit before issuing `list_dirs` and releases it before fanning out, so total in-flight requests stay bounded by `parallelism` regardless of depth. `buffer_unordered` is kept on the recursive stream — it now bounds memory of pending child futures, while the semaphore bounds actual I/O. `hoodie.plan.listing.parallelism` becomes a meaningful global knob: the default 10 is conservative for shallow trees but can be raised for wider or deeper layouts without risking exponential concurrency growth. --- crates/core/src/storage/mod.rs | 37 ++++++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/crates/core/src/storage/mod.rs b/crates/core/src/storage/mod.rs index e69fe24f..50ace651 100644 --- a/crates/core/src/storage/mod.rs +++ b/crates/core/src/storage/mod.rs @@ -35,6 +35,7 @@ use object_store::{ObjectStore, parse_url_opts}; use parquet::arrow::async_reader::ParquetObjectReader; use parquet::arrow::{ParquetRecordBatchStreamBuilder, parquet_to_arrow_schema}; use parquet::file::metadata::ParquetMetaData; +use tokio::sync::Semaphore; use url::Url; use crate::config::HudiConfigs; @@ -437,18 +438,21 @@ pub type LeafDirPredicate<'a> = &'a (dyn Fn(&str) -> bool + Send + Sync); /// Get relative paths of leaf directories under a given directory. /// -/// Walks subdirectories in parallel up to `parallelism` concurrent -/// `list_dirs` calls. An optional `predicate` prunes subtrees that cannot -/// match (e.g., partition pruning during file listing); subtrees rejected -/// by the predicate are skipped without listing their contents. +/// Walks subdirectories concurrently. Total in-flight `list_dirs` calls are +/// globally bounded by `parallelism` via a [`Semaphore`] shared across the +/// whole descent — recursion does not multiply concurrency. An optional +/// `predicate` prunes subtrees that cannot match (e.g., partition pruning +/// during file listing); subtrees rejected by the predicate are skipped +/// without listing their contents. /// /// # Parameters /// - `subdir`: starting directory relative to the storage base. `None` walks /// from the root. /// - `predicate`: if `Some`, called with each child path; subtrees where the /// predicate returns `false` are pruned. -/// - `parallelism`: maximum concurrent `list_dirs` calls. Values `< 1` are -/// clamped to 1 (sequential descent). +/// - `parallelism`: global upper bound on concurrent `list_dirs` calls +/// across all recursion levels. Values `< 1` are clamped to 1 (sequential +/// descent). /// /// **Example** /// - /usr/hudi/table_name @@ -463,11 +467,14 @@ pub fn get_leaf_dirs<'a>( predicate: Option>, parallelism: usize, ) -> BoxFuture<'a, Result>> { + let parallelism = parallelism.max(1); + let sem = Arc::new(Semaphore::new(parallelism)); get_leaf_dirs_inner( storage, subdir.map(String::from), predicate, - parallelism.max(1), + parallelism, + sem, ) } @@ -476,9 +483,19 @@ fn get_leaf_dirs_inner<'a>( subdir: Option, predicate: Option>, parallelism: usize, + sem: Arc, ) -> BoxFuture<'a, Result>> { Box::pin(async move { + // Acquire permit before issuing the list_dirs; release before + // recursing so children compete for the same global budget. + let permit = match sem.clone().acquire_owned().await { + Ok(p) => p, + // The semaphore is private to this descent and never closed. + Err(_) => unreachable!("get_leaf_dirs semaphore is owned by the descent"), + }; let child_dirs = storage.list_dirs(subdir.as_deref()).await?; + drop(permit); + if child_dirs.is_empty() { return Ok(vec![subdir.unwrap_or_default()]); } @@ -503,8 +520,12 @@ fn get_leaf_dirs_inner<'a>( return Ok(Vec::new()); } + // The semaphore caps actual list_dirs concurrency; buffer_unordered + // here only bounds memory of pending recursive futures per level. stream::iter(next_subdirs) - .map(move |path| get_leaf_dirs_inner(storage, Some(path), predicate, parallelism)) + .map(move |path| { + get_leaf_dirs_inner(storage, Some(path), predicate, parallelism, sem.clone()) + }) .buffer_unordered(parallelism) .try_fold(Vec::new(), |mut acc, mut sub| async move { acc.append(&mut sub);