Skip to content

Commit 66cf37a

Browse files
committed
feat: obstore support
1 parent 3f39ca8 commit 66cf37a

File tree

7 files changed

+112
-5
lines changed

7 files changed

+112
-5
lines changed

Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ opendal = { version = "0.54.0", features = ["services-http"] }
2323
tokio = { version = "1.41.1", features = ["rt-multi-thread"] }
2424
zarrs_opendal = "0.9.0"
2525
itertools = "0.14.0"
26+
pyo3-object_store = "0.7.0" # object_store 0.12
27+
zarrs_object_store = "0.5.0" # object_store 0.12
2628

2729
[profile.release]
2830
lto = true

README.md

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,11 @@ We export a `ZarrsCodecPipeline` class so that `zarr-python` can use the class b
2424

2525
At the moment, we only support a subset of the `zarr-python` stores:
2626

27-
- [x] [LocalStore](https://zarr.readthedocs.io/en/latest/_autoapi/zarr/storage/index.html#zarr.storage.LocalStore) (FileSystem)
28-
- [FsspecStore](https://zarr.readthedocs.io/en/latest/_autoapi/zarr/storage/index.html#zarr.storage.FsspecStore)
29-
- [x] [HTTPFileSystem](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.implementations.http.HTTPFileSystem)
27+
- [`LocalStore`](https://zarr.readthedocs.io/en/latest/api/zarr/storage/#zarr.storage.LocalStore) (local filesystem)
28+
- [`ObjectStore`](https://zarr.readthedocs.io/en/latest/user-guide/storage/#object-store) (cloud storage)
29+
- [`HTTPFileSystem`](https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.implementations.http.HTTPFileSystem) via [`FsspecStore`](https://zarr.readthedocs.io/en/latest/api/zarr/storage/#zarr.storage.FsspecStore)
3030

3131
A `NotImplementedError` will be raised if a store is not supported.
32-
We intend to support more stores in the future: https://github.com/zarrs/zarrs-python/issues/44.
3332

3433
### Configuration
3534

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ test = [
3030
"aiohttp",
3131
"fsspec>2024",
3232
"numcodecs>=0.16.1",
33+
"obstore>=0.8.2",
3334
"pytest",
3435
"pytest-asyncio",
3536
"pytest-xdist",

src/store.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use pyo3::{
66
exceptions::{PyNotImplementedError, PyValueError},
77
types::{PyAnyMethods, PyStringMethods, PyTypeMethods},
88
};
9+
use pyo3_object_store::PyExternalObjectStore;
910
use zarrs::storage::{
1011
ReadableWritableListableStorage, storage_adapter::async_to_sync::AsyncToSyncStorageAdapter,
1112
};
@@ -14,14 +15,17 @@ use crate::{runtime::tokio_block_on, utils::PyErrExt};
1415

1516
mod filesystem;
1617
mod http;
18+
mod obstore;
1719

1820
pub use self::filesystem::FilesystemStoreConfig;
1921
pub use self::http::HttpStoreConfig;
22+
pub use self::obstore::ObStoreConfig;
2023

2124
#[derive(Debug, Clone)]
2225
pub enum StoreConfig {
2326
Filesystem(FilesystemStoreConfig),
2427
Http(HttpStoreConfig),
28+
ObStore(ObStoreConfig),
2529
// TODO: Add support for more stores
2630
}
2731

@@ -53,6 +57,13 @@ impl<'py> FromPyObject<'_, 'py> for StoreConfig {
5357
))),
5458
}
5559
}
60+
"ObjectStore" => {
61+
let underlying_store = store.getattr("store")?;
62+
let external_object_store: PyExternalObjectStore = underlying_store.extract()?;
63+
let object_store: Arc<dyn zarrs_object_store::object_store::ObjectStore> =
64+
external_object_store.into_dyn();
65+
Ok(StoreConfig::ObStore(ObStoreConfig::new(object_store)))
66+
}
5667
_ => Err(PyErr::new::<PyNotImplementedError, _>(format!(
5768
"zarrs-python does not support {name} stores"
5869
))),
@@ -65,6 +76,7 @@ impl StoreConfig {
6576
match self {
6677
StoreConfig::Filesystem(config) => config.direct_io(flag),
6778
StoreConfig::Http(_config) => (),
79+
StoreConfig::ObStore(_config) => (),
6880
}
6981
}
7082
}
@@ -82,6 +94,7 @@ impl TryFrom<&StoreConfig> for ReadableWritableListableStorage {
8294
match value {
8395
StoreConfig::Filesystem(config) => config.try_into(),
8496
StoreConfig::Http(config) => config.try_into(),
97+
StoreConfig::ObStore(config) => config.try_into(),
8598
}
8699
}
87100
}

src/store/obstore.rs

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
use std::sync::Arc;
2+
3+
use pyo3::PyErr;
4+
use zarrs::storage::{
5+
ReadableWritableListableStorage, storage_adapter::async_to_sync::AsyncToSyncStorageAdapter,
6+
};
7+
use zarrs_object_store::{AsyncObjectStore, object_store::ObjectStore};
8+
9+
use crate::runtime::tokio_block_on;
10+
11+
#[derive(Debug, Clone)]
12+
pub struct ObStoreConfig {
13+
store: Arc<dyn ObjectStore>,
14+
}
15+
16+
impl ObStoreConfig {
17+
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
18+
Self { store }
19+
}
20+
}
21+
22+
impl TryInto<ReadableWritableListableStorage> for &ObStoreConfig {
23+
type Error = PyErr;
24+
25+
fn try_into(self) -> Result<ReadableWritableListableStorage, Self::Error> {
26+
let async_store = Arc::new(AsyncObjectStore::new(self.store.clone()));
27+
let sync_store = Arc::new(AsyncToSyncStorageAdapter::new(
28+
async_store,
29+
tokio_block_on(),
30+
));
31+
Ok(sync_store)
32+
}
33+
}

tests/test_obstore.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
import tempfile
2+
import warnings
3+
4+
import numpy as np
5+
import zarr
6+
from obstore.store import HTTPStore, LocalStore
7+
from zarr.storage import ObjectStore
8+
9+
from .test_zarrs_http import ARR_REF, URL
10+
11+
# Suppress the expected warning about cross-library object store usage
12+
warnings.filterwarnings(
13+
"ignore",
14+
message="Successfully reconstructed a store defined in another Python module",
15+
)
16+
17+
18+
def test_obstore_local_store():
19+
"""Test zarrs-python with obstore LocalStore"""
20+
with tempfile.TemporaryDirectory() as tmpdir:
21+
# Create obstore LocalStore and wrap with zarr's ObjectStore
22+
obstore_local = LocalStore(prefix=tmpdir)
23+
store = ObjectStore(obstore_local, read_only=False)
24+
25+
# Create zarr array with obstore
26+
# Expect a warning about cross-library object store usage
27+
with warnings.catch_warnings():
28+
warnings.simplefilter("ignore", RuntimeWarning)
29+
arr = zarr.open_array(
30+
store=store,
31+
mode="w",
32+
shape=(100, 100),
33+
chunks=(10, 10),
34+
dtype="f4",
35+
)
36+
37+
# Write data
38+
data = np.random.rand(100, 100).astype("f4")
39+
arr[:] = data
40+
41+
# Read back and verify
42+
arr2 = zarr.open_array(store=store, mode="r")
43+
assert arr2.shape == (100, 100)
44+
assert np.allclose(arr2[:], data)
45+
46+
47+
def test_obstore_http():
48+
"""Test zarrs-python with obstore HTTPStore - similar to test_zarrs_http"""
49+
# Create HTTPStore from the test URL
50+
http_store = HTTPStore.from_url(URL)
51+
store = ObjectStore(http_store, read_only=True)
52+
53+
with warnings.catch_warnings():
54+
warnings.simplefilter("ignore", RuntimeWarning)
55+
arr = zarr.open_array(store=store, mode="r")
56+
57+
# Verify shape and data match the reference
58+
assert arr.shape == (8, 8)
59+
assert np.allclose(arr[:], ARR_REF, equal_nan=True)

tests/test_zarrs_http.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
]
2020
)
2121

22-
URL = "https://raw.githubusercontent.com/LDeakin/zarrs/main/zarrs/tests/data/array_write_read.zarr/group/array"
22+
URL = "https://raw.githubusercontent.com/zarrs/zarrs/main/zarrs/tests/data/array_write_read.zarr/group/array"
2323

2424

2525
def test_zarrs_http():

0 commit comments

Comments
 (0)