Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

object-store-based Store implementation #1661

Open
wants to merge 25 commits into
base: main
Choose a base branch
from

Conversation

kylebarron
Copy link

@kylebarron kylebarron commented Feb 8, 2024

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
Copy link
Member

jhamman commented Feb 8, 2024

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

@kylebarron
Copy link
Author

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
Copy link
Member

martindurant/rfsspec#3

async def get_partial_values(
self, key_ranges: List[Tuple[str, Tuple[int, int]]]
) -> List[bytes]:
# TODO: use rust-based concurrency inside object-store
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Author

@kylebarron kylebarron Feb 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@normanrz
Copy link
Contributor

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

@martindurant
Copy link
Member

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
Copy link
Author

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
Copy link
Member

jhamman commented Feb 12, 2024

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
Copy link
Member

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
Copy link
Contributor

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.

@jhamman jhamman added the V3 Affects the v3 branch label Feb 13, 2024
@itsgifnotjiff
Copy link

This is awesome work, thank you all!!!

@kylebarron
Copy link
Author

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
Copy link
Member

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
Copy link
Author

kylebarron commented Oct 22, 2024

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
Copy link
Contributor

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?

@normanrz
Copy link
Contributor

The suffix request is required for sharding.

Is there an example of a suffix request somewhere in this repo, so we can see how the range is passed as an argument to the store?

https://github.com/zarr-developers/zarr-python/blob/main/src/zarr/codecs/sharding.py#L700-L702

@kylebarron
Copy link
Author

There's probably a bug or two in get_partial_values that I'd like to fix on Monday, but otherwise this PR should be ready to test with obstore v0.3.0-beta.2 if anyone is so inclined.

It's particularly nice that we're able to match the AsyncGenerator[str] semantics of the list operations.

@kylebarron
Copy link
Author

kylebarron commented Nov 7, 2024

I rewrote the implementation of get_partial_values, which should now support offset and suffix requests as well.

The type hinting in this PR passes. I'd love to get someone excited about trying this out.

I figure this may not be desired for merge into zarr-python proper. Anyone have a preference where this should live?

@martindurant
Copy link
Member

I'd love to get someone excited about trying this out.

Is there an obstore release with all of your latest changes?

@kylebarron
Copy link
Author

Is there an obstore release with all of your latest changes?

You can use https://pypi.org/project/obstore/0.3.0b2/ for now

@kylebarron kylebarron changed the title Prototype of object-store-based Store implementation object-store-based Store implementation Nov 18, 2024
@kylebarron
Copy link
Author

I'm not familiar enough with hatch to know what's going wrong with the CI

@TomAugspurger
Copy link
Contributor

https://github.com/zarr-developers/zarr-python/actions/runs/11931641418/job/33291263276?pr=1661, which failed with

Creating environment: test.py3.13-1.25-minimal
Installing project in development mode
ERROR: Ignored the following versions that require a different python version: 1.21.2 Requires-Python >=3.7,<3.11; 1.21.3 Requires-Python >=3.7,<3.11; 1.21.4 Requires-Python >=3.7,<3.11; 1.21.5 Requires-Python >=3.7,<3.11; 1.21.6 Requires-Python >=3.7,<3.11; 1.26.0 Requires-Python <3.13,>=3.9; 1.26.1 Requires-Python <3.13,>=3.9
ERROR: Could not find a version that satisfies the requirement obstore==0.3.0b5; extra == "test" (from zarr[extra,test]) (from versions: 0.1.0b1)

seems to be because there isn't a wheel for 3.13 at https://pypi.org/project/obstore/0.3.0b5/#files.

I think that https://github.com/zarr-developers/zarr-python/actions/runs/11931641418/job/33291259870?pr=1661 has just our required dependencies. Your test file will need something like a pytest.importorskip("obstore") before you import any module importing obstore

@kylebarron
Copy link
Author

I updated this to obstore 0.3.0-beta.8, which includes wheels for Python 3.13 (except for Windows, which only has up through 3.12)

@jhamman
Copy link
Member

jhamman commented Nov 22, 2024

@kylebarron - thanks for the continued work on this. We discussed this PR in our dev call today. Here's a brief summary of the outcome:

  1. We want to include this in Zarr-Python so long as obstore is an optional dependency
  2. We want to document the store as experimental at first (in the docstring)
  3. Over time, we will figure out a way to allow the user to select which remote store is used when a string like s3://bucket/my.zarr is used but for now, we'll continue to use the fsspec-based RemoteStore for this. For now, users will have to construct the ObjectStore themselves.

How does this sound to you?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
V3 Affects the v3 branch
Projects
Status: In review
Development

Successfully merging this pull request may close these issues.

10 participants