diff --git a/python/src/tiff.rs b/python/src/tiff.rs index af7623c..33d4f4b 100644 --- a/python/src/tiff.rs +++ b/python/src/tiff.rs @@ -1,4 +1,4 @@ -use async_tiff::{COGReader, ObjectReader}; +use async_tiff::{AsyncFileReader, COGReader, ObjectReader, PrefetchReader}; use pyo3::prelude::*; use pyo3::types::PyType; use pyo3_async_runtimes::tokio::future_into_py; @@ -12,16 +12,27 @@ pub(crate) struct PyTIFF(COGReader); #[pymethods] impl PyTIFF { #[classmethod] - #[pyo3(signature = (path, *, store))] + #[pyo3(signature = (path, *, store, prefetch=16384))] fn open<'py>( _cls: &'py Bound, py: Python<'py>, path: String, store: PyObjectStore, + prefetch: Option, ) -> PyResult> { let reader = ObjectReader::new(store.into_inner(), path.into()); + let cog_reader = future_into_py(py, async move { - Ok(PyTIFF(COGReader::try_open(Box::new(reader)).await.unwrap())) + let reader: Box = if let Some(prefetch) = prefetch { + Box::new( + PrefetchReader::new(Box::new(reader), prefetch) + .await + .unwrap(), + ) + } else { + Box::new(reader) + }; + Ok(PyTIFF(COGReader::try_open(reader).await.unwrap())) })?; Ok(cog_reader) } diff --git a/python/tests/test_cog.py b/python/tests/test_cog.py index 3cf0fe4..5fd9fb5 100644 --- a/python/tests/test_cog.py +++ b/python/tests/test_cog.py @@ -6,11 +6,16 @@ path = "sentinel-s2-l2a-cogs/12/S/UF/2022/6/S2B_12SUF_20220609_0_L2A/B04.tif" # 2 min, 15s -tiff = await TIFF.open(path, store=store) +tiff = await TIFF.open(path, store=store, prefetch=32768) ifds = tiff.ifds() ifd = ifds[0] +ifd.compression ifd.tile_height ifd.tile_width ifd.photometric_interpretation gkd = ifd.geo_key_directory gkd.citation +gkd.projected_type +gkd.citation + +dir(gkd) diff --git a/src/async_reader.rs b/src/async_reader.rs index 946f5d1..3bbf577 100644 --- a/src/async_reader.rs +++ b/src/async_reader.rs @@ -114,6 +114,44 @@ impl AsyncFileReader for ObjectReader { } } +pub struct PrefetchReader { + reader: Box, + buffer: Bytes, +} + +impl PrefetchReader { + pub async fn new(mut reader: Box, prefetch: u64) -> Result { + let buffer = reader.get_bytes(0..prefetch).await?; + Ok(Self { reader, buffer }) + } +} + +impl AsyncFileReader for PrefetchReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, Result> { + if range.start < self.buffer.len() as _ { + if range.end < self.buffer.len() as _ { + let usize_range = range.start as usize..range.end as usize; + let result = self.buffer.slice(usize_range); + async { Ok(result) }.boxed() + } else { + // TODO: reuse partial internal buffer + self.reader.get_bytes(range) + } + } else { + self.reader.get_bytes(range) + } + } + + fn get_byte_ranges(&mut self, ranges: Vec>) -> BoxFuture<'_, Result>> + where + Self: Send, + { + // In practice, get_byte_ranges is only used for fetching tiles, which are unlikely to + // overlap a metadata prefetch. + self.reader.get_byte_ranges(ranges) + } +} + #[derive(Debug, Clone, Copy, Default)] pub enum Endianness { #[default] diff --git a/src/lib.rs b/src/lib.rs index f0c3f21..48bb2a4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,6 @@ pub mod error; pub mod geo; mod ifd; -pub use async_reader::{AsyncFileReader, ObjectReader}; +pub use async_reader::{AsyncFileReader, ObjectReader, PrefetchReader}; pub use cog::COGReader; pub use ifd::{ImageFileDirectories, ImageFileDirectory};