This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 0216923f [rust] Parallelize OR global-index planning across files
(#1011)
0216923f is described below
commit 0216923f649a4dcf53853cb79ddfd19ee578c8c7
Author: XiaoHongbo <[email protected]>
AuthorDate: Sat Oct 3 18:07:00 2026 +0800
[rust] Parallelize OR global-index planning across files (#1011)
---
crates/paimon/src/table/global_index_scanner.rs | 28 ++-
.../src/table/global_index_scanner/evaluator.rs | 19 +-
.../src/table/global_index_scanner/reader.rs | 5 +
.../paimon/src/table/global_index_scanner/tests.rs | 235 +++++++++++++++++++++
4 files changed, 283 insertions(+), 4 deletions(-)
diff --git a/crates/paimon/src/table/global_index_scanner.rs
b/crates/paimon/src/table/global_index_scanner.rs
index fdd21a4e..09f838c3 100644
--- a/crates/paimon/src/table/global_index_scanner.rs
+++ b/crates/paimon/src/table/global_index_scanner.rs
@@ -52,7 +52,9 @@ use std::sync::{Arc, Mutex};
use tokio::sync::Semaphore;
#[cfg(test)]
-use std::sync::atomic::{AtomicUsize as TestAtomicUsize, Ordering as
TestOrdering};
+use std::sync::atomic::{
+ AtomicBool as TestAtomicBool, AtomicUsize as TestAtomicUsize, Ordering as
TestOrdering,
+};
type BoxedCmp = Box<dyn Fn(&[u8], &[u8]) -> Result<Ordering> + Send + Sync>;
@@ -63,8 +65,13 @@ const DELETION_VECTORS_INDEX_TYPE: &str = "DELETION_VECTORS";
struct QueryIoProbe {
active: TestAtomicUsize,
peak: TestAtomicUsize,
+ evaluate_futures: TestAtomicUsize,
predicate_queries: TestAtomicUsize,
range_queries: TestAtomicUsize,
+ btree_opens: TestAtomicUsize,
+ pause_on_enter: TestAtomicBool,
+ entered: tokio::sync::Notify,
+ resume: tokio::sync::Notify,
}
#[cfg(test)]
@@ -73,6 +80,10 @@ impl QueryIoProbe {
let current = self.active.fetch_add(1, TestOrdering::SeqCst) + 1;
self.peak.fetch_max(current, TestOrdering::SeqCst);
let guard = QueryIoProbeGuard { probe: self };
+ if self.pause_on_enter.load(TestOrdering::SeqCst) {
+ self.entered.notify_one();
+ self.resume.notified().await;
+ }
tokio::task::yield_now().await;
guard
}
@@ -118,8 +129,10 @@ pub(crate) struct GlobalIndexScanner {
coverage_by_field: HashMap<i32, Vec<RowRange>>,
/// Schema fields for field_id lookup.
schema_fields: Vec<DataField>,
- /// Cache of opened BTree readers, keyed by file name.
+ /// Cache of opened BTree readers, keyed by resolved path.
reader_cache: Mutex<HashMap<String, BTreeIndexReader<BoxedCmp>>>,
+ /// Serialize use of each BTree reader across concurrent predicate
branches.
+ btree_file_locks: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
#[cfg(test)]
query_io_probe: Option<Arc<QueryIoProbe>>,
}
@@ -319,10 +332,21 @@ impl GlobalIndexScanner {
coverage_by_field,
schema_fields: schema_fields.to_vec(),
reader_cache: Mutex::new(HashMap::new()),
+ btree_file_locks: Mutex::new(HashMap::new()),
#[cfg(test)]
query_io_probe: None,
}))
}
+
+ fn btree_file_lock(&self, entry: &GlobalIndexEntry) ->
Arc<tokio::sync::Mutex<()>> {
+ let path = entry.resolved_path(&self.table_path);
+ let mut locks = self.btree_file_locks.lock().unwrap();
+ Arc::clone(
+ locks
+ .entry(path)
+ .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))),
+ )
+ }
}
/// Create a GlobalIndexScanner and evaluate predicates, returning row ranges.
diff --git a/crates/paimon/src/table/global_index_scanner/evaluator.rs
b/crates/paimon/src/table/global_index_scanner/evaluator.rs
index b2ef50e5..6af6d379 100644
--- a/crates/paimon/src/table/global_index_scanner/evaluator.rs
+++ b/crates/paimon/src/table/global_index_scanner/evaluator.rs
@@ -82,6 +82,12 @@ impl GlobalIndexScanner {
/// Evaluate a predicate against the global indexes and return matching
row ranges.
/// Returns `None` if the predicate cannot be evaluated by the global
index.
pub(super) fn evaluate<'a>(&'a self, predicate: &'a Predicate) ->
EvaluateFuture<'a> {
+ #[cfg(test)]
+ if let Some(probe) = &self.query_io_probe {
+ probe
+ .evaluate_futures
+ .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+ }
Box::pin(async move {
match predicate {
Predicate::Leaf {
@@ -239,8 +245,12 @@ impl GlobalIndexScanner {
let mut all_ranges: Vec<RowRange> = Vec::new();
let mut evaluated_field_ids = HashSet::new();
let mut indexed_coverage: Option<Vec<RowRange>> = None;
- for child in children {
- match self.evaluate(child).await? {
+ let stream = futures::stream::iter(0..children.len())
+ .map(|index| self.evaluate(&children[index]))
+ .buffered(self.global_index_thread_num);
+ futures::pin_mut!(stream);
+ while let Some(result) = stream.try_next().await? {
+ match result {
Some(child_result) => {
all_ranges.extend(child_result.row_ranges);
evaluated_field_ids.extend(child_result.evaluated_field_ids);
@@ -563,6 +573,11 @@ impl GlobalIndexScanner {
let futures =
query_plans.into_iter().map(|plan| async move {
let entry = &entries[plan.entry_idx];
+ let _file_guard = if matches!(entry.index_type,
GlobalIndexFileKind::BTree) {
+ Some(self.btree_file_lock(entry).lock_owned().await)
+ } else {
+ None
+ };
let _permit =
self.query_semaphore.acquire().await.map_err(|error| {
Error::UnexpectedError {
message: "global-index query concurrency budget was
closed".to_string(),
diff --git a/crates/paimon/src/table/global_index_scanner/reader.rs
b/crates/paimon/src/table/global_index_scanner/reader.rs
index e877176f..dec996b1 100644
--- a/crates/paimon/src/table/global_index_scanner/reader.rs
+++ b/crates/paimon/src/table/global_index_scanner/reader.rs
@@ -235,6 +235,11 @@ impl GlobalIndexScanner {
}
}
+ #[cfg(test)]
+ if let Some(probe) = &self.query_io_probe {
+ probe.btree_opens.fetch_add(1, super::TestOrdering::SeqCst);
+ }
+
// Open new reader
let input = self.file_io.new_input(&resolved_path)?;
let file_size = if entry.file_size > 0 {
diff --git a/crates/paimon/src/table/global_index_scanner/tests.rs
b/crates/paimon/src/table/global_index_scanner/tests.rs
index c28a59b3..84ebbacd 100644
--- a/crates/paimon/src/table/global_index_scanner/tests.rs
+++ b/crates/paimon/src/table/global_index_scanner/tests.rs
@@ -2831,6 +2831,241 @@ async fn
test_multi_field_and_shares_query_concurrency_budget() {
}
}
+#[tokio::test]
+async fn test_or_and_nested_predicates_share_query_budget() {
+ let src = format!(
+ "{}/testdata/btree/btree_int_100_no_compress.bin",
+ env!("CARGO_MANIFEST_DIR")
+ );
+ let tmp = tempfile::tempdir().unwrap();
+ let index_dir = tmp.path().join("index");
+ std::fs::create_dir_all(&index_dir).unwrap();
+ let table_path = format!("file://{}", tmp.path().display());
+ let file_io = crate::io::FileIOBuilder::new("file").build().unwrap();
+ let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)),
false);
+ let fields: Vec<_> = (0..4)
+ .map(|index| {
+ DataField::new(
+ index + 1,
+ format!("field{}", index + 1),
+ DataType::Int(crate::spec::IntType::new()),
+ )
+ })
+ .collect();
+ let entries: Vec<_> = (0..4)
+ .map(|index| {
+ let name = format!("index_field{index}.bin");
+ std::fs::copy(&src, index_dir.join(&name)).unwrap();
+ make_global_index_entry(&name, index + 1, 0, 99, &meta)
+ })
+ .collect();
+ let leaves: Vec<_> = (0..4)
+ .map(|index| int_eq(&format!("field{}", index + 1), index as usize,
50))
+ .collect();
+ let predicates = [
+ Predicate::or(leaves.clone()),
+ Predicate::or(vec![
+ Predicate::and(leaves[..2].to_vec()),
+ Predicate::and(leaves[2..].to_vec()),
+ ]),
+ ];
+
+ for predicate in &predicates {
+ for limit in [1, 2, 4] {
+ let mut scanner = GlobalIndexScanner::create(
+ &file_io,
+ &table_path,
+ limit,
+ i64::MAX,
+ i64::MAX,
+ &entries,
+ &fields,
+ )
+ .unwrap()
+ .unwrap();
+ let probe = Arc::new(QueryIoProbe::default());
+ scanner.query_io_probe = Some(Arc::clone(&probe));
+
+ let result = scanner.evaluate(predicate).await.unwrap().unwrap();
+ assert_eq!(result.row_ranges, vec![RowRange::new(25, 25)]);
+ assert_eq!(result.indexed_coverage, vec![RowRange::new(0, 99)]);
+ assert_eq!(result.evaluated_field_ids, HashSet::from([1, 2, 3,
4]));
+ assert_eq!(probe.peak(), limit);
+ assert_eq!(probe.btree_opens.load(TestOrdering::SeqCst), 4);
+ }
+ }
+}
+
+#[tokio::test]
+async fn test_or_reuses_one_btree_reader_for_same_file() {
+ let (io, path, name, _tmp) =
setup_testdata_table("btree_int_100_no_compress.bin");
+ let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)),
false);
+ let entries = vec![make_global_index_entry(&name, 1, 0, 99, &meta)];
+ let mut scanner = GlobalIndexScanner::create(
+ &io,
+ &path,
+ 4,
+ i64::MAX,
+ i64::MAX,
+ &entries,
+ &int_schema_fields(),
+ )
+ .unwrap()
+ .unwrap();
+ let probe = Arc::new(QueryIoProbe::default());
+ scanner.query_io_probe = Some(Arc::clone(&probe));
+ let predicate = Predicate::or(
+ [10, 20, 30, 40]
+ .map(|value| int_eq("id", 0, value))
+ .to_vec(),
+ );
+
+ let result = scanner.evaluate(&predicate).await.unwrap().unwrap();
+ assert_eq!(
+ result.row_ranges,
+ [5, 10, 15, 20].map(|row| RowRange::new(row, row)).to_vec()
+ );
+ assert_eq!(probe.peak(), 1);
+ assert_eq!(probe.btree_opens.load(TestOrdering::SeqCst), 1);
+}
+
+#[tokio::test]
+async fn test_wide_or_fallback_creates_only_bounded_child_futures() {
+ let (io, path, name, _tmp) =
setup_testdata_table("btree_int_100_no_compress.bin");
+ let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)),
false);
+ let entries = vec![make_global_index_entry(&name, 1, 0, 99, &meta)];
+ let mut scanner = GlobalIndexScanner::create(
+ &io,
+ &path,
+ 4,
+ i64::MAX,
+ i64::MAX,
+ &entries,
+ &int_schema_fields(),
+ )
+ .unwrap()
+ .unwrap();
+ let probe = Arc::new(QueryIoProbe::default());
+ scanner.query_io_probe = Some(Arc::clone(&probe));
+
+ let mut children = vec![int_eq("missing", 1, 50)];
+ children.extend((0..1024).map(|_| int_eq("id", 0, 50)));
+ assert!(scanner
+ .evaluate(&Predicate::or(children))
+ .await
+ .unwrap()
+ .is_none());
+ assert!(probe.evaluate_futures.load(TestOrdering::SeqCst) <= 5);
+}
+
+#[tokio::test]
+async fn test_or_fallback_and_cancel_leave_reader_usable() {
+ let (io, path, name, _tmp) =
setup_testdata_table("btree_int_100_no_compress.bin");
+ let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)),
false);
+ let entries = vec![make_global_index_entry(&name, 1, 0, 99, &meta)];
+ let mut scanner = GlobalIndexScanner::create(
+ &io,
+ &path,
+ 2,
+ i64::MAX,
+ i64::MAX,
+ &entries,
+ &int_schema_fields(),
+ )
+ .unwrap()
+ .unwrap();
+ let probe = Arc::new(QueryIoProbe::default());
+ scanner.query_io_probe = Some(Arc::clone(&probe));
+ let scanner = Arc::new(scanner);
+ let valid = int_eq("id", 0, 50);
+ let unsupported = int_eq("missing", 1, 50);
+
+ for position in 0..3 {
+ let mut children = vec![valid.clone(), valid.clone()];
+ children.insert(position, unsupported.clone());
+ assert!(scanner
+ .evaluate(&Predicate::or(children))
+ .await
+ .unwrap()
+ .is_none());
+ }
+
+ probe.pause_on_enter.store(true, TestOrdering::SeqCst);
+ let entered = probe.entered.notified();
+ let cancelled_scanner = Arc::clone(&scanner);
+ let cancelled_predicate = Predicate::or(vec![valid.clone(),
valid.clone()]);
+ let task = tokio::spawn(async move {
cancelled_scanner.evaluate(&cancelled_predicate).await });
+ tokio::time::timeout(std::time::Duration::from_secs(5), entered)
+ .await
+ .unwrap();
+ task.abort();
+ assert!(matches!(task.await, Err(error) if error.is_cancelled()));
+ probe.pause_on_enter.store(false, TestOrdering::SeqCst);
+
+ let result = tokio::time::timeout(std::time::Duration::from_secs(5),
scanner.evaluate(&valid))
+ .await
+ .unwrap()
+ .unwrap()
+ .unwrap();
+ assert_eq!(result.row_ranges, vec![RowRange::new(25, 25)]);
+}
+
+#[tokio::test]
+async fn test_or_index_open_error_releases_file_lock_and_budget() {
+ let tmp = tempfile::tempdir().unwrap();
+ let index_dir = tmp.path().join("index");
+ std::fs::create_dir_all(&index_dir).unwrap();
+ let table_path = format!("file://{}", tmp.path().display());
+ let file_io = crate::io::FileIOBuilder::new("file").build().unwrap();
+ let meta = BTreeIndexMeta::new(Some(le_int_key(0)), Some(le_int_key(198)),
false);
+ let entries = vec![make_global_index_entry("late.bin", 1, 0, 99, &meta)];
+ let scanner = GlobalIndexScanner::create(
+ &file_io,
+ &table_path,
+ 2,
+ i64::MAX,
+ i64::MAX,
+ &entries,
+ &int_schema_fields(),
+ )
+ .unwrap()
+ .unwrap();
+ let predicate = Predicate::or(vec![int_eq("id", 0, 50), int_eq("id", 0,
60)]);
+ let unsupported = int_eq("unindexed", 1, 50);
+
+ assert!(scanner
+ .evaluate(&Predicate::or(vec![
+ unsupported.clone(),
+ int_eq("id", 0, 50)
+ ]))
+ .await
+ .unwrap()
+ .is_none());
+ assert!(scanner
+ .evaluate(&Predicate::or(vec![int_eq("id", 0, 50), unsupported]))
+ .await
+ .is_err());
+
+ assert!(scanner.evaluate(&predicate).await.is_err());
+ let src = format!(
+ "{}/testdata/btree/btree_int_100_no_compress.bin",
+ env!("CARGO_MANIFEST_DIR")
+ );
+ std::fs::copy(src, index_dir.join("late.bin")).unwrap();
+ let result = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ scanner.evaluate(&predicate),
+ )
+ .await
+ .unwrap()
+ .unwrap()
+ .unwrap();
+ assert_eq!(
+ result.row_ranges,
+ vec![RowRange::new(25, 25), RowRange::new(30, 30)]
+ );
+}
+
/// Regression for the Between+remaining bug in `evaluate_leaf`. When a
/// native `Between` leaf is paired with another conjunct (e.g. `id >= 0`),
/// and the file's b-tree key range falls **outside** the Between range
