JingsongLi commented on code in PR #229: URL: https://github.com/apache/paimon-rust/pull/229#discussion_r3055629402
########## crates/paimon/src/btree/reader.rs: ########## @@ -0,0 +1,560 @@ +// 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. + +//! BTree index reader compatible with Java Paimon's BTreeIndexReader. +//! +//! Supports: +//! - Point lookup (equal) +//! - Range queries (less than, greater than, between, etc.) +//! - Null bitmap reading +//! - Sequential iteration over all entries + +use crate::btree::block::{BlockHandle, BlockReader}; +use crate::btree::footer::{BTreeFileFooter, BTREE_FOOTER_ENCODED_LENGTH}; +use crate::btree::meta::BTreeIndexMeta; +use crate::btree::sst_file::{read_block_from_bytes, SstFileReader}; +use crate::btree::var_len::{decode_var_int, decode_var_long}; +use crate::io::FileRead; +use bytes::Bytes; +use roaring::RoaringTreemap; +use std::cmp::Ordering; +use std::io::{self, Cursor}; + +/// BTree index reader with on-demand async data block loading. +pub struct BTreeIndexReader<F: Fn(&[u8], &[u8]) -> Ordering> { + reader: Box<dyn FileRead>, + sst_reader: SstFileReader, + null_bitmap: RoaringTreemap, + min_key: Option<Vec<u8>>, + max_key: Option<Vec<u8>>, + key_comparator: F, +} + +/// A key and its associated row ids. +#[derive(Debug, Clone)] +pub struct KeyRowIds { + pub key: Vec<u8>, + pub row_ids: Vec<i64>, +} + +impl<F: Fn(&[u8], &[u8]) -> Ordering> BTreeIndexReader<F> { + /// Open a BTree index reader from a FileRead and file metadata. + /// Only reads footer, index block, and null bitmap on open. + /// Data blocks are read on demand during queries. + pub async fn open( + reader: Box<dyn FileRead>, + file_size: u64, + meta: &BTreeIndexMeta, + key_comparator: F, + ) -> io::Result<Self> { + if file_size < BTREE_FOOTER_ENCODED_LENGTH as u64 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "File too small for BTree footer", + )); + } + + // 1. Read footer (last 52 bytes) + let footer_start = file_size - BTREE_FOOTER_ENCODED_LENGTH as u64; + let footer_bytes = reader + .read(footer_start..file_size) + .await + .map_err(|e| io::Error::other(e.to_string()))?; + let footer = BTreeFileFooter::read_footer(&footer_bytes)?; + + // 2. Read index block + let idx = &footer.index_block_handle; + let idx_end = idx.offset + idx.full_block_size() as u64; + let index_bytes = reader + .read(idx.offset..idx_end) + .await + .map_err(|e| io::Error::other(e.to_string()))?; + let index_block = read_block_from_bytes(&index_bytes, idx.size)?; + let sst_reader = SstFileReader::from_index_block(index_block); + + // 3. Read null bitmap + let null_bitmap = match &footer.null_bitmap_handle { + Some(h) => read_null_bitmap(reader.as_ref(), h).await?, + None => RoaringTreemap::new(), + }; + + Ok(Self { + reader, + sst_reader, + null_bitmap, + min_key: meta.first_key.clone(), + max_key: meta.last_key.clone(), + key_comparator, + }) + } + + /// Open from in-memory file data (for tests and backward compatibility). + /// Supports `entry_iterator()` since full file data is available. + pub fn new(file_data: Vec<u8>, meta: &BTreeIndexMeta, key_comparator: F) -> io::Result<Self> { + let file_size = file_data.len(); + + if file_size < BTREE_FOOTER_ENCODED_LENGTH { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "File too small for BTree footer", + )); + } + let footer_start = file_size - BTREE_FOOTER_ENCODED_LENGTH; + let footer = BTreeFileFooter::read_footer(&file_data[footer_start..])?; + + let null_bitmap = match footer.null_bitmap_handle { + Some(handle) => read_null_bitmap_from_bytes(&file_data, &handle)?, + None => RoaringTreemap::new(), + }; + + // Use SstFileReader::new to keep full data for entry_iterator and sync queries + let sst_reader = SstFileReader::new(file_data.clone(), footer.index_block_handle)?; Review Comment: I have modified all the writers to use FileIO for writing, and remove byte[] input for SstFileReader. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
