(paimon-rust) branch main updated: [btree] Add support for prefix range scan on btree global index (#408)

[lzljs3620320]

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 bc439d7  [btree] Add support for prefix range scan on btree global 
index (#408)
bc439d7 is described below

commit bc439d7f15bbf7c925318a45225060160629de72
Author: Arnav Balyan <60175178+arnavbal...@users.noreply.github.com>
AuthorDate: Wed Jun 24 19:43:38 2026 -0700

    [btree] Add support for prefix range scan on btree global index (#408)
---
 crates/paimon/src/btree/reader.rs | 19 +++++++++++++++++++
 crates/paimon/src/btree/tests.rs  | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/crates/paimon/src/btree/reader.rs 
b/crates/paimon/src/btree/reader.rs
index e255728..6de0c3a 100644
--- a/crates/paimon/src/btree/reader.rs
+++ b/crates/paimon/src/btree/reader.rs
@@ -255,6 +255,25 @@ impl<F: Fn(&[u8], &[u8]) -> Ordering> BTreeIndexReader<F> {
         }
     }
 
+    pub async fn query_prefix(&self, prefix: &[u8]) -> 
io::Result<RoaringTreemap> {
+        match Self::prefix_successor(prefix) {
+            Some(upper) => self.range_query(prefix, &upper, true, false).await,
+            None => self.query_greater_or_equal(prefix).await,
+        }
+    }
+
+    fn prefix_successor(prefix: &[u8]) -> Option<Vec<u8>> {
+        let mut bound = prefix.to_vec();
+        while let Some(&last) = bound.last() {
+            if last != 0xFF {
+                *bound.last_mut().unwrap() = last + 1;
+                return Some(bound);
+            }
+            bound.pop();
+        }
+        None
+    }
+
     /// Between query (inclusive on both ends).
     pub async fn query_between(&self, from: &[u8], to: &[u8]) -> 
io::Result<RoaringTreemap> {
         self.range_query(from, to, true, true).await
diff --git a/crates/paimon/src/btree/tests.rs b/crates/paimon/src/btree/tests.rs
index 94349e3..0ef9687 100644
--- a/crates/paimon/src/btree/tests.rs
+++ b/crates/paimon/src/btree/tests.rs
@@ -183,6 +183,40 @@ async fn test_range_queries() {
     assert_eq!(bm.len(), 5);
 }
 
+#[tokio::test]
+async fn test_prefix_query() {
+    let buf = VecFileWrite::new();
+    let mut writer = BTreeIndexWriter::new(Box::new(buf.clone()), 64, 
BlockCompressionType::None);
+
+    let keys = ["apple", "apply", "apricot", "banana", "band", "bee"];
+    for (i, k) in keys.iter().enumerate() {
+        writer.write(Some(k.as_bytes()), i as i64).await.unwrap();
+    }
+
+    let result = writer.finish().await.unwrap();
+    let reader = write_and_open(&buf, &result, |a: &[u8], b: &[u8]| 
a.cmp(b)).await;
+
+    let bm = reader.query_prefix(b"ap").await.unwrap();
+    assert_eq!(bm.len(), 3);
+    assert!(bm.contains(0) && bm.contains(1) && bm.contains(2));
+
+    let bm = reader.query_prefix(b"app").await.unwrap();
+    assert_eq!(bm.len(), 2);
+    assert!(bm.contains(0) && bm.contains(1));
+
+    let bm = reader.query_prefix(b"ban").await.unwrap();
+    assert_eq!(bm.len(), 2);
+    assert!(bm.contains(3) && bm.contains(4));
+
+    let bm = reader.query_prefix(b"b").await.unwrap();
+    assert_eq!(bm.len(), 3);
+
+    let bm = reader.query_prefix(b"z").await.unwrap();
+    assert_eq!(bm.len(), 0);
+    let bm = reader.query_prefix(b"").await.unwrap();
+    assert_eq!(bm.len(), keys.len() as u64);
+}
+
 #[tokio::test]
 async fn test_not_equal_query() {
     let buf = VecFileWrite::new();