[rust] Reuse BTree data blocks across index lookups

crates/paimon/src/btree/block.rs

@@ -588,6 +588,16 @@ impl BlockReader {
         }
     }
 
+    pub(crate) fn retained_bytes(&self) -> usize {
+        let seek_bytes = match &self.seek_info {
+            SeekInfo::Aligned { .. } => 0,
+            SeekInfo::Unaligned { offsets } => offsets
+                .capacity()
+                .saturating_mul(std::mem::size_of::<i32>()),
+        };
+        self.data.capacity().saturating_add(seek_bytes)
+    }
+
     #[allow(dead_code)]
     pub fn record_count(&self) -> usize {
         self.record_count

crates/paimon/src/btree/data_block_cache.rs

@@ -0,0 +1,104 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::btree::block::BlockReader;
+use lru::LruCache;
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub(super) struct DataBlockCacheKey {
+    pub(super) file: Arc<str>,
+    pub(super) offset: u64,
+    pub(super) size: u32,
+}
+
+struct Entry {
+    block: Arc<BlockReader>,
+    retained_bytes: usize,
+}
+
+struct State {
+    entries: LruCache<DataBlockCacheKey, Entry>,
+    retained_bytes: usize,
+}
+
+pub(crate) struct BTreeDataBlockCache {
+    max_bytes: usize,
+    state: Mutex<State>,
+}
+
+impl BTreeDataBlockCache {
+    pub(crate) fn new(max_bytes: usize) -> Self {
+        Self {
+            max_bytes,
+            state: Mutex::new(State {
+                entries: LruCache::unbounded(),
+                retained_bytes: 0,
+            }),
+        }
+    }
+
+    pub(super) fn get(&self, key: &DataBlockCacheKey) -> Option<Arc<BlockReader>> {
+        if self.max_bytes == 0 {
+            return None;
+        }
+        self.state
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .entries
+            .get(key)
+            .map(|entry| Arc::clone(&entry.block))
+    }
+
+    pub(super) fn put(&self, key: DataBlockCacheKey, block: Arc<BlockReader>) -> Arc<BlockReader> {
+        if self.max_bytes == 0 {
+            return block;
+        }
+        let retained_bytes = block.retained_bytes();
+        if retained_bytes > self.max_bytes {
+            return block;
+        }
+
+        let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner());
+        if let Some(entry) = state.entries.get(&key) {
+            return Arc::clone(&entry.block);
+        }
+        state.retained_bytes = state.retained_bytes.saturating_add(retained_bytes);
+        state.entries.put(
+            key,
+            Entry {
+                block: Arc::clone(&block),
+                retained_bytes,
+            },
+        );
+        while state.retained_bytes > self.max_bytes {
+            let Some((_, entry)) = state.entries.pop_lru() else {
+                break;
+            };
+            state.retained_bytes = state.retained_bytes.saturating_sub(entry.retained_bytes);
+        }
+        block
+    }
+
+    #[cfg(test)]
+    pub(crate) fn retained_bytes(&self) -> usize {
+        self.state
+            .lock()
+            .unwrap_or_else(|error| error.into_inner())
+            .retained_bytes
+    }
+}

crates/paimon/src/btree/mod.rs

@@ -38,6 +38,7 @@
 
 mod block;
 mod bloom_filter;
+mod data_block_cache;
 mod footer;
 pub(crate) mod key_serde;
 mod meta;
@@ -52,6 +53,7 @@ pub use block::BlockCompressionType;
 pub(crate) use block::{
     compress_block, compress_codec_block, compute_crc32, decompress_block, decompress_codec_block,
 };
+pub(crate) use data_block_cache::BTreeDataBlockCache;
 pub use footer::BTreeFileFooter;
 pub use key_serde::{make_key_comparator, serialize_datum};
 pub use meta::BTreeIndexMeta;

crates/paimon/src/btree/reader.rs

@@ -25,6 +25,7 @@
 
 use crate::btree::block::{BlockHandle, BlockReader};
 use crate::btree::bloom_filter::BloomFilter;
+use crate::btree::data_block_cache::{BTreeDataBlockCache, DataBlockCacheKey};
 use crate::btree::footer::{BTreeFileFooter, BloomFilterHandle, BTREE_FOOTER_ENCODED_LENGTH};
 use crate::btree::key_serde::key_comparison_io_error;
 use crate::btree::meta::BTreeIndexMeta;
@@ -35,6 +36,7 @@ use crate::spec::murmur_hash::hash_bytes;
 use roaring::RoaringTreemap;
 use std::cmp::Ordering;
 use std::io;
+use std::sync::Arc;
 use tokio::sync::OnceCell;
 
 struct LazyBloomFilter {
@@ -72,6 +74,8 @@ pub struct BTreeIndexReader<F: Fn(&[u8], &[u8]) -> crate::Result<Ordering>> {
     key_comparator: F,
     bloom_filter: Option<LazyBloomFilter>,
     file_version: u32,
+    data_block_cache: Arc<BTreeDataBlockCache>,
+    data_block_cache_file: Arc<str>,
 }
 
 impl<F: Fn(&[u8], &[u8]) -> crate::Result<Ordering>> BTreeIndexReader<F> {
@@ -83,6 +87,25 @@ impl<F: Fn(&[u8], &[u8]) -> crate::Result<Ordering>> BTreeIndexReader<F> {
         file_size: u64,
         meta: &BTreeIndexMeta,
         key_comparator: F,
+    ) -> io::Result<Self> {
+        Self::open_with_data_block_cache(
+            reader,
+            file_size,
+            meta,
+            key_comparator,
+            Arc::new(BTreeDataBlockCache::new(0)),
+            Arc::from(""),
+        )
+        .await
+    }
+
+    pub(crate) async fn open_with_data_block_cache(
+        reader: Box<dyn FileRead>,
+        file_size: u64,
+        meta: &BTreeIndexMeta,
+        key_comparator: F,
+        data_block_cache: Arc<BTreeDataBlockCache>,
+        data_block_cache_file: Arc<str>,
     ) -> io::Result<Self> {
         if file_size < BTREE_FOOTER_ENCODED_LENGTH as u64 {
             return Err(io::Error::new(
@@ -128,6 +151,8 @@ impl<F: Fn(&[u8], &[u8]) -> crate::Result<Ordering>> BTreeIndexReader<F> {
             key_comparator,
             bloom_filter,
             file_version: footer.version,
+            data_block_cache,
+            data_block_cache_file,
         })
     }
 
@@ -287,14 +312,23 @@ impl<F: Fn(&[u8], &[u8]) -> crate::Result<Ordering>> BTreeIndexReader<F> {
     }
 
     /// Read a data block from the file on demand.
-    async fn read_data_block(&self, handle: &BlockHandle) -> io::Result<BlockReader> {
+    async fn read_data_block(&self, handle: &BlockHandle) -> io::Result<Arc<BlockReader>> {
+        let key = DataBlockCacheKey {
+            file: Arc::clone(&self.data_block_cache_file),
+            offset: handle.offset,
+            size: handle.size,
+        };
+        if let Some(block) = self.data_block_cache.get(&key) {
+            return Ok(block);
+        }
         let end = handle.offset + handle.full_block_size() as u64;
         let bytes = self
             .reader
             .read(handle.offset..end)
             .await
             .map_err(|e| io::Error::other(e.to_string()))?;
-        read_block_from_bytes(&bytes, handle.size)
+        let block = Arc::new(read_block_from_bytes(&bytes, handle.size)?);
+        Ok(self.data_block_cache.put(key, block))
     }
 
     async fn bloom_might_contain(&self, key: &[u8]) -> io::Result<bool> {