object-store-based Store implementation

[kylebarron]

A Zarr store based on obstore, which is a Python library that uses the Rust object_store crate under the hood.

object-store is a rust crate for interoperating with remote object stores like S3, GCS, Azure, etc. See the highlights section of its docs.

obstore maps async Rust functions to async Python functions, and is able to stream GET and LIST requests, which all make it a good candidate for use with the Zarr v3 Store protocol.

You should be able to test this branch with the latest pre-release version of obstore:

pip install --pre --upgrade obstore

TODO:

• Examples
• Add unit tests and/or doctests in docstrings
• Add docstrings and API docs for any new/modified user-facing classes and functions
• New/modified features documented in docs/tutorial.rst
• Changes documented in docs/release.rst
• GitHub Actions have all passed
• Test coverage is 100% (Codecov passes)

[jhamman]

Amazing @kylebarron! I'll spend some time playing with this today.

[kylebarron]

With roeap/object-store-python#9 it should be possible to fetch multiple ranges within a file concurrently with range coalescing (using get_ranges_async). Note that this object-store API accepts multiple ranges within one object, which is still not 100% aligned with the Zarr get_partial_values because that allows fetches across multiple objects.

That PR also adds a get_opts function which now supports "offset" and "suffix" ranges, of the sort Range:N- and Range:-N, which would allow removing the raise NotImplementedError on line 37.

[martindurant]

martindurant/rfsspec#3

[martindurant]

How I did it in rfsspec: https://github.com/martindurant/rfsspec/blob/main/src/lib.rs#L141

[kylebarron]

object-store has a built-in function for this: get_ranges. With the caveat that it only manages multiple ranges in a single file.

get_ranges also automatically handles request merging for nearby ranges in a file.

[martindurant]

Yes I know, but mine already did the whole thing, so I am showing how I did that.

[normanrz]

Great work @kylebarron!
What are everbody's thoughts on having this in zarr-python vs. spinning it out as a separate package?

[martindurant]

What are everbody's thoughts on having this in zarr-python vs. spinning it out as a separate package?

I suggest we see whether it makes any improvements first, so it's author's choice for now.

[kylebarron]

While @rabernat has seen some impressive perf improvements in some settings when making many requests with Rust's tokio runtime, which would possibly also trickle down to a Python binding, the biggest advantage I see is improved ease of use in installation.

A common hurdle I've seen is handling dependency management, especially around boto3, aioboto3, etc dependencies. Versions need to be compatible at runtime with any other libraries the user also has in their environment. And Python doesn't allow multiple versions of the same dependency at the same time in one environment. With a Python library wrapping a statically-linked Rust binary, you can remove all Python dependencies and remove this class of hardship.

The underlying Rust object-store crate is stable and under open governance via the Apache Arrow project. We'll just have to wait on some discussion in object-store-python for exactly where that should live.

I don't have an opinion myself on where this should live, but it should be on the order of 100 lines of code wherever it is (unless the v3 store api changes dramatically)

[jhamman]

I suggest we see whether it makes any improvements first, so it's author's choice for now.

👍

What are everbody's thoughts on having this in zarr-python vs. spinning it out as a separate package?

I want to keep an open mind about what the core stores provided by Zarr-Python are. My current thinking is that we should just do a MemoryStore and a LocalFilesystemStore. Everything else can be opt-in by installing a 3rd party package. That said, I like having a few additional stores in the mix as we develop the store interface since it helps us think about the design more broadly.

[martindurant]

A common hurdle I've seen is handling dependency management, especially around boto3, aioboto3, etc dependencies.

This is no longer an issue, s3fs has much more relaxed deps than it used to. Furthermore, it's very likely to be already part of an installation environment.

[normanrz]

I want to keep an open mind about what the core stores provided by Zarr-Python are. My current thinking is that we should just do a MemoryStore and a LocalFilesystemStore. Everything else can be opt-in by installing a 3rd party package.

I agree with that. I think it is beneficial to keep the number of dependencies of core zarr-python small. But, I am open for discussion.

That said, I like having a few additional stores in the mix as we develop the store interface since it helps us think about the design more broadly.

Sure! That is certainly useful.

[itsgifnotjiff]

This is awesome work, thank you all!!!

[kylebarron]

The object-store-python package is not very well maintained roeap/object-store-python#24, so I took a few days to implement my own wrapper around the Rust object_store crate: https://github.com/developmentseed/object-store-rs

I'd like to update this PR soonish to use that library instead.

[martindurant]

If the zarr group prefers object-store-rs, we can move it into the zarr-developers org, if you like. I would like to be involved in developing it, particularly if it can grow more explicit fsspec compatible functionality.

[kylebarron]

I have a few questions because the Store API has changed a bit since the spring.

• There's a new BufferPrototype object. Is the BufferPrototype chosen by the store implementation or the caller? It would be very nice if this prototype could be chosen by the store implementation, because then we could return a RustBuffer object that implements the Python buffer protocol, but doesn't need to copy the buffer into Python memory.
• Similarly for puts, is Buffer guaranteed to implement the buffer protocol? Contrary to fetching, we can't do zero-copy puts right now with object-store

I like that list now returns an AsyncGenerator. That aligns well with the underlying object-store rust API, but for technical reasons we can't expose that as an async iterable to Python yet (apache/arrow-rs#6587), even though we do expose the readable stream to Python as an async iterable.

[TomAugspurger]

Is the BufferPrototype chosen by the store implementation or the caller? It would be very nice if this prototype could be chosen by the store implementation, because then we could return a RustBuffer object that implements the Python buffer protocol, but doesn't need to copy the buffer into Python memory.

This came up in the discussion at https://github.com/zarr-developers/zarr-python/pull/2426/files/5e0ffe80d039d9261517d96ce87220ce8d48e4f2#diff-bb6bb03f87fe9491ef78156256160d798369749b4b35c06d4f275425bdb6c4ad. By default, it's passed as default_buffer_prototype though I think the user can override at the call site or globally.

Does it look compatible with what you need?

[maxrjones]

I'm making some progress on the failing tests in kylebarron#3 (currently up to 37/67 passing from 3/67).

I need some help understanding the expected behavior of this test:

zarr-python/src/zarr/testing/store.py

Lines 297 to 315 in 122760f

	async def test_list_dir(self, store: S) -> None:
	root = "foo"
	store_dict = {
	root + "/zarr.json": self.buffer_cls.from_bytes(b"bar"),
	root + "/c/1": self.buffer_cls.from_bytes(b"\x01"),
	}
	assert await _collect_aiterator(store.list_dir("")) == ()
	assert await _collect_aiterator(store.list_dir(root)) == ()
	await store._set_many(store_dict.items())
	keys_observed = await _collect_aiterator(store.list_dir(root))
	keys_expected = {k.removeprefix(root + "/").split("/")[0] for k in store_dict}
	assert sorted(keys_observed) == sorted(keys_expected)
	keys_observed = await _collect_aiterator(store.list_dir(root + "/"))
	assert sorted(keys_expected) == sorted(keys_observed)

The test writes the following local Zarr store.

.
└── foo
    ├── c
    │   └── 1
    └── zarr.json

The expected keys for store.list_dir('foo') are {'zarr.json', 'c'}. The returned keys for ObjectStore is ('zarr.json',). The docs state Retrieve all keys and prefixes with a given prefix and which do not contain the character “/” after the given prefix.. To me the obstore behavior seems correct and the fsspec behavior seems wrong based on this description because c is a folder (local store) and contains "/" after the given prefix. @jhamman @d-v-b could you please provide guidance here?

[jhamman]

Suggested change

	async def set_if_not_exists(self, key: str, value: Buffer) -> None:
	buf = value.to_bytes()
	await obs.put_async(self.store, key, buf, mode="create")
	async def set_if_not_exists(self, key: str, value: Buffer) -> None:
	buf = value.to_bytes()
	try:
	await obs.put_async(self.store, key, buf, mode="create")
	except AlreadyExistsError:
	pass

[maxrjones]

thanks Joe, pre-commit recommended a slightly different version (kylebarron@86951b8) but this fixed one of the failing tests

[maxrjones]

Leaving an update that we're down to 6 failing tests in kylebarron#3. I probably won't have time to work on the tests again for several days at least.

Test 1

The store will need to be make serializable for Dask to work
FAILED tests/test_store/test_object.py::TestObjectStore::test_serializable_store - TypeError: cannot pickle 'builtins.LocalStore' object

Test 2-4

Obstore has the end byte on a range request as exclusive whereas Zarr-Python tests have this as inclusive. I want to confirm the intended behavior before proposing any more changes.

FAILED tests/test_store/test_object.py::TestObjectStore::test_get[byte_range3-c/0] - AssertionError
FAILED tests/test_store/test_object.py::TestObjectStore::test_get[byte_range3-foo/c/0.0] - AssertionError
FAILED tests/test_store/test_object.py::TestObjectStore::test_get[byte_range3-foo/0/0] - AssertionError

Test 5

Obstore doesn't like None as the start of the range request. I proposed updating None be the same as 0 for get but refusing None for get_partial_values seems reasonable to me 🤷‍♂️

FAILED tests/test_store/test_object.py::TestObjectStore::test_get_partial_values[key_ranges3] - ValueError: Cannot pass `None` for the start of the range request.

Test 6

See #1661 (comment) for an explanation

FAILED tests/test_store/test_object.py::TestObjectStore::test_list_dir - AssertionError

[maxrjones]

I decided to try out an S3Store in https://gist.github.com/maxrjones/0318a87d6179faf4eacbe01beb6d3eb6 (all previous tests were run with local or memory stores). I ran into an issue where S3Store.from_url only seems to consider the bucket. In the example included in the gist, a store is created with from_url(nasa-veda-scratch/zarr-obstore-test/...). When listing objects within that store, objects with keys starting with nasa-veda-scratch/icechunk-test/... are also included in the results. @kylebarron do you know a way to specify a prefix to use for all objects when creating an ObjectStore?

Just for clarification, I know that list accepts as 'prefix' argument but the prefix I think needs to be configured at the store level since it should apply to all operations including get, set, etc.

[kylebarron]

Thanks for your help @maxrjones!

Test 1

The store will need to be make serializable for Dask to work FAILED tests/test_store/test_object.py::TestObjectStore::test_serializable_store - TypeError: cannot pickle 'builtins.LocalStore' object

I might suggest postponing this until we validate that obstore does provide a performance improvement. I'm not sure how to add pickle support.

Test 2-4

Obstore has the end byte on a range request as exclusive whereas Zarr-Python tests have this as inclusive. I want to confirm the intended behavior before proposing any more changes.

Obstore defines range requests as start and end where the end is exclusive. I do explicitly document that the range is exclusive in the docs. While the HTTP spec defines range request strings as inclusive byte ranges I believe. Perhaps fsspec and zarr-python implement the latter.

Test 5

Obstore doesn't like None as the start of the range request. I proposed updating None be the same as 0 for get but refusing None for get_partial_values seems reasonable to me 🤷‍♂️

I don't think there should be two different ways to specify "the beginning of the buffer". obstore.get_range only supports having integer start and end.

[d-v-b]

I don't think there should be two different ways to specify "the beginning of the buffer".

this also holds for zarr-python -- I think we currently use None and 0 to express the same thing, which should be fixed.

[maxrjones]

Test 2-4
Obstore has the end byte on a range request as exclusive whereas Zarr-Python tests have this as inclusive. I want to confirm the intended behavior before proposing any more changes.

Obstore defines range requests as start and end where the end is exclusive. I do explicitly document that the range is exclusive in the docs. While the HTTP spec defines range request strings as inclusive byte ranges I believe. Perhaps fsspec and zarr-python implement the latter.

Sorry for all the churn here, but I realized the actual difference seems to be that ByteRangeRequest is actually [start, length] rather than [start, end]. I guess this makes the inclusive/exclusive end debate moot, but was still super surprising to me. See #2437 for any further discussions on byte range representations.

[maxrjones]

An update for anyone following this PR following a sync with @kylebarron today - the implementations obstore.store.LocalStore and obstore.store.MemoryStore backed Zarr stores seem good to go, but AzureStore, GCSStore, and S3Stores will require more work to make the Zarr storage backend work relative to a specified root prefix (the object_store rust library only supports relativity to the bucket). Kyle and I plan to pair on finishing up this PR in the new year. I hope the work early next year will unblock performance testing of Zarr V3 + obstore relative to Zarr V2 and Zarr V3 + icechunk, and would like to collaborate with Earthmover + others on those performance tests.