feat: Exponential read-ahead cache

Cargo.toml

@@ -18,10 +18,7 @@ num_enum = "0.7.3"
 object_store = { version = "0.12", optional = true }
 reqwest = { version = "0.12", default-features = false, optional = true }
 thiserror = "1"
-tokio = { version = "1.43.0", optional = true, default-features = false, features = [
-    "io-util",
-    "sync",
-] }
+tokio = { version = "1.43.0", default-features = false, features = ["sync"] }
 weezl = "0.1.0"
 
 [dev-dependencies]
@@ -37,7 +34,7 @@ tokio-test = "0.4.4"
 
 [features]
 default = ["object_store", "reqwest"]
-tokio = ["dep:tokio"]
+tokio = ["tokio/io-util"]
 reqwest = ["dep:reqwest"]
 object_store = ["dep:object_store"]
 

python/DEVELOP.md

@@ -1,6 +1,6 @@
 ```
 uv sync --no-install-package async-tiff
-uv run --no-project maturin develop
+uv run --no-project maturin develop --uv
 uv run --no-project mkdocs serve
-uv run --no-project pytest --verbose
+uv run --no-project pytest tests --verbose
 ```

python/README.md

@@ -21,7 +21,7 @@ from async_tiff.store import S3Store
 store = S3Store("naip-visualization", region="us-west-2", request_payer=True)
 path = "ny/2022/60cm/rgb/40073/m_4007307_sw_18_060_20220803.tif"
 
-tiff = await TIFF.open(path, store=store, prefetch=32768)
+tiff = await TIFF.open(path, store=store)
 primary_ifd = tiff.ifds[0]
 
 primary_ifd.geo_key_directory.citation
@@ -68,7 +68,7 @@ from async_tiff.store import S3Store
 store = S3Store("sentinel-cogs", region="us-west-2", skip_signature=True)
 path = "sentinel-s2-l2a-cogs/12/S/UF/2022/6/S2B_12SUF_20220609_0_L2A/B04.tif"
 
-tiff = await TIFF.open(path, store=store, prefetch=32768)
+tiff = await TIFF.open(path, store=store)
 primary_ifd = tiff.ifds[0]
 # Text readable citation
 primary_ifd.geo_key_directory.citation

python/pyproject.toml

@@ -32,13 +32,12 @@ dev-dependencies = [
     "mkdocstrings>=0.27.0",
     "numpy>=1",
     "obstore>=0.5.1",
-    "pip>=24.2",
     "pytest-asyncio>=0.26.0",
     "pytest>=8.3.3",
     "ruff>=0.8.4",
 ]
 
 [tool.pytest.ini_options]
 addopts = "--color=yes"
-asyncio_default_fixture_loop_scope="function"
+asyncio_default_fixture_loop_scope = "function"
 asyncio_mode = "auto"

python/python/async_tiff/_tiff.pyi

@@ -17,13 +17,18 @@ class TIFF:
         *,
         store: ObjectStore | ObspecInput,
         prefetch: int = 32768,
+        multiplier: int | float = 2.0,
     ) -> TIFF:
         """Open a new TIFF.
 
         Args:
             path: The path within the store to read from.
             store: The backend to use for data fetching.
             prefetch: The number of initial bytes to read up front.
+            multiplier: The multiplier to use for readahead size growth. Must be
+                greater than 1.0. For example, for a value of `2.0`, the first metadata
+                read will be of size `prefetch`, and then the next read will be of size
+                `prefetch * 2`.
 
         Returns:
             A TIFF instance.

python/src/error.rs

@@ -0,0 +1,43 @@
+use async_tiff::error::AsyncTiffError;
+use pyo3::create_exception;
+use pyo3::exceptions::PyFileNotFoundError;
+use pyo3::prelude::*;
+
+// Base exception
+// Note that this is named `BaseError` instead of `ObstoreError` to not leak the name "obstore" to
+// other Rust-Python libraries using pyo3-object_store.
+create_exception!(
+    async_tiff,
+    AsyncTiffException,
+    pyo3::exceptions::PyException,
+    "A general error from the underlying Rust async-tiff library."
+);
+
+#[allow(missing_docs)]
+pub enum PyAsyncTiffError {
+    AsyncTiffError(async_tiff::error::AsyncTiffError),
+}
+
+impl From<AsyncTiffError> for PyAsyncTiffError {
+    fn from(value: AsyncTiffError) -> Self {
+        Self::AsyncTiffError(value)
+    }
+}
+
+impl From<PyAsyncTiffError> for PyErr {
+    fn from(error: PyAsyncTiffError) -> Self {
+        match error {
+            PyAsyncTiffError::AsyncTiffError(err) => match err {
+                AsyncTiffError::ObjectStore(err) => match err {
+                    object_store::Error::NotFound { path: _, source: _ } => {
+                        PyFileNotFoundError::new_err(err.to_string())
+                    }
+                    _ => AsyncTiffException::new_err(err.to_string()),
+                },
+                _ => AsyncTiffException::new_err(err.to_string()),
+            },
+        }
+    }
+}
+
+pub type PyAsyncTiffResult<T> = std::result::Result<T, PyAsyncTiffError>;

python/src/lib.rs

@@ -2,6 +2,7 @@
 
 mod decoder;
 mod enums;
+mod error;
 mod geo;
 mod ifd;
 mod reader;

python/src/tiff.rs

@@ -1,13 +1,15 @@
 use std::sync::Arc;
 
-use async_tiff::metadata::{PrefetchBuffer, TiffMetadataReader};
+use async_tiff::metadata::cache::ReadaheadMetadataCache;
+use async_tiff::metadata::TiffMetadataReader;
 use async_tiff::reader::AsyncFileReader;
 use async_tiff::TIFF;
-use pyo3::exceptions::{PyFileNotFoundError, PyIndexError, PyTypeError};
+use pyo3::exceptions::{PyIndexError, PyTypeError};
 use pyo3::prelude::*;
 use pyo3::types::PyType;
 use pyo3_async_runtimes::tokio::future_into_py;
 
+use crate::error::PyAsyncTiffResult;
 use crate::reader::StoreInput;
 use crate::tile::PyTile;
 use crate::PyImageFileDirectory;
@@ -18,31 +20,39 @@ pub(crate) struct PyTIFF {
     reader: Arc<dyn AsyncFileReader>,
 }
 
+async fn open(
+    reader: Arc<dyn AsyncFileReader>,
+    prefetch: u64,
+    multiplier: f64,
+) -> PyAsyncTiffResult<PyTIFF> {
+    let metadata_fetch = ReadaheadMetadataCache::new(reader.clone())
+        .with_initial_size(prefetch)
+        .with_multiplier(multiplier);
+    let mut metadata_reader = TiffMetadataReader::try_open(&metadata_fetch).await?;
+    let ifds = metadata_reader.read_all_ifds(&metadata_fetch).await?;
+    let tiff = TIFF::new(ifds);
+    Ok(PyTIFF { tiff, reader })
+}
+
 #[pymethods]
 impl PyTIFF {
     #[classmethod]
-    #[pyo3(signature = (path, *, store, prefetch=32768))]
+    #[pyo3(signature = (path, *, store, prefetch=32768, multiplier=2.0))]
     fn open<'py>(
         _cls: &'py Bound<PyType>,
         py: Python<'py>,
         path: String,
         store: StoreInput,
         prefetch: u64,
+        multiplier: f64,
     ) -> PyResult<Bound<'py, PyAny>> {
         let reader = store.into_async_file_reader(path);
 
-        let cog_reader = future_into_py(py, async move {
-            let metadata_fetch = PrefetchBuffer::new(reader.clone(), prefetch)
-                .await
-                .map_err(|err| PyFileNotFoundError::new_err(err.to_string()))?;
-            let mut metadata_reader = TiffMetadataReader::try_open(&metadata_fetch).await.unwrap();
-            let ifds = metadata_reader
-                .read_all_ifds(&metadata_fetch)
-                .await
-                .unwrap();
-            let tiff = TIFF::new(ifds);
-            Ok(PyTIFF { tiff, reader })
-        })?;
+        let cog_reader =
+            future_into_py(
+                py,
+                async move { Ok(open(reader, prefetch, multiplier).await?) },
+            )?;
         Ok(cog_reader)
     }
 

python/tests/test_cog.py

@@ -11,7 +11,7 @@ async def test_cog_s3():
     """
     path = "sentinel-s2-l2a-cogs/12/S/UF/2022/6/S2B_12SUF_20220609_0_L2A/B04.tif"
     store = S3Store("sentinel-cogs", region="us-west-2", skip_signature=True)
-    tiff = await TIFF.open(path=path, store=store, prefetch=32768)
+    tiff = await TIFF.open(path=path, store=store)
 
     ifds = tiff.ifds
     assert len(ifds) == 5

src/cog.rs

@@ -27,7 +27,8 @@ mod test {
     use tiff::decoder::{DecodingResult, Limits};
 
     use super::*;
-    use crate::metadata::{PrefetchBuffer, TiffMetadataReader};
+    use crate::metadata::cache::ReadaheadMetadataCache;
+    use crate::metadata::TiffMetadataReader;
     use crate::reader::{AsyncFileReader, ObjectReader};
 
     #[ignore = "local file"]
@@ -37,16 +38,9 @@ mod test {
         let path = object_store::path::Path::parse("m_4007307_sw_18_060_20220803.tif").unwrap();
         let store = Arc::new(LocalFileSystem::new_with_prefix(folder).unwrap());
         let reader = Arc::new(ObjectReader::new(store, path)) as Arc<dyn AsyncFileReader>;
-        let prefetch_reader = PrefetchBuffer::new(reader.clone(), 32 * 1024)
-            .await
-            .unwrap();
-        let mut metadata_reader = TiffMetadataReader::try_open(&prefetch_reader)
-            .await
-            .unwrap();
-        let ifds = metadata_reader
-            .read_all_ifds(&prefetch_reader)
-            .await
-            .unwrap();
+        let cached_reader = ReadaheadMetadataCache::new(reader.clone());
+        let mut metadata_reader = TiffMetadataReader::try_open(&cached_reader).await.unwrap();
+        let ifds = metadata_reader.read_all_ifds(&cached_reader).await.unwrap();
         let tiff = TIFF::new(ifds);
 
         let ifd = &tiff.ifds[1];