diff --git a/src/common/base/Cargo.toml b/src/common/base/Cargo.toml index 2d35ad5d31a..44eb5535ea0 100644 --- a/src/common/base/Cargo.toml +++ b/src/common/base/Cargo.toml @@ -4,6 +4,9 @@ version.workspace = true edition.workspace = true license.workspace = true +[features] +testing = [] + [lints] workspace = true diff --git a/src/common/base/src/range_read.rs b/src/common/base/src/range_read.rs index 53c26eeebde..fb0fc61fb03 100644 --- a/src/common/base/src/range_read.rs +++ b/src/common/base/src/range_read.rs @@ -17,6 +17,7 @@ use std::io; use std::ops::Range; use std::path::Path; use std::pin::Pin; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::task::{Context, Poll}; @@ -33,19 +34,22 @@ pub struct Metadata { pub content_length: u64, } -/// `RangeReader` reads a range of bytes from a source. -#[async_trait] -pub trait RangeReader: Send + Unpin { +/// `SizeAwareRangeReader` is a `RangeReader` that supports setting a file size hint. +pub trait SizeAwareRangeReader: RangeReader { /// Sets the file size hint for the reader. /// /// It's used to optimize the reading process by reducing the number of remote requests. fn with_file_size_hint(&mut self, file_size_hint: u64); +} +/// `RangeReader` reads a range of bytes from a source. +#[async_trait] +pub trait RangeReader: Sync + Send + Unpin { /// Returns the metadata of the source. - async fn metadata(&mut self) -> io::Result; + async fn metadata(&self) -> io::Result; /// Reads the bytes in the given range. - async fn read(&mut self, range: Range) -> io::Result; + async fn read(&self, range: Range) -> io::Result; /// Reads the bytes in the given range into the buffer. /// @@ -53,18 +57,14 @@ pub trait RangeReader: Send + Unpin { /// - If the buffer is insufficient to hold the bytes, it will either: /// - Allocate additional space (e.g., for `Vec`) /// - Panic (e.g., for `&mut [u8]`) - async fn read_into( - &mut self, - range: Range, - buf: &mut (impl BufMut + Send), - ) -> io::Result<()> { + async fn read_into(&self, range: Range, buf: &mut (impl BufMut + Send)) -> io::Result<()> { let bytes = self.read(range).await?; buf.put_slice(&bytes); Ok(()) } /// Reads the bytes in the given ranges. - async fn read_vec(&mut self, ranges: &[Range]) -> io::Result> { + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { let mut result = Vec::with_capacity(ranges.len()); for range in ranges { result.push(self.read(range.clone()).await?); @@ -74,25 +74,20 @@ pub trait RangeReader: Send + Unpin { } #[async_trait] -impl RangeReader for &mut R { - fn with_file_size_hint(&mut self, file_size_hint: u64) { - (*self).with_file_size_hint(file_size_hint) - } - - async fn metadata(&mut self) -> io::Result { +impl RangeReader for &R { + async fn metadata(&self) -> io::Result { (*self).metadata().await } - async fn read(&mut self, range: Range) -> io::Result { + + async fn read(&self, range: Range) -> io::Result { (*self).read(range).await } - async fn read_into( - &mut self, - range: Range, - buf: &mut (impl BufMut + Send), - ) -> io::Result<()> { + + async fn read_into(&self, range: Range, buf: &mut (impl BufMut + Send)) -> io::Result<()> { (*self).read_into(range, buf).await } - async fn read_vec(&mut self, ranges: &[Range]) -> io::Result> { + + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { (*self).read_vec(ranges).await } } @@ -120,7 +115,7 @@ pub struct AsyncReadAdapter { impl AsyncReadAdapter { pub async fn new(inner: R) -> io::Result { - let mut inner = inner; + let inner = inner; let metadata = inner.metadata().await?; Ok(AsyncReadAdapter { inner: Arc::new(Mutex::new(inner)), @@ -160,7 +155,7 @@ impl AsyncRead for AsyncReadAdapter { let range = *this.position..(*this.position + size); let inner = this.inner.clone(); let fut = async move { - let mut inner = inner.lock().await; + let inner = inner.lock().await; inner.read(range).await }; @@ -195,27 +190,24 @@ impl AsyncRead for AsyncReadAdapter { #[async_trait] impl RangeReader for Vec { - fn with_file_size_hint(&mut self, _file_size_hint: u64) { - // do nothing - } - - async fn metadata(&mut self) -> io::Result { + async fn metadata(&self) -> io::Result { Ok(Metadata { content_length: self.len() as u64, }) } - async fn read(&mut self, range: Range) -> io::Result { + async fn read(&self, range: Range) -> io::Result { let bytes = Bytes::copy_from_slice(&self[range.start as usize..range.end as usize]); Ok(bytes) } } +// TODO(weny): considers replacing `tokio::fs::File` with opendal reader. /// `FileReader` is a `RangeReader` for reading a file. pub struct FileReader { content_length: u64, - position: u64, - file: tokio::fs::File, + position: AtomicU64, + file: Mutex, } impl FileReader { @@ -225,32 +217,37 @@ impl FileReader { let metadata = file.metadata().await?; Ok(FileReader { content_length: metadata.len(), - position: 0, - file, + position: AtomicU64::new(0), + file: Mutex::new(file), }) } } +#[cfg(any(test, feature = "testing"))] +impl SizeAwareRangeReader for FileReader { + fn with_file_size_hint(&mut self, _file_size_hint: u64) { + // do nothing + } +} + #[async_trait] impl RangeReader for FileReader { - fn with_file_size_hint(&mut self, _file_size_hint: u64) { - // do nothing - } - - async fn metadata(&mut self) -> io::Result { + async fn metadata(&self) -> io::Result { Ok(Metadata { content_length: self.content_length, }) } - async fn read(&mut self, mut range: Range) -> io::Result { - if range.start != self.position { - self.file.seek(io::SeekFrom::Start(range.start)).await?; - self.position = range.start; + async fn read(&self, mut range: Range) -> io::Result { + let mut file = self.file.lock().await; + + if range.start != self.position.load(Ordering::Relaxed) { + file.seek(io::SeekFrom::Start(range.start)).await?; + self.position.store(range.start, Ordering::Relaxed); } range.end = range.end.min(self.content_length); - if range.end <= self.position { + if range.end <= self.position.load(Ordering::Relaxed) { return Err(io::Error::new( io::ErrorKind::UnexpectedEof, "Start of range is out of bounds", @@ -259,8 +256,8 @@ impl RangeReader for FileReader { let mut buf = vec![0; (range.end - range.start) as usize]; - self.file.read_exact(&mut buf).await?; - self.position = range.end; + file.read_exact(&mut buf).await?; + self.position.store(range.end, Ordering::Relaxed); Ok(Bytes::from(buf)) } @@ -301,7 +298,7 @@ mod tests { let data = b"hello world"; tokio::fs::write(path, data).await.unwrap(); - let mut reader = FileReader::new(path).await.unwrap(); + let reader = FileReader::new(path).await.unwrap(); let metadata = reader.metadata().await.unwrap(); assert_eq!(metadata.content_length, data.len() as u64); diff --git a/src/index/src/bloom_filter/applier.rs b/src/index/src/bloom_filter/applier.rs index 2750cbb92b6..e9b11a8f1d6 100644 --- a/src/index/src/bloom_filter/applier.rs +++ b/src/index/src/bloom_filter/applier.rs @@ -42,7 +42,7 @@ pub struct BloomFilterApplier { } impl BloomFilterApplier { - pub async fn new(mut reader: Box) -> Result { + pub async fn new(reader: Box) -> Result { let meta = reader.metadata().await?; Ok(Self { reader, meta }) diff --git a/src/index/src/bloom_filter/reader.rs b/src/index/src/bloom_filter/reader.rs index 02085fa671f..e816ab01c60 100644 --- a/src/index/src/bloom_filter/reader.rs +++ b/src/index/src/bloom_filter/reader.rs @@ -29,16 +29,16 @@ use crate::bloom_filter::{BloomFilterMeta, BloomFilterSegmentLocation, SEED}; const BLOOM_META_LEN_SIZE: u64 = 4; /// Default prefetch size of bloom filter meta. -pub const DEFAULT_PREFETCH_SIZE: u64 = 1024; // 1KiB +pub const DEFAULT_PREFETCH_SIZE: u64 = 8192; // 8KiB /// `BloomFilterReader` reads the bloom filter from the file. #[async_trait] -pub trait BloomFilterReader { +pub trait BloomFilterReader: Sync { /// Reads range of bytes from the file. - async fn range_read(&mut self, offset: u64, size: u32) -> Result; + async fn range_read(&self, offset: u64, size: u32) -> Result; /// Reads bunch of ranges from the file. - async fn read_vec(&mut self, ranges: &[Range]) -> Result> { + async fn read_vec(&self, ranges: &[Range]) -> Result> { let mut results = Vec::with_capacity(ranges.len()); for range in ranges { let size = (range.end - range.start) as u32; @@ -49,10 +49,10 @@ pub trait BloomFilterReader { } /// Reads the meta information of the bloom filter. - async fn metadata(&mut self) -> Result; + async fn metadata(&self) -> Result; /// Reads a bloom filter with the given location. - async fn bloom_filter(&mut self, loc: &BloomFilterSegmentLocation) -> Result { + async fn bloom_filter(&self, loc: &BloomFilterSegmentLocation) -> Result { let bytes = self.range_read(loc.offset, loc.size as _).await?; let vec = bytes .chunks_exact(std::mem::size_of::()) @@ -80,23 +80,23 @@ impl BloomFilterReaderImpl { #[async_trait] impl BloomFilterReader for BloomFilterReaderImpl { - async fn range_read(&mut self, offset: u64, size: u32) -> Result { + async fn range_read(&self, offset: u64, size: u32) -> Result { self.reader .read(offset..offset + size as u64) .await .context(IoSnafu) } - async fn read_vec(&mut self, ranges: &[Range]) -> Result> { + async fn read_vec(&self, ranges: &[Range]) -> Result> { self.reader.read_vec(ranges).await.context(IoSnafu) } - async fn metadata(&mut self) -> Result { + async fn metadata(&self) -> Result { let metadata = self.reader.metadata().await.context(IoSnafu)?; let file_size = metadata.content_length; let mut meta_reader = - BloomFilterMetaReader::new(&mut self.reader, file_size, Some(DEFAULT_PREFETCH_SIZE)); + BloomFilterMetaReader::new(&self.reader, file_size, Some(DEFAULT_PREFETCH_SIZE)); meta_reader.metadata().await } } @@ -250,7 +250,7 @@ mod tests { async fn test_bloom_filter_reader() { let bytes = mock_bloom_filter_bytes().await; - let mut reader = BloomFilterReaderImpl::new(bytes); + let reader = BloomFilterReaderImpl::new(bytes); let meta = reader.metadata().await.unwrap(); assert_eq!(meta.bloom_filter_segments.len(), 2); diff --git a/src/index/src/inverted_index/format/reader.rs b/src/index/src/inverted_index/format/reader.rs index 24f21504d00..268b169979e 100644 --- a/src/index/src/inverted_index/format/reader.rs +++ b/src/index/src/inverted_index/format/reader.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::VecDeque; use std::ops::Range; use std::sync::Arc; @@ -33,10 +34,10 @@ mod footer; #[async_trait] pub trait InvertedIndexReader: Send + Sync { /// Seeks to given offset and reads data with exact size as provided. - async fn range_read(&mut self, offset: u64, size: u32) -> Result>; + async fn range_read(&self, offset: u64, size: u32) -> Result>; /// Reads the bytes in the given ranges. - async fn read_vec(&mut self, ranges: &[Range]) -> Result> { + async fn read_vec(&self, ranges: &[Range]) -> Result> { let mut result = Vec::with_capacity(ranges.len()); for range in ranges { let data = self @@ -48,16 +49,35 @@ pub trait InvertedIndexReader: Send + Sync { } /// Retrieves metadata of all inverted indices stored within the blob. - async fn metadata(&mut self) -> Result>; + async fn metadata(&self) -> Result>; /// Retrieves the finite state transducer (FST) map from the given offset and size. - async fn fst(&mut self, offset: u64, size: u32) -> Result { + async fn fst(&self, offset: u64, size: u32) -> Result { let fst_data = self.range_read(offset, size).await?; FstMap::new(fst_data).context(DecodeFstSnafu) } + /// Retrieves the multiple finite state transducer (FST) maps from the given ranges. + async fn fst_vec(&mut self, ranges: &[Range]) -> Result> { + self.read_vec(ranges) + .await? + .into_iter() + .map(|bytes| FstMap::new(bytes.to_vec()).context(DecodeFstSnafu)) + .collect::>>() + } + /// Retrieves the bitmap from the given offset and size. - async fn bitmap(&mut self, offset: u64, size: u32) -> Result { + async fn bitmap(&self, offset: u64, size: u32) -> Result { self.range_read(offset, size).await.map(BitVec::from_vec) } + + /// Retrieves the multiple bitmaps from the given ranges. + async fn bitmap_deque(&mut self, ranges: &[Range]) -> Result> { + Ok(self + .read_vec(ranges) + .await? + .into_iter() + .map(|bytes| BitVec::from_slice(bytes.as_ref())) + .collect::>()) + } } diff --git a/src/index/src/inverted_index/format/reader/blob.rs b/src/index/src/inverted_index/format/reader/blob.rs index 73d98835794..b48a3224844 100644 --- a/src/index/src/inverted_index/format/reader/blob.rs +++ b/src/index/src/inverted_index/format/reader/blob.rs @@ -52,7 +52,7 @@ impl InvertedIndexBlobReader { #[async_trait] impl InvertedIndexReader for InvertedIndexBlobReader { - async fn range_read(&mut self, offset: u64, size: u32) -> Result> { + async fn range_read(&self, offset: u64, size: u32) -> Result> { let buf = self .source .read(offset..offset + size as u64) @@ -61,16 +61,16 @@ impl InvertedIndexReader for InvertedIndexBlobReader { Ok(buf.into()) } - async fn read_vec(&mut self, ranges: &[Range]) -> Result> { + async fn read_vec(&self, ranges: &[Range]) -> Result> { self.source.read_vec(ranges).await.context(CommonIoSnafu) } - async fn metadata(&mut self) -> Result> { + async fn metadata(&self) -> Result> { let metadata = self.source.metadata().await.context(CommonIoSnafu)?; let blob_size = metadata.content_length; Self::validate_blob_size(blob_size)?; - let mut footer_reader = InvertedIndexFooterReader::new(&mut self.source, blob_size) + let mut footer_reader = InvertedIndexFooterReader::new(&self.source, blob_size) .with_prefetch_size(DEFAULT_PREFETCH_SIZE); footer_reader.metadata().await.map(Arc::new) } @@ -160,7 +160,7 @@ mod tests { #[tokio::test] async fn test_inverted_index_blob_reader_metadata() { let blob = create_inverted_index_blob(); - let mut blob_reader = InvertedIndexBlobReader::new(blob); + let blob_reader = InvertedIndexBlobReader::new(blob); let metas = blob_reader.metadata().await.unwrap(); assert_eq!(metas.metas.len(), 2); @@ -187,7 +187,7 @@ mod tests { #[tokio::test] async fn test_inverted_index_blob_reader_fst() { let blob = create_inverted_index_blob(); - let mut blob_reader = InvertedIndexBlobReader::new(blob); + let blob_reader = InvertedIndexBlobReader::new(blob); let metas = blob_reader.metadata().await.unwrap(); let meta = metas.metas.get("tag0").unwrap(); @@ -219,7 +219,7 @@ mod tests { #[tokio::test] async fn test_inverted_index_blob_reader_bitmap() { let blob = create_inverted_index_blob(); - let mut blob_reader = InvertedIndexBlobReader::new(blob); + let blob_reader = InvertedIndexBlobReader::new(blob); let metas = blob_reader.metadata().await.unwrap(); let meta = metas.metas.get("tag0").unwrap(); diff --git a/src/index/src/inverted_index/format/reader/footer.rs b/src/index/src/inverted_index/format/reader/footer.rs index c025ecf52ec..b5a855a4559 100644 --- a/src/index/src/inverted_index/format/reader/footer.rs +++ b/src/index/src/inverted_index/format/reader/footer.rs @@ -24,7 +24,7 @@ use crate::inverted_index::error::{ }; use crate::inverted_index::format::FOOTER_PAYLOAD_SIZE_SIZE; -pub const DEFAULT_PREFETCH_SIZE: u64 = 1024; // 1KiB +pub const DEFAULT_PREFETCH_SIZE: u64 = 8192; // 8KiB /// InvertedIndexFooterReader is for reading the footer section of the blob. pub struct InvertedIndexFooterReader { @@ -177,11 +177,11 @@ mod tests { ..Default::default() }; - let mut payload_buf = create_test_payload(meta); + let payload_buf = create_test_payload(meta); let blob_size = payload_buf.len() as u64; for prefetch in [0, blob_size / 2, blob_size, blob_size + 10] { - let mut reader = InvertedIndexFooterReader::new(&mut payload_buf, blob_size); + let mut reader = InvertedIndexFooterReader::new(&payload_buf, blob_size); if prefetch > 0 { reader = reader.with_prefetch_size(prefetch); } @@ -205,7 +205,7 @@ mod tests { for prefetch in [0, blob_size / 2, blob_size, blob_size + 10] { let blob_size = payload_buf.len() as u64; - let mut reader = InvertedIndexFooterReader::new(&mut payload_buf, blob_size); + let mut reader = InvertedIndexFooterReader::new(&payload_buf, blob_size); if prefetch > 0 { reader = reader.with_prefetch_size(prefetch); } @@ -224,11 +224,11 @@ mod tests { ..Default::default() }; - let mut payload_buf = create_test_payload(meta); + let payload_buf = create_test_payload(meta); let blob_size = payload_buf.len() as u64; for prefetch in [0, blob_size / 2, blob_size, blob_size + 10] { - let mut reader = InvertedIndexFooterReader::new(&mut payload_buf, blob_size); + let mut reader = InvertedIndexFooterReader::new(&payload_buf, blob_size); if prefetch > 0 { reader = reader.with_prefetch_size(prefetch); } diff --git a/src/index/src/inverted_index/format/writer/blob.rs b/src/index/src/inverted_index/format/writer/blob.rs index 7b95b8deafa..d53dfee855a 100644 --- a/src/index/src/inverted_index/format/writer/blob.rs +++ b/src/index/src/inverted_index/format/writer/blob.rs @@ -118,7 +118,7 @@ mod tests { .await .unwrap(); - let mut reader = InvertedIndexBlobReader::new(blob); + let reader = InvertedIndexBlobReader::new(blob); let metadata = reader.metadata().await.unwrap(); assert_eq!(metadata.total_row_count, 8); assert_eq!(metadata.segment_row_count, 1); @@ -158,7 +158,7 @@ mod tests { .await .unwrap(); - let mut reader = InvertedIndexBlobReader::new(blob); + let reader = InvertedIndexBlobReader::new(blob); let metadata = reader.metadata().await.unwrap(); assert_eq!(metadata.total_row_count, 8); assert_eq!(metadata.segment_row_count, 1); diff --git a/src/index/src/inverted_index/search/fst_values_mapper.rs b/src/index/src/inverted_index/search/fst_values_mapper.rs index 71665315212..54a842de020 100644 --- a/src/index/src/inverted_index/search/fst_values_mapper.rs +++ b/src/index/src/inverted_index/search/fst_values_mapper.rs @@ -18,55 +18,75 @@ use greptime_proto::v1::index::InvertedIndexMeta; use crate::inverted_index::error::Result; use crate::inverted_index::format::reader::InvertedIndexReader; -/// `FstValuesMapper` maps FST-encoded u64 values to their corresponding bitmaps -/// within an inverted index. +/// `ParallelFstValuesMapper` enables parallel mapping of multiple FST value groups to their +/// corresponding bitmaps within an inverted index. /// -/// The higher 32 bits of each u64 value represent the -/// bitmap offset and the lower 32 bits represent its size. This mapper uses these -/// combined offset-size pairs to fetch and union multiple bitmaps into a single `BitVec`. -pub struct FstValuesMapper<'a> { - /// `reader` retrieves bitmap data using offsets and sizes from FST values. +/// This mapper processes multiple groups of FST values in parallel, where each group is associated +/// with its own metadata. It optimizes bitmap retrieval by batching requests across all groups +/// before combining them into separate result bitmaps. +pub struct ParallelFstValuesMapper<'a> { reader: &'a mut dyn InvertedIndexReader, - - /// `metadata` provides context for interpreting the index structures. - metadata: &'a InvertedIndexMeta, } -impl<'a> FstValuesMapper<'a> { - pub fn new( - reader: &'a mut dyn InvertedIndexReader, - metadata: &'a InvertedIndexMeta, - ) -> FstValuesMapper<'a> { - FstValuesMapper { reader, metadata } +impl<'a> ParallelFstValuesMapper<'a> { + pub fn new(reader: &'a mut dyn InvertedIndexReader) -> Self { + Self { reader } } - /// Maps an array of FST values to a `BitVec` by retrieving and combining bitmaps. - pub async fn map_values(&mut self, values: &[u64]) -> Result { - let mut bitmap = BitVec::new(); + pub async fn map_values_vec( + &mut self, + value_and_meta_vec: &[(Vec, &'a InvertedIndexMeta)], + ) -> Result> { + let groups = value_and_meta_vec + .iter() + .map(|(values, _)| values.len()) + .collect::>(); + let len = groups.iter().sum::(); + let mut fetch_ranges = Vec::with_capacity(len); - for value in values { - // relative_offset (higher 32 bits), size (lower 32 bits) - let [relative_offset, size] = bytemuck::cast::(*value); - - let bm = self - .reader - .bitmap(self.metadata.base_offset + relative_offset as u64, size) - .await?; - - // Ensure the longest BitVec is the left operand to prevent truncation during OR. - if bm.len() > bitmap.len() { - bitmap = bm | bitmap - } else { - bitmap |= bm + for (values, meta) in value_and_meta_vec { + for value in values { + // The higher 32 bits of each u64 value represent the + // bitmap offset and the lower 32 bits represent its size. This mapper uses these + // combined offset-size pairs to fetch and union multiple bitmaps into a single `BitVec`. + let [relative_offset, size] = bytemuck::cast::(*value); + fetch_ranges.push( + meta.base_offset + relative_offset as u64 + ..meta.base_offset + relative_offset as u64 + size as u64, + ); } } - Ok(bitmap) + if fetch_ranges.is_empty() { + return Ok(vec![BitVec::new()]); + } + + common_telemetry::debug!("fetch ranges: {:?}", fetch_ranges); + let mut bitmaps = self.reader.bitmap_deque(&fetch_ranges).await?; + let mut output = Vec::with_capacity(groups.len()); + + for counter in groups { + let mut bitmap = BitVec::new(); + for _ in 0..counter { + let bm = bitmaps.pop_front().unwrap(); + if bm.len() > bitmap.len() { + bitmap = bm | bitmap + } else { + bitmap |= bm + } + } + + output.push(bitmap); + } + + Ok(output) } } #[cfg(test)] mod tests { + use std::collections::VecDeque; + use common_base::bit_vec::prelude::*; use super::*; @@ -77,38 +97,70 @@ mod tests { } #[tokio::test] - async fn test_map_values() { + async fn test_map_values_vec() { let mut mock_reader = MockInvertedIndexReader::new(); - mock_reader - .expect_bitmap() - .returning(|offset, size| match (offset, size) { - (1, 1) => Ok(bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1]), - (2, 1) => Ok(bitvec![u8, Lsb0; 0, 1, 0, 1, 0, 1, 0, 1]), - _ => unreachable!(), - }); + mock_reader.expect_bitmap_deque().returning(|ranges| { + let mut output = VecDeque::new(); + for range in ranges { + let offset = range.start; + let size = range.end - range.start; + match (offset, size) { + (1, 1) => output.push_back(bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1]), + (2, 1) => output.push_back(bitvec![u8, Lsb0; 0, 1, 0, 1, 0, 1, 0, 1]), + _ => unreachable!(), + } + } + Ok(output) + }); let meta = InvertedIndexMeta::default(); - let mut values_mapper = FstValuesMapper::new(&mut mock_reader, &meta); - - let result = values_mapper.map_values(&[]).await.unwrap(); - assert_eq!(result.count_ones(), 0); - - let result = values_mapper.map_values(&[value(1, 1)]).await.unwrap(); - assert_eq!(result, bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1]); - - let result = values_mapper.map_values(&[value(2, 1)]).await.unwrap(); - assert_eq!(result, bitvec![u8, Lsb0; 0, 1, 0, 1, 0, 1, 0, 1]); + let mut values_mapper = ParallelFstValuesMapper::new(&mut mock_reader); let result = values_mapper - .map_values(&[value(1, 1), value(2, 1)]) + .map_values_vec(&[(vec![], &meta)]) .await .unwrap(); - assert_eq!(result, bitvec![u8, Lsb0; 1, 1, 1, 1, 1, 1, 1, 1]); + assert_eq!(result[0].count_ones(), 0); let result = values_mapper - .map_values(&[value(2, 1), value(1, 1)]) + .map_values_vec(&[(vec![value(1, 1)], &meta)]) .await .unwrap(); - assert_eq!(result, bitvec![u8, Lsb0; 1, 1, 1, 1, 1, 1, 1, 1]); + assert_eq!(result[0], bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1]); + + let result = values_mapper + .map_values_vec(&[(vec![value(2, 1)], &meta)]) + .await + .unwrap(); + assert_eq!(result[0], bitvec![u8, Lsb0; 0, 1, 0, 1, 0, 1, 0, 1]); + + let result = values_mapper + .map_values_vec(&[(vec![value(1, 1), value(2, 1)], &meta)]) + .await + .unwrap(); + assert_eq!(result[0], bitvec![u8, Lsb0; 1, 1, 1, 1, 1, 1, 1, 1]); + + let result = values_mapper + .map_values_vec(&[(vec![value(2, 1), value(1, 1)], &meta)]) + .await + .unwrap(); + assert_eq!(result[0], bitvec![u8, Lsb0; 1, 1, 1, 1, 1, 1, 1, 1]); + + let result = values_mapper + .map_values_vec(&[(vec![value(2, 1)], &meta), (vec![value(1, 1)], &meta)]) + .await + .unwrap(); + assert_eq!(result[0], bitvec![u8, Lsb0; 0, 1, 0, 1, 0, 1, 0, 1]); + assert_eq!(result[1], bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1]); + + let result = values_mapper + .map_values_vec(&[ + (vec![value(2, 1), value(1, 1)], &meta), + (vec![value(1, 1)], &meta), + ]) + .await + .unwrap(); + assert_eq!(result[0], bitvec![u8, Lsb0; 1, 1, 1, 1, 1, 1, 1, 1]); + assert_eq!(result[1], bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1]); } } diff --git a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs index bead0761d7b..cf8e30085f0 100644 --- a/src/index/src/inverted_index/search/index_apply/predicates_apply.rs +++ b/src/index/src/inverted_index/search/index_apply/predicates_apply.rs @@ -23,7 +23,7 @@ use crate::inverted_index::format::reader::InvertedIndexReader; use crate::inverted_index::search::fst_apply::{ FstApplier, IntersectionFstApplier, KeysFstApplier, }; -use crate::inverted_index::search::fst_values_mapper::FstValuesMapper; +use crate::inverted_index::search::fst_values_mapper::ParallelFstValuesMapper; use crate::inverted_index::search::index_apply::{ ApplyOutput, IndexApplier, IndexNotFoundStrategy, SearchContext, }; @@ -57,11 +57,10 @@ impl IndexApplier for PredicatesIndexApplier { let mut bitmap = Self::bitmap_full_range(&metadata); // TODO(zhongzc): optimize the order of applying to make it quicker to return empty. - for (name, fst_applier) in &self.fst_appliers { - if bitmap.count_ones() == 0 { - break; - } + let mut appliers = Vec::with_capacity(self.fst_appliers.len()); + let mut fst_ranges = Vec::with_capacity(self.fst_appliers.len()); + for (name, fst_applier) in &self.fst_appliers { let Some(meta) = metadata.metas.get(name) else { match context.index_not_found_strategy { IndexNotFoundStrategy::ReturnEmpty => { @@ -75,14 +74,31 @@ impl IndexApplier for PredicatesIndexApplier { } } }; - let fst_offset = meta.base_offset + meta.relative_fst_offset as u64; - let fst_size = meta.fst_size; - let fst = reader.fst(fst_offset, fst_size).await?; - let values = fst_applier.apply(&fst); + let fst_size = meta.fst_size as u64; + appliers.push((fst_applier, meta)); + fst_ranges.push(fst_offset..fst_offset + fst_size); + } - let mut mapper = FstValuesMapper::new(&mut *reader, meta); - let bm = mapper.map_values(&values).await?; + if fst_ranges.is_empty() { + output.matched_segment_ids = bitmap; + return Ok(output); + } + + let fsts = reader.fst_vec(&fst_ranges).await?; + let value_and_meta_vec = fsts + .into_iter() + .zip(appliers) + .map(|(fst, (fst_applier, meta))| (fst_applier.apply(&fst), meta)) + .collect::>(); + + let mut mapper = ParallelFstValuesMapper::new(reader); + let bm_vec = mapper.map_values_vec(&value_and_meta_vec).await?; + + for bm in bm_vec { + if bitmap.count_ones() == 0 { + break; + } bitmap &= bm; } @@ -148,6 +164,7 @@ impl TryFrom)>> for PredicatesIndexApplier { #[cfg(test)] mod tests { + use std::collections::VecDeque; use std::sync::Arc; use common_base::bit_vec::prelude::*; @@ -204,15 +221,19 @@ mod tests { mock_reader .expect_metadata() .returning(|| Ok(mock_metas([("tag-0", 0)]))); - mock_reader.expect_fst().returning(|_offset, _size| { - Ok(FstMap::from_iter([(b"tag-0_value-0", fst_value(2, 1))]).unwrap()) + mock_reader.expect_fst_vec().returning(|_ranges| { + Ok(vec![FstMap::from_iter([( + b"tag-0_value-0", + fst_value(2, 1), + )]) + .unwrap()]) + }); + + mock_reader.expect_bitmap_deque().returning(|range| { + assert_eq!(range.len(), 1); + assert_eq!(range[0], 2..3); + Ok(VecDeque::from([bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1, 0]])) }); - mock_reader - .expect_bitmap() - .returning(|offset, size| match (offset, size) { - (2, 1) => Ok(bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1, 0]), - _ => unreachable!(), - }); let output = applier .apply(SearchContext::default(), &mut mock_reader) .await @@ -227,8 +248,12 @@ mod tests { mock_reader .expect_metadata() .returning(|| Ok(mock_metas([("tag-0", 0)]))); - mock_reader.expect_fst().returning(|_offset, _size| { - Ok(FstMap::from_iter([(b"tag-0_value-1", fst_value(2, 1))]).unwrap()) + mock_reader.expect_fst_vec().returning(|_range| { + Ok(vec![FstMap::from_iter([( + b"tag-0_value-1", + fst_value(2, 1), + )]) + .unwrap()]) }); let output = applier .apply(SearchContext::default(), &mut mock_reader) @@ -252,20 +277,33 @@ mod tests { mock_reader .expect_metadata() .returning(|| Ok(mock_metas([("tag-0", 0), ("tag-1", 1)]))); - mock_reader - .expect_fst() - .returning(|offset, _size| match offset { - 0 => Ok(FstMap::from_iter([(b"tag-0_value-0", fst_value(1, 1))]).unwrap()), - 1 => Ok(FstMap::from_iter([(b"tag-1_value-a", fst_value(2, 1))]).unwrap()), - _ => unreachable!(), - }); - mock_reader - .expect_bitmap() - .returning(|offset, size| match (offset, size) { - (1, 1) => Ok(bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1, 0]), - (2, 1) => Ok(bitvec![u8, Lsb0; 1, 1, 0, 1, 1, 0, 1, 1]), - _ => unreachable!(), - }); + mock_reader.expect_fst_vec().returning(|ranges| { + let mut output = vec![]; + for range in ranges { + match range.start { + 0 => output + .push(FstMap::from_iter([(b"tag-0_value-0", fst_value(1, 1))]).unwrap()), + 1 => output + .push(FstMap::from_iter([(b"tag-1_value-a", fst_value(2, 1))]).unwrap()), + _ => unreachable!(), + } + } + Ok(output) + }); + mock_reader.expect_bitmap_deque().returning(|ranges| { + let mut output = VecDeque::new(); + for range in ranges { + let offset = range.start; + let size = range.end - range.start; + match (offset, size) { + (1, 1) => output.push_back(bitvec![u8, Lsb0; 1, 0, 1, 0, 1, 0, 1, 0]), + (2, 1) => output.push_back(bitvec![u8, Lsb0; 1, 1, 0, 1, 1, 0, 1, 1]), + _ => unreachable!(), + } + } + + Ok(output) + }); let output = applier .apply(SearchContext::default(), &mut mock_reader) diff --git a/src/mito2/src/cache/index.rs b/src/mito2/src/cache/index.rs index 137dc3d8745..c28ac063b32 100644 --- a/src/mito2/src/cache/index.rs +++ b/src/mito2/src/cache/index.rs @@ -157,7 +157,7 @@ where load: F, ) -> Result, E> where - F: FnOnce(Vec>) -> Fut, + F: Fn(Vec>) -> Fut, Fut: Future, E>>, E: std::error::Error, { diff --git a/src/mito2/src/cache/index/bloom_filter_index.rs b/src/mito2/src/cache/index/bloom_filter_index.rs index 08ac7e8273d..c5b95a5b843 100644 --- a/src/mito2/src/cache/index/bloom_filter_index.rs +++ b/src/mito2/src/cache/index/bloom_filter_index.rs @@ -87,8 +87,8 @@ impl CachedBloomFilterIndexBlobReader { #[async_trait] impl BloomFilterReader for CachedBloomFilterIndexBlobReader { - async fn range_read(&mut self, offset: u64, size: u32) -> Result { - let inner = &mut self.inner; + async fn range_read(&self, offset: u64, size: u32) -> Result { + let inner = &self.inner; self.cache .get_or_load( (self.file_id, self.column_id), @@ -102,7 +102,7 @@ impl BloomFilterReader for CachedBloomFilterIndexBl } /// Reads the meta information of the bloom filter. - async fn metadata(&mut self) -> Result { + async fn metadata(&self) -> Result { if let Some(cached) = self.cache.get_metadata((self.file_id, self.column_id)) { CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc(); Ok((*cached).clone()) diff --git a/src/mito2/src/cache/index/inverted_index.rs b/src/mito2/src/cache/index/inverted_index.rs index 53929c467a9..aaedcd8f89f 100644 --- a/src/mito2/src/cache/index/inverted_index.rs +++ b/src/mito2/src/cache/index/inverted_index.rs @@ -12,11 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. +use core::ops::Range; use std::sync::Arc; use api::v1::index::InvertedIndexMetas; use async_trait::async_trait; use bytes::Bytes; +use futures::future::try_join_all; use index::inverted_index::error::Result; use index::inverted_index::format::reader::InvertedIndexReader; use prost::Message; @@ -77,8 +79,8 @@ impl CachedInvertedIndexBlobReader { #[async_trait] impl InvertedIndexReader for CachedInvertedIndexBlobReader { - async fn range_read(&mut self, offset: u64, size: u32) -> Result> { - let inner = &mut self.inner; + async fn range_read(&self, offset: u64, size: u32) -> Result> { + let inner = &self.inner; self.cache .get_or_load( self.file_id, @@ -90,7 +92,25 @@ impl InvertedIndexReader for CachedInvertedIndexBlobRead .await } - async fn metadata(&mut self) -> Result> { + async fn read_vec(&self, ranges: &[Range]) -> Result> { + let fetch = ranges.iter().map(|range| { + let inner = &self.inner; + self.cache.get_or_load( + self.file_id, + self.blob_size, + range.start, + (range.end - range.start) as u32, + move |ranges| async move { inner.read_vec(&ranges).await }, + ) + }); + Ok(try_join_all(fetch) + .await? + .into_iter() + .map(Bytes::from) + .collect::>()) + } + + async fn metadata(&self) -> Result> { if let Some(cached) = self.cache.get_metadata(self.file_id) { CACHE_HIT.with_label_values(&[INDEX_METADATA_TYPE]).inc(); Ok(cached) @@ -220,7 +240,7 @@ mod test { .unwrap(); let reader = InvertedIndexBlobReader::new(range_reader); - let mut cached_reader = CachedInvertedIndexBlobReader::new( + let cached_reader = CachedInvertedIndexBlobReader::new( FileId::random(), file_size, reader, @@ -304,7 +324,7 @@ mod test { let offset = rng.gen_range(0..file_size); let size = rng.gen_range(0..file_size as u32 - offset as u32); let expected = cached_reader.range_read(offset, size).await.unwrap(); - let inner = &mut cached_reader.inner; + let inner = &cached_reader.inner; let read = cached_reader .cache .get_or_load( diff --git a/src/mito2/src/sst/index/bloom_filter/applier.rs b/src/mito2/src/sst/index/bloom_filter/applier.rs index 1e7533a7044..2f024d5cb71 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier.rs @@ -118,7 +118,7 @@ impl BloomFilterIndexApplier { .start_timer(); for (column_id, predicates) in &self.filters { - let mut blob = match self + let blob = match self .blob_reader(file_id, *column_id, file_size_hint) .await? { diff --git a/src/mito2/src/sst/index/bloom_filter/creator.rs b/src/mito2/src/sst/index/bloom_filter/creator.rs index d1b73a0bde2..01c79b408ee 100644 --- a/src/mito2/src/sst/index/bloom_filter/creator.rs +++ b/src/mito2/src/sst/index/bloom_filter/creator.rs @@ -481,7 +481,7 @@ mod tests { .await .unwrap(); let reader = blob_guard.reader().await.unwrap(); - let mut bloom_filter = BloomFilterReaderImpl::new(reader); + let bloom_filter = BloomFilterReaderImpl::new(reader); let metadata = bloom_filter.metadata().await.unwrap(); assert_eq!(metadata.bloom_filter_segments.len(), 10); @@ -510,7 +510,7 @@ mod tests { .await .unwrap(); let reader = blob_guard.reader().await.unwrap(); - let mut bloom_filter = BloomFilterReaderImpl::new(reader); + let bloom_filter = BloomFilterReaderImpl::new(reader); let metadata = bloom_filter.metadata().await.unwrap(); assert_eq!(metadata.bloom_filter_segments.len(), 5); diff --git a/src/mito2/src/sst/index/inverted_index/applier.rs b/src/mito2/src/sst/index/inverted_index/applier.rs index 73785f3264a..61865f76f4c 100644 --- a/src/mito2/src/sst/index/inverted_index/applier.rs +++ b/src/mito2/src/sst/index/inverted_index/applier.rs @@ -126,7 +126,7 @@ impl InvertedIndexApplier { index_not_found_strategy: IndexNotFoundStrategy::ReturnEmpty, }; - let mut blob = match self.cached_blob_reader(file_id).await { + let blob = match self.cached_blob_reader(file_id, file_size_hint).await { Ok(Some(puffin_reader)) => puffin_reader, other => { if let Err(err) = other { @@ -158,7 +158,11 @@ impl InvertedIndexApplier { } /// Creates a blob reader from the cached index file. - async fn cached_blob_reader(&self, file_id: FileId) -> Result> { + async fn cached_blob_reader( + &self, + file_id: FileId, + file_size_hint: Option, + ) -> Result> { let Some(file_cache) = &self.file_cache else { return Ok(None); }; @@ -171,10 +175,12 @@ impl InvertedIndexApplier { let puffin_manager = self.puffin_manager_factory.build(file_cache.local_store()); let puffin_file_name = file_cache.cache_file_path(index_key); + // Adds file size hint to the puffin reader to avoid extra metadata read. let reader = puffin_manager .reader(&puffin_file_name) .await .context(PuffinBuildReaderSnafu)? + .with_file_size_hint(file_size_hint) .blob(INDEX_BLOB_TYPE) .await .context(PuffinReadBlobSnafu)? diff --git a/src/mito2/src/sst/index/puffin_manager.rs b/src/mito2/src/sst/index/puffin_manager.rs index 8bfa25de1e8..498ed254f2f 100644 --- a/src/mito2/src/sst/index/puffin_manager.rs +++ b/src/mito2/src/sst/index/puffin_manager.rs @@ -191,7 +191,7 @@ mod tests { let reader = manager.reader(file_name).await.unwrap(); let blob_guard = reader.blob(blob_key).await.unwrap(); - let mut blob_reader = blob_guard.reader().await.unwrap(); + let blob_reader = blob_guard.reader().await.unwrap(); let meta = blob_reader.metadata().await.unwrap(); let bs = blob_reader.read(0..meta.content_length).await.unwrap(); assert_eq!(&*bs, raw_data); diff --git a/src/mito2/src/sst/index/store.rs b/src/mito2/src/sst/index/store.rs index 7322bd4db49..f37fdc5c7a9 100644 --- a/src/mito2/src/sst/index/store.rs +++ b/src/mito2/src/sst/index/store.rs @@ -19,7 +19,7 @@ use std::task::{Context, Poll}; use async_trait::async_trait; use bytes::{BufMut, Bytes}; -use common_base::range_read::{Metadata, RangeReader}; +use common_base::range_read::{Metadata, RangeReader, SizeAwareRangeReader}; use futures::{AsyncRead, AsyncSeek, AsyncWrite}; use object_store::ObjectStore; use pin_project::pin_project; @@ -266,13 +266,15 @@ pub(crate) struct InstrumentedRangeReader<'a> { file_size_hint: Option, } -#[async_trait] -impl RangeReader for InstrumentedRangeReader<'_> { +impl SizeAwareRangeReader for InstrumentedRangeReader<'_> { fn with_file_size_hint(&mut self, file_size_hint: u64) { self.file_size_hint = Some(file_size_hint); } +} - async fn metadata(&mut self) -> io::Result { +#[async_trait] +impl RangeReader for InstrumentedRangeReader<'_> { + async fn metadata(&self) -> io::Result { match self.file_size_hint { Some(file_size_hint) => Ok(Metadata { content_length: file_size_hint, @@ -286,18 +288,14 @@ impl RangeReader for InstrumentedRangeReader<'_> { } } - async fn read(&mut self, range: Range) -> io::Result { + async fn read(&self, range: Range) -> io::Result { let buf = self.store.reader(&self.path).await?.read(range).await?; self.read_byte_count.inc_by(buf.len() as _); self.read_count.inc_by(1); Ok(buf.to_bytes()) } - async fn read_into( - &mut self, - range: Range, - buf: &mut (impl BufMut + Send), - ) -> io::Result<()> { + async fn read_into(&self, range: Range, buf: &mut (impl BufMut + Send)) -> io::Result<()> { let reader = self.store.reader(&self.path).await?; let size = reader.read_into(buf, range).await?; self.read_byte_count.inc_by(size as _); @@ -305,7 +303,7 @@ impl RangeReader for InstrumentedRangeReader<'_> { Ok(()) } - async fn read_vec(&mut self, ranges: &[Range]) -> io::Result> { + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { let bufs = self .store .reader(&self.path) diff --git a/src/puffin/Cargo.toml b/src/puffin/Cargo.toml index 31c92ba4f97..7116bbef52b 100644 --- a/src/puffin/Cargo.toml +++ b/src/puffin/Cargo.toml @@ -35,4 +35,5 @@ tokio-util.workspace = true uuid.workspace = true [dev-dependencies] +common-base = { workspace = true, features = ["testing"] } common-test-util.workspace = true diff --git a/src/puffin/src/file_format/reader/file.rs b/src/puffin/src/file_format/reader/file.rs index 9ed40a7f181..27a9853acdf 100644 --- a/src/puffin/src/file_format/reader/file.rs +++ b/src/puffin/src/file_format/reader/file.rs @@ -74,14 +74,14 @@ impl PuffinFileReader { #[async_trait] impl<'a, R: RangeReader + 'a> AsyncReader<'a> for PuffinFileReader { - type Reader = PartialReader<&'a mut R>; + type Reader = PartialReader<&'a R>; async fn metadata(&'a mut self) -> Result { if let Some(metadata) = &self.metadata { return Ok(metadata.clone()); } let file_size = self.get_file_size_async().await?; - let mut reader = PuffinFileFooterReader::new(&mut self.source, file_size) + let mut reader = PuffinFileFooterReader::new(&self.source, file_size) .with_prefetch_size(DEFAULT_PREFETCH_SIZE); let metadata = reader.metadata().await?; self.metadata = Some(metadata.clone()); @@ -90,7 +90,7 @@ impl<'a, R: RangeReader + 'a> AsyncReader<'a> for PuffinFileReader { fn blob_reader(&'a mut self, blob_metadata: &BlobMetadata) -> Result { Ok(PartialReader::new( - &mut self.source, + &self.source, blob_metadata.offset as _, blob_metadata.length as _, )) diff --git a/src/puffin/src/file_format/reader/footer.rs b/src/puffin/src/file_format/reader/footer.rs index d0cd1e8ed4f..c7a7eab08b0 100644 --- a/src/puffin/src/file_format/reader/footer.rs +++ b/src/puffin/src/file_format/reader/footer.rs @@ -25,7 +25,7 @@ use crate::file_format::{Flags, FLAGS_SIZE, MAGIC, MAGIC_SIZE, MIN_FILE_SIZE, PA use crate::file_metadata::FileMetadata; /// The default prefetch size for the footer reader. -pub const DEFAULT_PREFETCH_SIZE: u64 = 1024; // 1KiB +pub const DEFAULT_PREFETCH_SIZE: u64 = 8192; // 8KiB /// Reader for the footer of a Puffin data file /// diff --git a/src/puffin/src/partial_reader.rs b/src/puffin/src/partial_reader.rs index ef481567944..7954f705ef1 100644 --- a/src/puffin/src/partial_reader.rs +++ b/src/puffin/src/partial_reader.rs @@ -13,8 +13,6 @@ // limitations under the License. mod r#async; -mod position; -mod sync; use pin_project::pin_project; @@ -31,13 +29,6 @@ pub struct PartialReader { /// The `offset` and `size` fields are used to determine the slice of `source` to read. #[pin] source: R, - - /// The current position within the portion. - /// - /// A `None` value indicates that no read operations have been performed yet on this portion. - /// Before a read operation can be performed, the resource must be positioned at the correct offset in the portion. - /// After the first read operation, this field will be set to `Some(_)`, representing the current read position in the portion. - position_in_portion: Option, } impl PartialReader { @@ -47,15 +38,9 @@ impl PartialReader { offset, size, source, - position_in_portion: None, } } - /// Returns the current position in the portion. - pub fn position(&self) -> u64 { - self.position_in_portion.unwrap_or_default() - } - /// Returns the size of the portion in portion. pub fn size(&self) -> u64 { self.size @@ -65,11 +50,6 @@ impl PartialReader { pub fn is_empty(&self) -> bool { self.size == 0 } - - /// Returns whether the current position is at the end of the portion. - pub fn is_eof(&self) -> bool { - self.position() == self.size - } } #[cfg(test)] @@ -83,7 +63,6 @@ mod tests { let data: Vec = (0..100).collect(); let reader = PartialReader::new(Cursor::new(data), 10, 0); assert!(reader.is_empty()); - assert!(reader.is_eof()); } #[test] diff --git a/src/puffin/src/partial_reader/async.rs b/src/puffin/src/partial_reader/async.rs index 4eedd1ee31f..0b9a5c54c71 100644 --- a/src/puffin/src/partial_reader/async.rs +++ b/src/puffin/src/partial_reader/async.rs @@ -23,17 +23,13 @@ use crate::partial_reader::PartialReader; #[async_trait] impl RangeReader for PartialReader { - fn with_file_size_hint(&mut self, _file_size_hint: u64) { - // do nothing - } - - async fn metadata(&mut self) -> io::Result { + async fn metadata(&self) -> io::Result { Ok(Metadata { content_length: self.size, }) } - async fn read(&mut self, range: Range) -> io::Result { + async fn read(&self, range: Range) -> io::Result { let absolute_range_start = self.offset + range.start; if absolute_range_start >= self.offset + self.size { return Err(io::Error::new( @@ -45,15 +41,10 @@ impl RangeReader for PartialReader { let absolute_range = absolute_range_start..absolute_range_end; let result = self.source.read(absolute_range.clone()).await?; - self.position_in_portion = Some(absolute_range.end); Ok(result) } - async fn read_into( - &mut self, - range: Range, - buf: &mut (impl BufMut + Send), - ) -> io::Result<()> { + async fn read_into(&self, range: Range, buf: &mut (impl BufMut + Send)) -> io::Result<()> { let absolute_range_start = self.offset + range.start; if absolute_range_start >= self.offset + self.size { return Err(io::Error::new( @@ -65,11 +56,10 @@ impl RangeReader for PartialReader { let absolute_range = absolute_range_start..absolute_range_end; self.source.read_into(absolute_range.clone(), buf).await?; - self.position_in_portion = Some(absolute_range.end); Ok(()) } - async fn read_vec(&mut self, ranges: &[Range]) -> io::Result> { + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { let absolute_ranges = ranges .iter() .map(|range| { @@ -88,9 +78,6 @@ impl RangeReader for PartialReader { .collect::>>()?; let results = self.source.read_vec(&absolute_ranges).await?; - if let Some(last_range) = absolute_ranges.last() { - self.position_in_portion = Some(last_range.end); - } Ok(results) } @@ -103,7 +90,7 @@ mod tests { #[tokio::test] async fn read_all_data_in_portion() { let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(data.clone(), 0, 100); + let reader = PartialReader::new(data.clone(), 0, 100); let buf = reader.read(0..100).await.unwrap(); assert_eq!(*buf, data); } @@ -111,7 +98,7 @@ mod tests { #[tokio::test] async fn read_part_of_data_in_portion() { let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(data, 10, 30); + let reader = PartialReader::new(data, 10, 30); let buf = reader.read(0..30).await.unwrap(); assert_eq!(*buf, (10..40).collect::>()); } @@ -119,7 +106,7 @@ mod tests { #[tokio::test] async fn seek_past_end_of_portion_returns_error() { let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(data, 10, 30); + let reader = PartialReader::new(data, 10, 30); // seeking past the portion returns an error assert!(reader.read(31..32).await.is_err()); } @@ -127,11 +114,7 @@ mod tests { #[tokio::test] async fn is_eof_returns_true_at_end_of_portion() { let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(data, 10, 30); - // we are not at the end of the portion - assert!(!reader.is_eof()); + let reader = PartialReader::new(data, 10, 30); let _ = reader.read(0..20).await.unwrap(); - // we are at the end of the portion - assert!(reader.is_eof()); } } diff --git a/src/puffin/src/partial_reader/position.rs b/src/puffin/src/partial_reader/position.rs deleted file mode 100644 index e57817c493a..00000000000 --- a/src/puffin/src/partial_reader/position.rs +++ /dev/null @@ -1,102 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed 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 std::io; - -/// Calculates the new position after seeking. It checks if the new position -/// is valid (within the portion bounds) before returning it. -pub fn position_after_seek( - seek_from: io::SeekFrom, - position_in_portion: u64, - size_of_portion: u64, -) -> io::Result { - let new_position = match seek_from { - io::SeekFrom::Start(offset) => offset, - io::SeekFrom::Current(offset) => { - let next = (position_in_portion as i64) + offset; - if next < 0 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "invalid seek to a negative or overflowing position", - )); - } - next as u64 - } - io::SeekFrom::End(offset) => { - let end = size_of_portion as i64; - (end + offset) as u64 - } - }; - - if new_position > size_of_portion { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "invalid seek to a position beyond the end of the portion", - )); - } - - Ok(new_position) -} - -#[cfg(test)] -mod tests { - use std::io::ErrorKind; - - use super::*; - - #[test] - fn test_position_after_seek_from_start() { - let result = position_after_seek(io::SeekFrom::Start(10), 0, 20).unwrap(); - assert_eq!(result, 10); - } - - #[test] - fn test_position_after_seek_from_start_out_of_bounds() { - let result = position_after_seek(io::SeekFrom::Start(30), 0, 20); - assert!(result.is_err()); - assert_eq!(result.unwrap_err().kind(), ErrorKind::InvalidInput); - } - - #[test] - fn test_position_after_seek_from_current() { - let result = position_after_seek(io::SeekFrom::Current(10), 10, 30).unwrap(); - assert_eq!(result, 20); - } - - #[test] - fn test_position_after_seek_from_current_negative_position_within_bounds() { - let result = position_after_seek(io::SeekFrom::Current(-10), 15, 20).unwrap(); - assert_eq!(result, 5); - } - - #[test] - fn test_position_after_seek_from_current_negative_position() { - let result = position_after_seek(io::SeekFrom::Current(-10), 5, 20); - assert!(result.is_err()); - assert_eq!(result.unwrap_err().kind(), ErrorKind::InvalidInput); - } - - #[test] - fn test_position_after_seek_from_end() { - let result = position_after_seek(io::SeekFrom::End(-10), 0, 30).unwrap(); - assert_eq!(result, 20); - } - - #[test] - fn test_position_after_seek_from_end_out_of_bounds() { - let result = position_after_seek(io::SeekFrom::End(10), 0, 20); - assert!(result.is_err()); - assert_eq!(result.unwrap_err().kind(), ErrorKind::InvalidInput); - } -} diff --git a/src/puffin/src/partial_reader/sync.rs b/src/puffin/src/partial_reader/sync.rs deleted file mode 100644 index 1b778154397..00000000000 --- a/src/puffin/src/partial_reader/sync.rs +++ /dev/null @@ -1,180 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed 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 std::io; - -use crate::partial_reader::position::position_after_seek; -use crate::partial_reader::PartialReader; - -impl io::Read for PartialReader { - fn read(&mut self, buf: &mut [u8]) -> io::Result { - // past end of portion - if self.position() > self.size() { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "invalid read past the end of the portion", - )); - } - - // end of portion - if self.is_eof() { - return Ok(0); - } - - // haven't read from the portion yet, need to seek to the start of it. - if self.position_in_portion.is_none() { - self.source.seek(io::SeekFrom::Start(self.offset))?; - self.position_in_portion = Some(0); - } - - // prevent reading over the end - let max_len = (self.size() - self.position_in_portion.unwrap()) as usize; - let actual_len = max_len.min(buf.len()); - - // create a limited reader - let target_buf = &mut buf[..actual_len]; - - // perform the actual read from the source and update the position. - let read_bytes = self.source.read(target_buf)?; - self.position_in_portion = Some(self.position_in_portion.unwrap() + read_bytes as u64); - - Ok(read_bytes) - } -} - -impl io::Seek for PartialReader { - fn seek(&mut self, pos: io::SeekFrom) -> io::Result { - let new_position = position_after_seek(pos, self.position(), self.size())?; - let pos = io::SeekFrom::Start(self.offset + new_position); - self.source.seek(pos)?; - - self.position_in_portion = Some(new_position); - Ok(new_position) - } -} - -#[cfg(test)] -mod tests { - use std::io::{Cursor, Read, Seek, SeekFrom}; - - use super::*; - - #[test] - fn read_all_data_in_portion() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data.clone()), 0, 100); - let mut buf = vec![0; 100]; - assert_eq!(reader.read(&mut buf).unwrap(), 100); - assert_eq!(buf, data); - } - - #[test] - fn read_part_of_data_in_portion() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - let mut buf = vec![0; 30]; - assert_eq!(reader.read(&mut buf).unwrap(), 30); - assert_eq!(buf, (10..40).collect::>()); - } - - #[test] - fn seek_and_read_data_in_portion() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - assert_eq!(reader.seek(SeekFrom::Start(10)).unwrap(), 10); - let mut buf = vec![0; 10]; - assert_eq!(reader.read(&mut buf).unwrap(), 10); - assert_eq!(buf, (20..30).collect::>()); - } - - #[test] - fn read_past_end_of_portion_is_eof() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - let mut buf = vec![0; 50]; - assert_eq!(reader.read(&mut buf).unwrap(), 30); - assert_eq!(reader.read(&mut buf).unwrap(), 0); // hit EOF - } - - #[test] - fn seek_past_end_of_portion_returns_error() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - // seeking past the portion returns an error - assert!(reader.seek(SeekFrom::Start(31)).is_err()); - } - - #[test] - fn seek_to_negative_position_returns_error() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - assert_eq!(reader.seek(SeekFrom::Start(10)).unwrap(), 10); - // seeking back to the start of the portion - assert_eq!(reader.seek(SeekFrom::Current(-10)).unwrap(), 0); - // seeking to a negative position returns an error - assert!(reader.seek(SeekFrom::Current(-1)).is_err()); - } - - #[test] - fn seek_from_end_of_portion() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - let mut buf = vec![0; 10]; - // seek to 10 bytes before the end of the portion - assert_eq!(reader.seek(SeekFrom::End(-10)).unwrap(), 20); - assert_eq!(reader.read(&mut buf).unwrap(), 10); - // the final 10 bytes of the portion - assert_eq!(buf, (30..40).collect::>()); - assert!(reader.is_eof()); - } - - #[test] - fn seek_from_end_to_negative_position_returns_error() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data.clone()), 10, 30); - // seeking to a negative position returns an error - assert!(reader.seek(SeekFrom::End(-31)).is_err()); - } - - #[test] - fn zero_length_portion_returns_zero_on_read() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 0); - let mut buf = vec![0; 10]; - // reading a portion with zero length returns 0 bytes - assert_eq!(reader.read(&mut buf).unwrap(), 0); - } - - #[test] - fn is_eof_returns_true_at_end_of_portion() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - // we are not at the end of the portion - assert!(!reader.is_eof()); - let mut buf = vec![0; 30]; - assert_eq!(reader.read(&mut buf).unwrap(), 30); - // we are at the end of the portion - assert!(reader.is_eof()); - } - - #[test] - fn position_resets_after_seek_to_start() { - let data: Vec = (0..100).collect(); - let mut reader = PartialReader::new(Cursor::new(data), 10, 30); - assert_eq!(reader.seek(SeekFrom::Start(10)).unwrap(), 10); - assert_eq!(reader.position(), 10); - assert_eq!(reader.seek(SeekFrom::Start(0)).unwrap(), 0); - assert_eq!(reader.position(), 0); - } -} diff --git a/src/puffin/src/puffin_manager/file_accessor.rs b/src/puffin/src/puffin_manager/file_accessor.rs index dc32db6db50..351423b054e 100644 --- a/src/puffin/src/puffin_manager/file_accessor.rs +++ b/src/puffin/src/puffin_manager/file_accessor.rs @@ -13,7 +13,7 @@ // limitations under the License. use async_trait::async_trait; -use common_base::range_read::RangeReader; +use common_base::range_read::SizeAwareRangeReader; use futures::AsyncWrite; use crate::error::Result; @@ -22,7 +22,7 @@ use crate::error::Result; #[async_trait] #[auto_impl::auto_impl(Arc)] pub trait PuffinFileAccessor: Send + Sync + 'static { - type Reader: RangeReader + Sync; + type Reader: SizeAwareRangeReader + Sync; type Writer: AsyncWrite + Unpin + Send; /// Opens a reader for the given puffin file. diff --git a/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs b/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs index a5da2f75f85..9228c0b5942 100644 --- a/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs +++ b/src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use async_compression::futures::bufread::ZstdDecoder; use async_trait::async_trait; use bytes::{BufMut, Bytes}; -use common_base::range_read::{AsyncReadAdapter, Metadata, RangeReader}; +use common_base::range_read::{AsyncReadAdapter, Metadata, RangeReader, SizeAwareRangeReader}; use futures::io::BufReader; use futures::{AsyncRead, AsyncWrite}; use snafu::{ensure, OptionExt, ResultExt}; @@ -202,7 +202,7 @@ where .find(|m| m.blob_type == key.as_str()) .context(BlobNotFoundSnafu { blob: key })?; - let mut reader = file.blob_reader(blob_metadata)?; + let reader = file.blob_reader(blob_metadata)?; let meta = reader.metadata().await.context(MetadataSnafu)?; let buf = reader .read(0..meta.content_length) @@ -315,36 +315,25 @@ where A: RangeReader, B: RangeReader, { - fn with_file_size_hint(&mut self, file_size_hint: u64) { - match self { - Either::L(a) => a.with_file_size_hint(file_size_hint), - Either::R(b) => b.with_file_size_hint(file_size_hint), - } - } - - async fn metadata(&mut self) -> io::Result { + async fn metadata(&self) -> io::Result { match self { Either::L(a) => a.metadata().await, Either::R(b) => b.metadata().await, } } - async fn read(&mut self, range: Range) -> io::Result { + async fn read(&self, range: Range) -> io::Result { match self { Either::L(a) => a.read(range).await, Either::R(b) => b.read(range).await, } } - async fn read_into( - &mut self, - range: Range, - buf: &mut (impl BufMut + Send), - ) -> io::Result<()> { + async fn read_into(&self, range: Range, buf: &mut (impl BufMut + Send)) -> io::Result<()> { match self { Either::L(a) => a.read_into(range, buf).await, Either::R(b) => b.read_into(range, buf).await, } } - async fn read_vec(&mut self, ranges: &[Range]) -> io::Result> { + async fn read_vec(&self, ranges: &[Range]) -> io::Result> { match self { Either::L(a) => a.read_vec(ranges).await, Either::R(b) => b.read_vec(ranges).await, diff --git a/src/puffin/src/puffin_manager/stager/bounded_stager.rs b/src/puffin/src/puffin_manager/stager/bounded_stager.rs index 19fd891438b..476eaf9bb62 100644 --- a/src/puffin/src/puffin_manager/stager/bounded_stager.rs +++ b/src/puffin/src/puffin_manager/stager/bounded_stager.rs @@ -548,7 +548,7 @@ mod tests { let puffin_file_name = "test_get_blob"; let key = "key"; - let mut reader = stager + let reader = stager .get_blob( puffin_file_name, key, @@ -685,7 +685,7 @@ mod tests { .await .unwrap(); - let mut reader = stager + let reader = stager .get_blob( puffin_file_name, blob_key, @@ -732,7 +732,7 @@ mod tests { let blob_key = "blob_key"; // First time to get the blob - let mut reader = stager + let reader = stager .get_blob( puffin_file_name, blob_key, @@ -758,7 +758,7 @@ mod tests { assert_eq!(&*buf, b"Hello world"); // Second time to get the blob, get from recycle bin - let mut reader = stager + let reader = stager .get_blob( puffin_file_name, blob_key, diff --git a/src/puffin/src/puffin_manager/tests.rs b/src/puffin/src/puffin_manager/tests.rs index 37b3ce7c65a..c4057a5f5bc 100644 --- a/src/puffin/src/puffin_manager/tests.rs +++ b/src/puffin/src/puffin_manager/tests.rs @@ -297,7 +297,7 @@ async fn check_blob( compressed: bool, ) { let blob = puffin_reader.blob(key).await.unwrap(); - let mut reader = blob.reader().await.unwrap(); + let reader = blob.reader().await.unwrap(); let meta = reader.metadata().await.unwrap(); let bs = reader.read(0..meta.content_length).await.unwrap(); assert_eq!(&*bs, raw_data); diff --git a/src/puffin/src/tests.rs b/src/puffin/src/tests.rs index a3bb4858792..5eb31511c89 100644 --- a/src/puffin/src/tests.rs +++ b/src/puffin/src/tests.rs @@ -57,7 +57,7 @@ async fn test_read_puffin_file_metadata_async() { "src/tests/resources/sample-metric-data-uncompressed.puffin", ]; for path in paths { - let mut reader = FileReader::new(path).await.unwrap(); + let reader = FileReader::new(path).await.unwrap(); let file_size = reader.metadata().await.unwrap().content_length; let mut reader = PuffinFileReader::new(reader); let metadata = reader.metadata().await.unwrap();