├── tests └── it │ ├── main.rs │ ├── sync.rs │ ├── parquet_s3.rs │ ├── parquet_s3_async.rs │ └── stream.rs ├── .gitignore ├── Cargo.toml ├── .github └── workflows │ └── test.yaml ├── src ├── lib.rs ├── sync.rs └── stream.rs └── README.md /tests/it/main.rs: -------------------------------------------------------------------------------- 1 | mod stream; 2 | mod sync; 3 | 4 | mod parquet_s3; 5 | mod parquet_s3_async; 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "range-reader" 3 | version = "0.2.0" 4 | license = "Apache-2.0" 5 | description = "Converts low-level APIs to read ranges of bytes to `Read + Seek`" 6 | homepage = "https://github.com/DataEngineeringLabs/ranged-reader-rs" 7 | repository = "https://github.com/DataEngineeringLabs/ranged-reader-rs" 8 | authors = ["Jorge C. Leitao "] 9 | edition = "2018" 10 | 11 | [dependencies] 12 | futures = { version = "0.3", optional = true } 13 | 14 | [features] 15 | default = ["sync", "async"] 16 | sync = [] 17 | async = ["futures"] 18 | 19 | [dev-dependencies] 20 | tokio = { version = "1", features = ["macros", "rt", "fs"] } 21 | tokio-util = { version = "0.6", features = ["compat"] } 22 | parquet2 = { version = "0.7" } 23 | rust-s3 = { version = "0.27", features = ["blocking", "futures"] } 24 | -------------------------------------------------------------------------------- /.github/workflows/test.yaml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | test: 7 | name: Test 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v2 11 | - uses: Swatinem/rust-cache@v1 12 | - uses: actions-rs/toolchain@v1 13 | with: 14 | toolchain: stable 15 | - name: test 16 | run: cargo test 17 | 18 | coverage: 19 | name: Coverage 20 | runs-on: ubuntu-latest 21 | steps: 22 | - uses: actions/checkout@v2 23 | - uses: Swatinem/rust-cache@v1 24 | - uses: actions-rs/toolchain@v1 25 | with: 26 | toolchain: stable 27 | - name: Install tarpaulin 28 | run: cargo install cargo-tarpaulin 29 | - name: Run coverage 30 | run: | 31 | cargo tarpaulin --ignore-tests --out Xml 32 | - name: Report coverage 33 | continue-on-error: true 34 | run: bash <(curl -s https://codecov.io/bash) 35 | -------------------------------------------------------------------------------- /tests/it/sync.rs: -------------------------------------------------------------------------------- 1 | use std::io::Read; 2 | 3 | use range_reader::RangedReader; 4 | 5 | fn test(calls: usize, call_size: usize, buffer: usize) { 6 | let length = 100; 7 | let range_fn = Box::new(move |start: usize, buf: &mut [u8]| { 8 | let iter = (start..start + buf.len()).map(|x| x as u8); 9 | buf.iter_mut().zip(iter).for_each(|(x, v)| *x = v); 10 | Ok(()) 11 | }); 12 | 13 | let mut reader = RangedReader::new(length, range_fn, vec![0; buffer]); 14 | 15 | let mut to = vec![0; call_size]; 16 | let mut result = vec![]; 17 | (0..calls).for_each(|i| { 18 | let _ = reader.read(&mut to); 19 | result.extend_from_slice(&to); 20 | assert_eq!( 21 | result, 22 | (0..(i + 1) * call_size) 23 | .map(|x| x as u8) 24 | .collect::>() 25 | ); 26 | }); 27 | } 28 | 29 | #[test] 30 | fn basics() { 31 | test(10, 5, 10); 32 | test(5, 20, 10); 33 | test(10, 7, 10); 34 | } 35 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #![deny(missing_docs)] 2 | #![forbid(unsafe_code)] 3 | //! Offers two structs: 4 | //! * [`RangedReader`], that implements [`std::io::Read`] and [`std::io::Seek`] from a function 5 | //! returning a [`Vec`] from a ranged request `(start, length)`. 6 | //! * [`RangedAsyncReader`] that implements [`futures::io::AsyncRead`] and [`futures::io::AsyncRead`] 7 | //! from a asynchronous function returning (a future of) [`Vec`] from a ranged request `(start, length)`. 8 | //! 9 | //! A common use use-case for this crate is to perform ranged queries from remote blob storages 10 | //! such as AWS s3, Azure blob, and Google's cloud storage. 11 | 12 | #[cfg(feature = "sync")] 13 | #[cfg_attr(docsrs, doc(cfg(feature = "sync")))] 14 | mod sync; 15 | #[cfg(feature = "sync")] 16 | #[cfg_attr(docsrs, doc(cfg(feature = "sync")))] 17 | pub use sync::*; 18 | #[cfg(feature = "async")] 19 | #[cfg_attr(docsrs, doc(cfg(feature = "async")))] 20 | mod stream; 21 | #[cfg(feature = "async")] 22 | #[cfg_attr(docsrs, doc(cfg(feature = "async")))] 23 | pub use stream::*; 24 | -------------------------------------------------------------------------------- /tests/it/parquet_s3.rs: -------------------------------------------------------------------------------- 1 | use parquet2::read::read_metadata; 2 | use range_reader::RangedReader; 3 | use s3::Bucket; 4 | 5 | #[test] 6 | fn main() { 7 | let bucket_name = "dev-jorgecardleitao"; 8 | let region = "eu-central-1".parse().unwrap(); 9 | let bucket = Bucket::new_public(bucket_name, region).unwrap(); 10 | let path = "benches_65536.parquet".to_string(); 11 | 12 | let (data, _) = bucket.head_object_blocking(&path).unwrap(); 13 | let length = data.content_length.unwrap() as usize; 14 | 15 | let range_fn = Box::new(move |start: usize, buf: &mut [u8]| { 16 | let (mut data, _) = bucket 17 | // -1 because ranges are inclusive in `get_object_range` 18 | .get_object_range_blocking( 19 | &path, 20 | start as u64, 21 | Some(start as u64 + buf.len() as u64 - 1), 22 | ) 23 | .map_err(|x| std::io::Error::new(std::io::ErrorKind::Other, x.to_string()))?; 24 | data.truncate(buf.len()); 25 | buf[..data.len()].copy_from_slice(&data); 26 | Ok(()) 27 | }); 28 | 29 | let buffer = 1024 * 4; // 4 kb per request. 30 | 31 | let mut reader = RangedReader::new(length, range_fn, vec![0; buffer]); 32 | 33 | let metadata = read_metadata(&mut reader).unwrap(); 34 | 35 | let num_rows: usize = metadata 36 | .row_groups 37 | .iter() 38 | .map(|group| group.num_rows() as usize) 39 | .sum(); 40 | assert_eq!(num_rows, 524288); 41 | assert_eq!(metadata.row_groups.len(), 1); 42 | } 43 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Ranged reader 2 | 3 | [![test](https://github.com/DataEngineeringLabs/ranged-reader-rs/actions/workflows/test.yaml/badge.svg)](https://github.com/DataEngineeringLabs/ranged-reader-rs/actions/workflows/test.yaml) 4 | [![codecov](https://codecov.io/gh/DataEngineeringLabs/ranged-reader-rs/branch/main/graph/badge.svg?token=AgyTF60R3D)](https://codecov.io/gh/DataEngineeringLabs/ranged-reader-rs) 5 | 6 | Convert low-level APIs to read ranges of files into structs that implement 7 | `Read + Seek` and `AsyncRead + AsyncSeek`. See 8 | [parquet_s3_async.rs](tests/it/parquet_s3_async.rs) for an example of this API to 9 | read parts of a large parquet file from s3 asynchronously. 10 | 11 | ### Rational 12 | 13 | Blob storage https APIs offer the ability to read ranges of bytes from a single blob, 14 | i.e. functions of the form 15 | 16 | ```rust 17 | fn read_range_blocking(path: &str, start: usize, length: usize) -> Vec; 18 | async fn read_range(path: &str, start: usize, length: usize) -> Vec; 19 | ``` 20 | 21 | together with its total size, 22 | 23 | ```rust 24 | async fn length(path: &str) -> usize; 25 | fn length(path: &str) -> usize; 26 | ``` 27 | 28 | These APIs are usually IO-bounded - they wait for network. 29 | 30 | Some file formats (e.g. Apache Parquet, Apache Avro, Apache Arrow IPC) allow reading 31 | parts of a file for filter and projection push down. 32 | 33 | This crate offers 2 structs, `RangedReader` and `RangedStreamer` that implement 34 | `Read + Seek` and `AsyncRead + AsyncSeek` respectively, to bridge the blob storage 35 | APIs mentioned above to the traits used by most Rust APIs to read bytes. 36 | 37 | ## License 38 | 39 | Licensed under either of 40 | 41 | * Apache License, Version 2.0 ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0) 42 | * MIT license ([LICENSE-MIT](LICENSE-MIT) or http://opensource.org/licenses/MIT) 43 | 44 | at your option. 45 | 46 | ### Contribution 47 | 48 | Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in the work by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions. 49 | -------------------------------------------------------------------------------- /tests/it/parquet_s3_async.rs: -------------------------------------------------------------------------------- 1 | use std::io::Result; 2 | use std::sync::Arc; 3 | 4 | use futures::pin_mut; 5 | use futures::{future::BoxFuture, StreamExt}; 6 | use s3::Bucket; 7 | 8 | use parquet2::read::{get_page_stream, read_metadata_async}; 9 | use range_reader::{RangeOutput, RangedAsyncReader}; 10 | 11 | #[tokio::test] 12 | async fn main() -> Result<()> { 13 | let bucket_name = "dev-jorgecardleitao"; 14 | let region = "eu-central-1".parse().unwrap(); 15 | let bucket = Bucket::new_public(bucket_name, region).unwrap(); 16 | let path = "benches_65536.parquet".to_string(); 17 | 18 | let (data, _) = bucket.head_object(&path).await.unwrap(); 19 | let length = data.content_length.unwrap() as usize; 20 | println!("total size in bytes: {}", length); 21 | 22 | let range_get = Box::new(move |start: u64, length: usize| { 23 | let bucket = bucket.clone(); 24 | let path = path.clone(); 25 | Box::pin(async move { 26 | let bucket = bucket.clone(); 27 | let path = path.clone(); 28 | // to get a sense of what is being queried in s3 29 | let (mut data, _) = bucket 30 | // -1 because ranges are inclusive in `get_object_range` 31 | .get_object_range(&path, start, Some(start + length as u64 - 1)) 32 | .await 33 | .map_err(|x| std::io::Error::new(std::io::ErrorKind::Other, x.to_string()))?; 34 | data.truncate(length); 35 | Ok(RangeOutput { start, data }) 36 | }) as BoxFuture<'static, std::io::Result> 37 | }); 38 | 39 | // at least 4kb per s3 request. Adjust as you like. 40 | let mut reader = RangedAsyncReader::new(length, 4 * 1024, range_get); 41 | 42 | let metadata = read_metadata_async(&mut reader) 43 | .await 44 | .map_err(|x| std::io::Error::new(std::io::ErrorKind::Other, x.to_string()))?; 45 | 46 | assert_eq!(524288, metadata.num_rows); 47 | 48 | // pages of the first row group and first column 49 | let column_metadata = &metadata.row_groups[0].columns()[0]; 50 | let pages = get_page_stream(column_metadata, &mut reader, vec![], Arc::new(|_, _| true)) 51 | .await 52 | .map_err(|x| std::io::Error::new(std::io::ErrorKind::Other, x.to_string()))?; 53 | 54 | pin_mut!(pages); 55 | while let Some(maybe_page) = pages.next().await { 56 | let page = maybe_page 57 | .map_err(|x| std::io::Error::new(std::io::ErrorKind::Other, x.to_string()))?; 58 | assert_eq!(page.num_values(), 524288); 59 | } 60 | Ok(()) 61 | } 62 | -------------------------------------------------------------------------------- /src/sync.rs: -------------------------------------------------------------------------------- 1 | use std::io::{Read, Result, Seek, SeekFrom}; 2 | 3 | type FnAPI = Box Result<()>>; 4 | 5 | /// Implements `Read + Seek` for a (blocking) function that reads ranges of bytes. 6 | /// # Implementation 7 | /// This struct has an internal `Vec` that buffers calls. 8 | pub struct RangedReader { 9 | pos: u64, // position of the seek 10 | length: u64, // total size 11 | buffer: Vec, // a ring 12 | offset: usize, // offset in the ring: buffer[:offset] have been read 13 | range_fn: FnAPI, 14 | } 15 | 16 | impl RangedReader { 17 | /// Creates a new [`RangedReader`] with internal buffer `buffer` 18 | pub fn new(length: usize, range_fn: FnAPI, mut buffer: Vec) -> Self { 19 | let length = length as u64; 20 | buffer.clear(); 21 | Self { 22 | pos: 0, 23 | range_fn, 24 | length, 25 | buffer, 26 | offset: 0, 27 | } 28 | } 29 | 30 | fn read_more(&mut self, to_consume: usize) -> Result<()> { 31 | let remaining = self.buffer.len() - self.offset; 32 | 33 | if to_consume < remaining { 34 | return Ok(()); 35 | } 36 | let to_read = std::cmp::max( 37 | std::cmp::max(self.offset, to_consume), 38 | self.buffer.capacity(), 39 | ) - remaining; 40 | 41 | self.buffer.rotate_left(self.offset); 42 | self.buffer.resize(remaining + to_read, 0); 43 | 44 | (self.range_fn)(self.pos as usize, &mut self.buffer[remaining..])?; 45 | self.pos += to_read as u64; 46 | self.offset = 0; 47 | Ok(()) 48 | } 49 | } 50 | 51 | impl Read for RangedReader { 52 | fn read(&mut self, buf: &mut [u8]) -> Result { 53 | let to_consume = buf.len(); 54 | self.read_more(to_consume)?; 55 | 56 | // copy from the internal buffer. 57 | buf[..to_consume].copy_from_slice(&self.buffer[self.offset..self.offset + to_consume]); 58 | // and offset 59 | self.offset += to_consume; 60 | Ok(to_consume) 61 | } 62 | } 63 | 64 | impl Seek for RangedReader { 65 | fn seek(&mut self, pos: SeekFrom) -> Result { 66 | match pos { 67 | SeekFrom::Start(pos) => self.pos = pos, 68 | SeekFrom::End(pos) => self.pos = (self.length as i64 + pos) as u64, 69 | SeekFrom::Current(pos) => self.pos = (self.pos as i64 + pos) as u64, 70 | }; 71 | // todo: optimize: do not clear buffer and instead check whether we can just move the offset. 72 | self.offset = 0; 73 | self.buffer.clear(); 74 | Ok(self.pos) 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /tests/it/stream.rs: -------------------------------------------------------------------------------- 1 | use std::io::SeekFrom; 2 | 3 | use futures::{future::BoxFuture, AsyncReadExt, AsyncSeekExt}; 4 | 5 | use range_reader::{RangeOutput, RangedAsyncReader}; 6 | 7 | // range that returns a vector of `(x % 255) for x in start..start+length` 8 | fn get_mock_range(min_size: usize) -> RangedAsyncReader { 9 | let ranged_future = Box::new(move |start: u64, length: usize| { 10 | Box::pin(async move { 11 | let data = (start..start + length as u64) 12 | .map(|x| (x % 255) as u8) 13 | .collect(); 14 | Ok(RangeOutput { start, data }) 15 | }) as BoxFuture<'static, std::io::Result> 16 | }); 17 | 18 | let length = 100; 19 | RangedAsyncReader::new(length, min_size, ranged_future) 20 | } 21 | 22 | // performs `calls` calls of length `call_size` and checks that they return the expected result. 23 | async fn test(calls: usize, call_size: usize, min_size: usize) { 24 | let mut reader = get_mock_range(min_size); 25 | 26 | let mut to = vec![0; call_size]; 27 | let mut result = vec![]; 28 | for i in 0..calls { 29 | let _ = reader.read(&mut to).await; 30 | result.extend_from_slice(&to); 31 | assert_eq!( 32 | result, 33 | (0..(i + 1) * call_size) 34 | .map(|x| x as u8) 35 | .collect::>() 36 | ); 37 | } 38 | } 39 | 40 | #[tokio::test] 41 | async fn basics() { 42 | test(10, 5, 10).await; 43 | test(5, 20, 10).await; 44 | test(10, 7, 10).await; 45 | } 46 | 47 | // tests that multiple `seek`s followed by `read` are correct 48 | #[tokio::test] 49 | async fn seek_inside() -> std::io::Result<()> { 50 | let mut reader = get_mock_range(100); 51 | 52 | for seek in 0u8..99 { 53 | let length = 100 - seek; 54 | let mut result = vec![0; length as usize]; 55 | reader.seek(SeekFrom::Start(seek as u64)).await?; 56 | reader.read_exact(&mut result).await?; 57 | assert_eq!(result, (seek..100).collect::>()); 58 | } 59 | 60 | Ok(()) 61 | } 62 | 63 | // tests that `seek` followed by `read` is correct 64 | #[tokio::test] 65 | async fn seek_new() -> std::io::Result<()> { 66 | for seek in 0u8..99 { 67 | let mut reader = get_mock_range(100); 68 | let length = 100 - seek; 69 | let mut result = vec![0; length as usize]; 70 | reader.seek(SeekFrom::Start(seek as u64)).await?; 71 | reader.read_exact(&mut result).await?; 72 | assert_eq!(result, (seek..100).collect::>()); 73 | } 74 | 75 | Ok(()) 76 | } 77 | 78 | // tests that reading incomplete overlapping regions work as expected 79 | #[tokio::test] 80 | async fn seek_split() -> std::io::Result<()> { 81 | let mut reader = get_mock_range(2); 82 | let length = 20; 83 | let mut result = vec![0; length as usize]; 84 | reader.read_exact(&mut result).await?; 85 | assert_eq!(result, (0..length).collect::>()); 86 | 87 | reader.seek(SeekFrom::Start(10)).await?; 88 | let mut result = vec![0; length as usize]; 89 | reader.read_exact(&mut result).await?; 90 | assert_eq!(result, (10..10 + length).collect::>()); 91 | 92 | Ok(()) 93 | } 94 | -------------------------------------------------------------------------------- /src/stream.rs: -------------------------------------------------------------------------------- 1 | use std::io::{Result, SeekFrom}; 2 | use std::pin::Pin; 3 | 4 | use futures::{ 5 | future::BoxFuture, 6 | io::{AsyncRead, AsyncSeek}, 7 | Future, 8 | }; 9 | 10 | /// A range of bytes with a known starting position. 11 | #[derive(Debug, Clone, PartialEq, Eq, Hash)] 12 | pub struct RangeOutput { 13 | /// the start 14 | pub start: u64, 15 | /// the data 16 | pub data: Vec, 17 | } 18 | 19 | /// A function that returns a [`BoxFuture`] of [`RangeOutput`]. 20 | /// For example, an async request to return a range of bytes from a blob from the internet. 21 | pub type RangedFuture = 22 | Box BoxFuture<'static, std::io::Result> + Send + Sync>; 23 | 24 | /// A struct that converts [`RangedFuture`] to a `AsyncRead + AsyncSeek` with an internal buffer. 25 | pub struct RangedAsyncReader { 26 | pos: u64, 27 | length: u64, // total size 28 | state: State, 29 | ranged_future: RangedFuture, 30 | min_request_size: usize, // requests have at least this size 31 | } 32 | 33 | enum State { 34 | HasChunk(RangeOutput), 35 | Seeking(BoxFuture<'static, std::io::Result>), 36 | } 37 | 38 | impl RangedAsyncReader { 39 | /// Creates a new [`RangedAsyncReader`]. `length` is the total size of the blob being seeked, 40 | /// `min_request_size` is the minimum number of bytes allowed to be requested to `range_get`. 41 | pub fn new(length: usize, min_request_size: usize, ranged_future: RangedFuture) -> Self { 42 | let length = length as u64; 43 | Self { 44 | pos: 0, 45 | length, 46 | state: State::HasChunk(RangeOutput { 47 | start: 0, 48 | data: vec![], 49 | }), 50 | ranged_future, 51 | min_request_size, 52 | } 53 | } 54 | } 55 | 56 | // whether `test_interval` is inside `a` (start, length). 57 | // a = [ ] 58 | // test = [ ] 59 | // returns true 60 | fn range_includes(a: (usize, usize), test_interval: (usize, usize)) -> bool { 61 | if test_interval.0 < a.0 { 62 | return false; 63 | } 64 | let test_end = test_interval.0 + test_interval.1; 65 | let a_end = a.0 + a.1; 66 | if test_end > a_end { 67 | return false; 68 | } 69 | true 70 | } 71 | 72 | impl AsyncRead for RangedAsyncReader { 73 | fn poll_read( 74 | mut self: std::pin::Pin<&mut Self>, 75 | cx: &mut std::task::Context<'_>, 76 | buf: &mut [u8], 77 | ) -> std::task::Poll> { 78 | let requested_range = (self.pos as usize, buf.len()); 79 | let min_request_size = self.min_request_size; 80 | match &mut self.state { 81 | State::HasChunk(output) => { 82 | let existing_range = (output.start as usize, output.data.len()); 83 | if range_includes(existing_range, requested_range) { 84 | let offset = requested_range.0 - existing_range.0; 85 | buf.copy_from_slice(&output.data[offset..offset + buf.len()]); 86 | self.pos += buf.len() as u64; 87 | std::task::Poll::Ready(Ok(buf.len())) 88 | } else { 89 | let start = requested_range.0 as u64; 90 | let length = std::cmp::max(min_request_size, requested_range.1); 91 | let future = (self.ranged_future)(start, length); 92 | self.state = State::Seeking(Box::pin(future)); 93 | self.poll_read(cx, buf) 94 | } 95 | } 96 | State::Seeking(ref mut future) => match Pin::new(future).poll(cx) { 97 | std::task::Poll::Ready(v) => { 98 | match v { 99 | Ok(output) => self.state = State::HasChunk(output), 100 | Err(e) => return std::task::Poll::Ready(Err(e)), 101 | }; 102 | self.poll_read(cx, buf) 103 | } 104 | std::task::Poll::Pending => std::task::Poll::Pending, 105 | }, 106 | } 107 | } 108 | } 109 | 110 | impl AsyncSeek for RangedAsyncReader { 111 | fn poll_seek( 112 | mut self: std::pin::Pin<&mut Self>, 113 | _: &mut std::task::Context<'_>, 114 | pos: SeekFrom, 115 | ) -> std::task::Poll> { 116 | match pos { 117 | SeekFrom::Start(pos) => self.pos = pos, 118 | SeekFrom::End(pos) => self.pos = (self.length as i64 + pos) as u64, 119 | SeekFrom::Current(pos) => self.pos = (self.pos as i64 + pos) as u64, 120 | }; 121 | std::task::Poll::Ready(Ok(self.pos)) 122 | } 123 | } 124 | --------------------------------------------------------------------------------