rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-26 21:56:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: introduce ParallelFstValuesMapper (#5276)

Browse Source 
* refactor: `RangeReader` to use `&self`

* refactor: `InvertedIndexReader` to use `&self`

* refactor: refactor: `BloomFilterReader` to use `&self`

* feat: introduce `ParallelFstValuesMapper`

* chore: change prefetch size to 8KiB

* chore: add `file_size_hint` for cached blob reader

* chore: fix clippy

* refactor: remove `FstValuesMapper`

* chore: apply suggestions from CR
This commit is contained in:
Weny Xu
2025-01-06 07:33:35 +00:00
committed by GitHub
parent bbbba29afc
commit 96b2a5fb28
30 changed files with 359 additions and 555 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/base/Cargo.toml
+3
View File
@@ -4,6 +4,9 @@ version.workspace = true
edition.workspace = true
license.workspace = true
[features]
testing = []
[lints]
workspace = true
src/common/base/src/range_read.rs
+46 -49
View File
@@ -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<Metadata>;
async fn metadata(&self) -> io::Result<Metadata>;
/// Reads the bytes in the given range.
async fn read(&mut self, range: Range<u64>) -> io::Result<Bytes>;
async fn read(&self, range: Range<u64>) -> io::Result<Bytes>;
/// 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<u8>`)
/// - Panic (e.g., for `&mut [u8]`)
async fn read_into(
&mut self,
range: Range<u64>,
buf: &mut (impl BufMut + Send),
) -> io::Result<()> {
async fn read_into(&self, range: Range<u64>, 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<u64>]) -> io::Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
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<R: ?Sized + RangeReader> 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<Metadata> {
impl<R: ?Sized + RangeReader> RangeReader for &R {
async fn metadata(&self) -> io::Result<Metadata> {
(*self).metadata().await
}
async fn read(&mut self, range: Range<u64>) -> io::Result<Bytes> {
async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
(*self).read(range).await
}
async fn read_into(
&mut self,
range: Range<u64>,
buf: &mut (impl BufMut + Send),
) -> io::Result<()> {
async fn read_into(&self, range: Range<u64>, buf: &mut (impl BufMut + Send)) -> io::Result<()> {
(*self).read_into(range, buf).await
}
async fn read_vec(&mut self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
(*self).read_vec(ranges).await
}
}
@@ -120,7 +115,7 @@ pub struct AsyncReadAdapter<R> {
impl<R: RangeReader + 'static> AsyncReadAdapter<R> {
pub async fn new(inner: R) -> io::Result<Self> {
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<R: RangeReader + 'static> AsyncRead for AsyncReadAdapter<R> {
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<R: RangeReader + 'static> AsyncRead for AsyncReadAdapter<R> {
#[async_trait]
impl RangeReader for Vec<u8> {
fn with_file_size_hint(&mut self, _file_size_hint: u64) {
// do nothing
}
async fn metadata(&mut self) -> io::Result<Metadata> {
async fn metadata(&self) -> io::Result<Metadata> {
Ok(Metadata {
content_length: self.len() as u64,
})
}
async fn read(&mut self, range: Range<u64>) -> io::Result<Bytes> {
async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
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<tokio::fs::File>,
}
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<Metadata> {
async fn metadata(&self) -> io::Result<Metadata> {
Ok(Metadata {
content_length: self.content_length,
})
}
async fn read(&mut self, mut range: Range<u64>) -> io::Result<Bytes> {
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<u64>) -> io::Result<Bytes> {
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);
src/index/src/bloom_filter/applier.rs
+1 -1
View File
@@ -42,7 +42,7 @@ pub struct BloomFilterApplier {
}
impl BloomFilterApplier {
pub async fn new(mut reader: Box<dyn BloomFilterReader + Send>) -> Result<Self> {
pub async fn new(reader: Box<dyn BloomFilterReader + Send>) -> Result<Self> {
let meta = reader.metadata().await?;
Ok(Self { reader, meta })
src/index/src/bloom_filter/reader.rs
+11 -11
View File
@@ -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<Bytes>;
async fn range_read(&self, offset: u64, size: u32) -> Result<Bytes>;
/// Reads bunch of ranges from the file.
async fn read_vec(&mut self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
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<BloomFilterMeta>;
async fn metadata(&self) -> Result<BloomFilterMeta>;
/// Reads a bloom filter with the given location.
async fn bloom_filter(&mut self, loc: &BloomFilterSegmentLocation) -> Result<BloomFilter> {
async fn bloom_filter(&self, loc: &BloomFilterSegmentLocation) -> Result<BloomFilter> {
let bytes = self.range_read(loc.offset, loc.size as _).await?;
let vec = bytes
.chunks_exact(std::mem::size_of::<u64>())
@@ -80,23 +80,23 @@ impl<R: RangeReader> BloomFilterReaderImpl<R> {
#[async_trait]
impl<R: RangeReader> BloomFilterReader for BloomFilterReaderImpl<R> {
async fn range_read(&mut self, offset: u64, size: u32) -> Result<Bytes> {
async fn range_read(&self, offset: u64, size: u32) -> Result<Bytes> {
self.reader
.read(offset..offset + size as u64)
.await
.context(IoSnafu)
}
async fn read_vec(&mut self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
self.reader.read_vec(ranges).await.context(IoSnafu)
}
async fn metadata(&mut self) -> Result<BloomFilterMeta> {
async fn metadata(&self) -> Result<BloomFilterMeta> {
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);
src/index/src/inverted_index/format/reader.rs
+25 -5
View File
@@ -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<Vec<u8>>;
async fn range_read(&self, offset: u64, size: u32) -> Result<Vec<u8>>;
/// Reads the bytes in the given ranges.
async fn read_vec(&mut self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
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<Arc<InvertedIndexMetas>>;
async fn metadata(&self) -> Result<Arc<InvertedIndexMetas>>;
/// Retrieves the finite state transducer (FST) map from the given offset and size.
async fn fst(&mut self, offset: u64, size: u32) -> Result<FstMap> {
async fn fst(&self, offset: u64, size: u32) -> Result<FstMap> {
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<u64>]) -> Result<Vec<FstMap>> {
self.read_vec(ranges)
.await?
.into_iter()
.map(|bytes| FstMap::new(bytes.to_vec()).context(DecodeFstSnafu))
.collect::<Result<Vec<_>>>()
}
/// Retrieves the bitmap from the given offset and size.
async fn bitmap(&mut self, offset: u64, size: u32) -> Result<BitVec> {
async fn bitmap(&self, offset: u64, size: u32) -> Result<BitVec> {
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<u64>]) -> Result<VecDeque<BitVec>> {
Ok(self
.read_vec(ranges)
.await?
.into_iter()
.map(|bytes| BitVec::from_slice(bytes.as_ref()))
.collect::<VecDeque<_>>())
}
}
src/index/src/inverted_index/format/reader/blob.rs
+7 -7
View File
@@ -52,7 +52,7 @@ impl<R> InvertedIndexBlobReader<R> {
#[async_trait]
impl<R: RangeReader + Sync> InvertedIndexReader for InvertedIndexBlobReader<R> {
async fn range_read(&mut self, offset: u64, size: u32) -> Result<Vec<u8>> {
async fn range_read(&self, offset: u64, size: u32) -> Result<Vec<u8>> {
let buf = self
.source
.read(offset..offset + size as u64)
@@ -61,16 +61,16 @@ impl<R: RangeReader + Sync> InvertedIndexReader for InvertedIndexBlobReader<R> {
Ok(buf.into())
}
async fn read_vec(&mut self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
self.source.read_vec(ranges).await.context(CommonIoSnafu)
}
async fn metadata(&mut self) -> Result<Arc<InvertedIndexMetas>> {
async fn metadata(&self) -> Result<Arc<InvertedIndexMetas>> {
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();
src/index/src/inverted_index/format/reader/footer.rs
+6 -6
View File
@@ -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<R> {
@@ -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);
}
src/index/src/inverted_index/format/writer/blob.rs
+2 -2
View File
@@ -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);
src/index/src/inverted_index/search/fst_values_mapper.rs
+108 -56
View File
@@ -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<BitVec> {
let mut bitmap = BitVec::new();
pub async fn map_values_vec(
&mut self,
value_and_meta_vec: &[(Vec<u64>, &'a InvertedIndexMeta)],
) -> Result<Vec<BitVec>> {
let groups = value_and_meta_vec
.iter()
.map(|(values, _)| values.len())
.collect::<Vec<_>>();
let len = groups.iter().sum::<usize>();
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::<u64, [u32; 2]>(*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::<u64, [u32; 2]>(*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]);
}
}
src/index/src/inverted_index/search/index_apply/predicates_apply.rs
+73 -35
View File
@@ -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::<Vec<_>>();
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<Vec<(String, Vec<Predicate>)>> 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)
src/mito2/src/cache/index.rs
Vendored
+1 -1
View File
@@ -157,7 +157,7 @@ where
load: F,
) -> Result<Vec<u8>, E>
where
F: FnOnce(Vec<Range<u64>>) -> Fut,
F: Fn(Vec<Range<u64>>) -> Fut,
Fut: Future<Output = Result<Vec<Bytes>, E>>,
E: std::error::Error,
{
src/mito2/src/cache/index/bloom_filter_index.rs
Vendored
+3 -3
View File
@@ -87,8 +87,8 @@ impl<R> CachedBloomFilterIndexBlobReader<R> {
#[async_trait]
impl<R: BloomFilterReader + Send> BloomFilterReader for CachedBloomFilterIndexBlobReader<R> {
async fn range_read(&mut self, offset: u64, size: u32) -> Result<Bytes> {
let inner = &mut self.inner;
async fn range_read(&self, offset: u64, size: u32) -> Result<Bytes> {
let inner = &self.inner;
self.cache
.get_or_load(
(self.file_id, self.column_id),
@@ -102,7 +102,7 @@ impl<R: BloomFilterReader + Send> BloomFilterReader for CachedBloomFilterIndexBl
}
/// Reads the meta information of the bloom filter.
async fn metadata(&mut self) -> Result<BloomFilterMeta> {
async fn metadata(&self) -> Result<BloomFilterMeta> {
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())
src/mito2/src/cache/index/inverted_index.rs
Vendored
+25 -5
View File
@@ -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<R> CachedInvertedIndexBlobReader<R> {
#[async_trait]
impl<R: InvertedIndexReader> InvertedIndexReader for CachedInvertedIndexBlobReader<R> {
async fn range_read(&mut self, offset: u64, size: u32) -> Result<Vec<u8>> {
let inner = &mut self.inner;
async fn range_read(&self, offset: u64, size: u32) -> Result<Vec<u8>> {
let inner = &self.inner;
self.cache
.get_or_load(
self.file_id,
@@ -90,7 +92,25 @@ impl<R: InvertedIndexReader> InvertedIndexReader for CachedInvertedIndexBlobRead
.await
}
async fn metadata(&mut self) -> Result<Arc<InvertedIndexMetas>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
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::<Vec<_>>())
}
async fn metadata(&self) -> Result<Arc<InvertedIndexMetas>> {
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(
src/mito2/src/sst/index/bloom_filter/applier.rs
+1 -1
View File
@@ -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?
{
src/mito2/src/sst/index/bloom_filter/creator.rs
+2 -2
View File
@@ -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);
src/mito2/src/sst/index/inverted_index/applier.rs
+8 -2
View File
@@ -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<Option<BlobReader>> {
async fn cached_blob_reader(
&self,
file_id: FileId,
file_size_hint: Option<u64>,
) -> Result<Option<BlobReader>> {
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)?
src/mito2/src/sst/index/puffin_manager.rs
+1 -1
View File
@@ -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);
src/mito2/src/sst/index/store.rs
+9 -11
View File
@@ -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<u64>,
}
#[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<Metadata> {
#[async_trait]
impl RangeReader for InstrumentedRangeReader<'_> {
async fn metadata(&self) -> io::Result<Metadata> {
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<u64>) -> io::Result<Bytes> {
async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
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<u64>,
buf: &mut (impl BufMut + Send),
) -> io::Result<()> {
async fn read_into(&self, range: Range<u64>, 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<u64>]) -> io::Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
let bufs = self
.store
.reader(&self.path)
src/puffin/Cargo.toml
+1
View File
@@ -35,4 +35,5 @@ tokio-util.workspace = true
uuid.workspace = true
[dev-dependencies]
common-base = { workspace = true, features = ["testing"] }
common-test-util.workspace = true
src/puffin/src/file_format/reader/file.rs
+3 -3
View File
@@ -74,14 +74,14 @@ impl<R> PuffinFileReader<R> {
#[async_trait]
impl<'a, R: RangeReader + 'a> AsyncReader<'a> for PuffinFileReader<R> {
type Reader = PartialReader<&'a mut R>;
type Reader = PartialReader<&'a R>;
async fn metadata(&'a mut self) -> Result<FileMetadata> {
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<R> {
fn blob_reader(&'a mut self, blob_metadata: &BlobMetadata) -> Result<Self::Reader> {
Ok(PartialReader::new(
&mut self.source,
&self.source,
blob_metadata.offset as _,
blob_metadata.length as _,
))
src/puffin/src/file_format/reader/footer.rs
+1 -1
View File
@@ -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
///
src/puffin/src/partial_reader.rs
-21
View File
@@ -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<R> {
/// 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<u64>,
}
impl<R> PartialReader<R> {
@@ -47,15 +38,9 @@ impl<R> PartialReader<R> {
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<R> PartialReader<R> {
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<u8> = (0..100).collect();
let reader = PartialReader::new(Cursor::new(data), 10, 0);
assert!(reader.is_empty());
assert!(reader.is_eof());
}
#[test]
src/puffin/src/partial_reader/async.rs
+8 -25
View File
@@ -23,17 +23,13 @@ use crate::partial_reader::PartialReader;
#[async_trait]
impl<R: RangeReader> RangeReader for PartialReader<R> {
fn with_file_size_hint(&mut self, _file_size_hint: u64) {
// do nothing
}
async fn metadata(&mut self) -> io::Result<Metadata> {
async fn metadata(&self) -> io::Result<Metadata> {
Ok(Metadata {
content_length: self.size,
})
}
async fn read(&mut self, range: Range<u64>) -> io::Result<Bytes> {
async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
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<R: RangeReader> RangeReader for PartialReader<R> {
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<u64>,
buf: &mut (impl BufMut + Send),
) -> io::Result<()> {
async fn read_into(&self, range: Range<u64>, 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<R: RangeReader> RangeReader for PartialReader<R> {
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<u64>]) -> io::Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
let absolute_ranges = ranges
.iter()
.map(|range| {
@@ -88,9 +78,6 @@ impl<R: RangeReader> RangeReader for PartialReader<R> {
.collect::<io::Result<Vec<_>>>()?;
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<u8> = (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<u8> = (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::<Vec<u8>>());
}
@@ -119,7 +106,7 @@ mod tests {
#[tokio::test]
async fn seek_past_end_of_portion_returns_error() {
let data: Vec<u8> = (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<u8> = (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());
}
}
src/puffin/src/partial_reader/position.rs
-102
View File
@@ -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<u64> {
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);
}
}
src/puffin/src/partial_reader/sync.rs
-180
View File
@@ -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<R: io::Read + io::Seek> io::Read for PartialReader<R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
// 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<R: io::Read + io::Seek> io::Seek for PartialReader<R> {
fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
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<u8> = (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<u8> = (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::<Vec<u8>>());
}
#[test]
fn seek_and_read_data_in_portion() {
let data: Vec<u8> = (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::<Vec<u8>>());
}
#[test]
fn read_past_end_of_portion_is_eof() {
let data: Vec<u8> = (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<u8> = (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<u8> = (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<u8> = (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::<Vec<u8>>());
assert!(reader.is_eof());
}
#[test]
fn seek_from_end_to_negative_position_returns_error() {
let data: Vec<u8> = (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<u8> = (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<u8> = (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<u8> = (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);
}
}
src/puffin/src/puffin_manager/file_accessor.rs
+2 -2
View File
@@ -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.
src/puffin/src/puffin_manager/fs_puffin_manager/reader.rs
+6 -17
View File
@@ -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<Metadata> {
async fn metadata(&self) -> io::Result<Metadata> {
match self {
Either::L(a) => a.metadata().await,
Either::R(b) => b.metadata().await,
}
}
async fn read(&mut self, range: Range<u64>) -> io::Result<Bytes> {
async fn read(&self, range: Range<u64>) -> io::Result<Bytes> {
match self {
Either::L(a) => a.read(range).await,
Either::R(b) => b.read(range).await,
}
}
async fn read_into(
&mut self,
range: Range<u64>,
buf: &mut (impl BufMut + Send),
) -> io::Result<()> {
async fn read_into(&self, range: Range<u64>, 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<u64>]) -> io::Result<Vec<Bytes>> {
async fn read_vec(&self, ranges: &[Range<u64>]) -> io::Result<Vec<Bytes>> {
match self {
Either::L(a) => a.read_vec(ranges).await,
Either::R(b) => b.read_vec(ranges).await,
src/puffin/src/puffin_manager/stager/bounded_stager.rs
+4 -4
View File
@@ -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,
src/puffin/src/puffin_manager/tests.rs
+1 -1
View File
@@ -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);
src/puffin/src/tests.rs
+1 -1
View File
@@ -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();