use std::sync::Arc;
use bytes::Bytes;
use futures::future::{BoxFuture, FutureExt};
use parquet::{
arrow::{arrow_reader::ArrowReaderOptions, async_reader::AsyncFileReader},
errors::Result as ParquetResult,
file::metadata::{PageIndexPolicy, ParquetMetaData},
};
use crate::data_file::{ParquetFile, auxiliary_cache};
#[derive(Clone)]
pub(super) struct ObjectStoreReader {
pub(super) file: ParquetFile,
pub(super) page_index: PageIndexPolicy,
}
impl ObjectStoreReader {
pub(super) fn new(file: &ParquetFile, page_index: PageIndexPolicy) -> Self {
Self {
file: file.clone(),
page_index,
}
}
}
fn footer_prefetch_size(footer_size: u64, file_size: u64) -> Option<usize> {
(footer_size > 0)
.then_some(footer_size)
.and_then(|size| size.checked_add(8))
.filter(|size| *size <= file_size)
.and_then(|size| usize::try_from(size).ok())
}
impl AsyncFileReader for ObjectStoreReader {
fn get_bytes(&mut self, range: std::ops::Range<u64>) -> BoxFuture<'_, ParquetResult<Bytes>> {
let file = self.file.clone();
async move { auxiliary_cache::shared().range(&file, range).await }.boxed()
}
fn get_byte_ranges(
&mut self,
ranges: Vec<std::ops::Range<u64>>,
) -> BoxFuture<'_, ParquetResult<Vec<Bytes>>> {
let file = self.file.clone();
async move { auxiliary_cache::shared().ranges(&file, ranges).await }.boxed()
}
fn get_metadata<'a>(
&'a mut self,
_options: Option<&'a ArrowReaderOptions>,
) -> BoxFuture<'a, ParquetResult<Arc<ParquetMetaData>>> {
let prefetch = footer_prefetch_size(self.file.footer_size, self.file.file_size);
async move { auxiliary_cache::shared().metadata(self, prefetch).await }.boxed()
}
}