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 d5229495 perf(btree): optimize sparse IN queries (#868)
d5229495 is described below
commit d5229495a206e49a9ad4b5c397a506937baddc22
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Sep 18 09:17:11 2026 +0800
perf(btree): optimize sparse IN queries (#868)
---
crates/paimon/src/btree/reader.rs | 111 +++++++++++++-------------------------
crates/paimon/src/btree/tests.rs | 100 ++++++++++++++++++++++++++++++++++
2 files changed, 137 insertions(+), 74 deletions(-)
diff --git a/crates/paimon/src/btree/reader.rs
b/crates/paimon/src/btree/reader.rs
index f9836332..11d48ebd 100644
--- a/crates/paimon/src/btree/reader.rs
+++ b/crates/paimon/src/btree/reader.rs
@@ -309,103 +309,66 @@ impl<F: Fn(&[u8], &[u8]) -> Ordering>
BTreeIndexReader<F> {
self.range_query(from, to, true, true).await
}
- /// In query: sort keys and do a single sequential scan (merge-join style).
+ /// IN query: group sorted target keys by data block, then seek each key
within its block.
+ /// This reads every target block at most once and skips unrelated blocks
and entries between
+ /// sparse target keys.
pub async fn query_in(&self, keys: &[&[u8]]) -> io::Result<RoaringTreemap>
{
if keys.is_empty() {
return Ok(RoaringTreemap::new());
}
let cmp = &self.key_comparator;
+ let (Some(min_key), Some(max_key)) = (self.min_key.as_deref(),
self.max_key.as_deref())
+ else {
+ return Ok(RoaringTreemap::new());
+ };
- // Sort query keys
+ // Sort, deduplicate, and discard keys outside this file's bounds
before resolving blocks.
let mut sorted_keys: Vec<&[u8]> = keys.to_vec();
sorted_keys.sort_by(|a, b| cmp(a, b));
sorted_keys.dedup_by(|a, b| cmp(a, b) == Ordering::Equal);
+ sorted_keys.retain(|key| {
+ cmp(key, min_key) != Ordering::Less && cmp(key, max_key) !=
Ordering::Greater
+ });
+ if sorted_keys.is_empty() {
+ return Ok(RoaringTreemap::new());
+ }
- let mut result = RoaringTreemap::new();
- let mut key_idx = 0;
-
- // Seek in index block to the first data block
+ // The index block is already resident in memory. Resolve every target
key to its first
+ // possible data block and coalesce adjacent targets that share the
same block handle.
let index_block = self.sst_reader.index_block();
- let (_, mut index_iter) = index_block.seek_and_iter(sorted_keys[0],
cmp);
+ let mut target_blocks: Vec<(BlockHandle, Vec<&[u8]>)> = Vec::new();
+ for key in sorted_keys {
+ let (_, mut index_iter) = index_block.seek_and_iter(key, cmp);
+ let Some((_last_key, handle_bytes)) = index_iter.next() else {
+ break;
+ };
+ let handle = BlockHandle::decode(handle_bytes)?;
- // First block: seek within
- let first_block = match index_iter.next() {
- Some((_key, handle_bytes)) => {
- let handle = BlockHandle::decode(handle_bytes)?;
- self.read_data_block(&handle).await?
+ match target_blocks.last_mut() {
+ Some((current, block_keys))
+ if current.offset == handle.offset && current.size ==
handle.size =>
+ {
+ block_keys.push(key);
+ }
+ _ => target_blocks.push((handle, vec![key])),
}
- None => return Ok(result),
- };
-
- let (_, seeked) = first_block.seek_and_iter(sorted_keys[0], cmp);
- let mut offset = seeked.offset;
-
- if self.scan_block_in(
- &first_block,
- &mut offset,
- &sorted_keys,
- &mut key_idx,
- &mut result,
- )? {
- return Ok(result);
}
- while let Some((_key, handle_bytes)) = index_iter.next() {
- let handle = BlockHandle::decode(handle_bytes)?;
+ let mut result = RoaringTreemap::new();
+ for (handle, block_keys) in target_blocks {
let block = self.read_data_block(&handle).await?;
- let mut block_offset = 0;
- if self.scan_block_in(
- &block,
- &mut block_offset,
- &sorted_keys,
- &mut key_idx,
- &mut result,
- )? {
- return Ok(result);
+ for key in block_keys {
+ let (found, mut entry_iter) = block.seek_and_iter(key, cmp);
+ if let (true, Some((_entry_key, value))) = (found,
entry_iter.next()) {
+ insert_row_ids_into(value, &mut result)?;
+ }
}
}
Ok(result)
}
- /// Scan a block for IN query. Returns true when all keys are consumed.
- fn scan_block_in(
- &self,
- block: &BlockReader,
- offset: &mut usize,
- sorted_keys: &[&[u8]],
- key_idx: &mut usize,
- result: &mut RoaringTreemap,
- ) -> io::Result<bool> {
- let cmp = &self.key_comparator;
- while *offset < block.data.len() {
- let (entry_key, value, next_offset) = block.read_entry_at(*offset);
- *offset = next_offset;
-
- // Advance key_idx past keys smaller than current entry
- while *key_idx < sorted_keys.len()
- && cmp(sorted_keys[*key_idx], entry_key) == Ordering::Less
- {
- *key_idx += 1;
- }
-
- if *key_idx >= sorted_keys.len() {
- return Ok(true);
- }
-
- // Past the last query key
- if cmp(entry_key, sorted_keys[sorted_keys.len() - 1]) ==
Ordering::Greater {
- return Ok(true);
- }
-
- if cmp(entry_key, sorted_keys[*key_idx]) == Ordering::Equal {
- insert_row_ids_into(value, result)?;
- }
- }
- Ok(false)
- }
-
/// Not equal query.
pub async fn query_not_equal(&self, key: &[u8]) ->
io::Result<RoaringTreemap> {
let mut result = self.all_non_null_rows().await?;
diff --git a/crates/paimon/src/btree/tests.rs b/crates/paimon/src/btree/tests.rs
index 873c7124..252ff883 100644
--- a/crates/paimon/src/btree/tests.rs
+++ b/crates/paimon/src/btree/tests.rs
@@ -21,8 +21,26 @@ use crate::btree::query::IndexQuery;
use crate::btree::reader::BTreeIndexReader;
use crate::btree::test_util::{BytesFileRead, VecFileWrite};
use crate::btree::writer::BTreeIndexWriter;
+use crate::io::FileRead;
use crate::spec::{DataType, Datum, PredicateOperator, VarCharType};
use bytes::Bytes;
+use std::ops::Range;
+use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone)]
+struct RecordingFileRead {
+ data: Bytes,
+ ranges: Arc<Mutex<Vec<Range<u64>>>>,
+}
+
+#[async_trait::async_trait]
+impl FileRead for RecordingFileRead {
+ async fn read(&self, range: Range<u64>) -> crate::Result<Bytes> {
+ self.ranges.lock().unwrap().push(range.clone());
+ Ok(self.data.slice(range.start as usize..range.end as usize))
+ }
+}
fn int_key(v: i32) -> Vec<u8> {
v.to_be_bytes().to_vec()
@@ -310,6 +328,88 @@ async fn test_in_query() {
assert!(bm.contains(8));
}
+#[tokio::test]
+async fn test_sparse_in_query_reads_only_target_blocks_once() {
+ let buf = VecFileWrite::new();
+ let mut writer = BTreeIndexWriter::new(Box::new(buf.clone()), 64,
BlockCompressionType::None);
+
+ for i in 0..1000 {
+ writer.write(Some(&int_key(i * 2)), i as i64).await.unwrap();
+ }
+
+ let write_result = writer.finish().await.unwrap();
+ let data = Bytes::from(buf.to_vec());
+ let file_size = data.len() as u64;
+ let ranges = Arc::new(Mutex::new(Vec::new()));
+ let reader = BTreeIndexReader::open(
+ Box::new(RecordingFileRead {
+ data,
+ ranges: ranges.clone(),
+ }),
+ file_size,
+ &write_result.meta,
+ int_cmp,
+ )
+ .await
+ .unwrap();
+ ranges.lock().unwrap().clear();
+
+ // Unsorted, duplicate, missing, and out-of-bounds keys exercise filtering
as well as lookup.
+ // Each nearby group belongs to one target block, with large gaps between
the groups.
+ let keys = [1998, -2, 2, 1000, 3, 1001, 1002, 4, 1996, 2001, 1000]
+ .into_iter()
+ .map(int_key)
+ .collect::<Vec<_>>();
+ let key_refs = keys.iter().map(Vec::as_slice).collect::<Vec<_>>();
+ let rows = reader.query_in(&key_refs).await.unwrap();
+
+ assert_eq!(
+ rows.iter().collect::<Vec<_>>(),
+ vec![1, 2, 500, 501, 998, 999]
+ );
+ assert_eq!(ranges.lock().unwrap().len(), 3);
+}
+
+#[tokio::test]
+async fn test_sparse_in_query_seeks_each_key_within_target_block() {
+ let buf = VecFileWrite::new();
+ let mut writer =
+ BTreeIndexWriter::new(Box::new(buf.clone()), 1_000_000,
BlockCompressionType::None);
+
+ for i in 0..4096 {
+ writer.write(Some(&int_key(i)), i as i64).await.unwrap();
+ }
+
+ let write_result = writer.finish().await.unwrap();
+ let data = Bytes::from(buf.to_vec());
+ let file_size = data.len() as u64;
+ let comparisons = Arc::new(AtomicUsize::new(0));
+ let query_comparisons = comparisons.clone();
+ let reader = BTreeIndexReader::open(
+ Box::new(BytesFileRead(data)),
+ file_size,
+ &write_result.meta,
+ move |a, b| {
+ query_comparisons.fetch_add(1, AtomicOrdering::Relaxed);
+ int_cmp(a, b)
+ },
+ )
+ .await
+ .unwrap();
+ comparisons.store(0, AtomicOrdering::Relaxed);
+
+ let keys = [int_key(1), int_key(4094)];
+ let key_refs = keys.iter().map(Vec::as_slice).collect::<Vec<_>>();
+ let rows = reader.query_in(&key_refs).await.unwrap();
+
+ assert_eq!(rows.iter().collect::<Vec<_>>(), vec![1, 4094]);
+ let comparison_count = comparisons.load(AtomicOrdering::Relaxed);
+ assert!(
+ comparison_count < 100,
+ "sparse IN should seek target keys instead of scanning the block;
comparisons: {comparison_count}"
+ );
+}
+
#[tokio::test]
async fn test_string_keys() {
let buf = VecFileWrite::new();
