Prefetch reader (#18)

Author: kylebarron

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<PyType>,
         py: Python<'py>,
         path: String,
         store: PyObjectStore,
+        prefetch: Option<u64>,
     ) -> PyResult<Bound<'py, PyAny>> {
         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<dyn AsyncFileReader> = 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)
     }

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)

src/async_reader.rs

@@ -114,6 +114,44 @@ impl AsyncFileReader for ObjectReader {
     }
 }
 
+pub struct PrefetchReader {
+    reader: Box<dyn AsyncFileReader>,
+    buffer: Bytes,
+}
+
+impl PrefetchReader {
+    pub async fn new(mut reader: Box<dyn AsyncFileReader>, prefetch: u64) -> Result<Self> {
+        let buffer = reader.get_bytes(0..prefetch).await?;
+        Ok(Self { reader, buffer })
+    }
+}
+
+impl AsyncFileReader for PrefetchReader {
+    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
+        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<Range<u64>>) -> BoxFuture<'_, Result<Vec<Bytes>>>
+    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]

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};