Skip to content

Commit

Permalink
api: make buffer public api
Browse files Browse the repository at this point in the history
  • Loading branch information
jhamman committed Jan 7, 2025
1 parent bb8ab0f commit 97ddaf8
Show file tree
Hide file tree
Showing 58 changed files with 159 additions and 151 deletions.
4 changes: 2 additions & 2 deletions docs/user-guide/config.rst
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ This is the current default configuration::
'string': {'name': 'vlen-utf8'}},
'write_empty_chunks': False},
'async': {'concurrency': 10, 'timeout': None},
'buffer': 'zarr.core.buffer.cpu.Buffer',
'buffer': 'zarr.buffer.cpu.Buffer',
'codec_pipeline': {'batch_size': 1,
'path': 'zarr.core.codec_pipeline.BatchedCodecPipeline'},
'codecs': {'blosc': 'zarr.codecs.blosc.BloscCodec',
Expand All @@ -85,5 +85,5 @@ This is the current default configuration::
'zstd': 'zarr.codecs.zstd.ZstdCodec'},
'default_zarr_format': 3,
'json_indent': 2,
'ndbuffer': 'zarr.core.buffer.cpu.NDBuffer',
'ndbuffer': 'zarr.buffer.cpu.NDBuffer',
'threading': {'max_workers': None}}
43 changes: 19 additions & 24 deletions src/zarr/core/buffer/core.py → src/zarr/abc/buffer.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,23 @@
from zarr.codecs.bytes import Endian
from zarr.core.common import BytesLike, ChunkCoords

# Everything here is imported into ``zarr.core.buffer`` namespace.
__all__: list[str] = []
__all__ = [
"ArrayLike",
"Buffer",
"BufferPrototype",
"NDArrayLike",
"NDBuffer",
]


def _check_item_key_is_1d_contiguous(key: Any) -> None:
"""Raises error if `key` isn't a 1d contiguous slice"""
if not isinstance(key, slice):
raise TypeError(
f"Item key has incorrect type (expected slice, got {key.__class__.__name__})"
)
if not (key.step is None or key.step == 1):
raise ValueError("slice must be contiguous")


@runtime_checkable
Expand Down Expand Up @@ -105,16 +120,6 @@ def __eq__(self, other: object) -> Self: # type: ignore[explicit-override, over
"""


def check_item_key_is_1d_contiguous(key: Any) -> None:
"""Raises error if `key` isn't a 1d contiguous slice"""
if not isinstance(key, slice):
raise TypeError(
f"Item key has incorrect type (expected slice, got {key.__class__.__name__})"
)
if not (key.step is None or key.step == 1):
raise ValueError("slice must be contiguous")


class Buffer(ABC):
"""A flat contiguous memory block
Expand Down Expand Up @@ -266,11 +271,11 @@ def to_bytes(self) -> bytes:
return bytes(self.as_numpy_array())

def __getitem__(self, key: slice) -> Self:
check_item_key_is_1d_contiguous(key)
_check_item_key_is_1d_contiguous(key)
return self.__class__(self._data.__getitem__(key))

def __setitem__(self, key: slice, value: Any) -> None:
check_item_key_is_1d_contiguous(key)
_check_item_key_is_1d_contiguous(key)
self._data.__setitem__(key, value)

def __len__(self) -> int:
Expand Down Expand Up @@ -498,13 +503,3 @@ class BufferPrototype(NamedTuple):

buffer: type[Buffer]
nd_buffer: type[NDBuffer]


# The default buffer prototype used throughout the Zarr codebase.
def default_buffer_prototype() -> BufferPrototype:
from zarr.registry import (
get_buffer_class,
get_ndbuffer_class,
)

return BufferPrototype(buffer=get_buffer_class(), nd_buffer=get_ndbuffer_class())
2 changes: 1 addition & 1 deletion src/zarr/abc/codec.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
from abc import abstractmethod
from typing import TYPE_CHECKING, Any, Generic, TypeVar

from zarr.abc.buffer import Buffer, NDBuffer
from zarr.abc.metadata import Metadata
from zarr.core.buffer import Buffer, NDBuffer
from zarr.core.common import ChunkCoords, concurrent_map
from zarr.core.config import config

Expand Down
4 changes: 2 additions & 2 deletions src/zarr/abc/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from itertools import starmap
from typing import TYPE_CHECKING, Protocol, runtime_checkable

from zarr.core.buffer.core import default_buffer_prototype
from zarr.buffer import default_buffer_prototype
from zarr.core.common import concurrent_map
from zarr.core.config import config

Expand All @@ -14,7 +14,7 @@
from types import TracebackType
from typing import Any, Self, TypeAlias

from zarr.core.buffer import Buffer, BufferPrototype
from zarr.abc.buffer import Buffer, BufferPrototype
from zarr.core.common import BytesLike

__all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
Expand Down
2 changes: 1 addition & 1 deletion src/zarr/api/asynchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@
import numpy.typing as npt
from typing_extensions import deprecated

from zarr.abc.buffer import NDArrayLike
from zarr.core.array import Array, AsyncArray, create_array, get_array_metadata
from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
from zarr.core.buffer import NDArrayLike
from zarr.core.common import (
JSON,
AccessModeLiteral,
Expand Down
2 changes: 1 addition & 1 deletion src/zarr/api/synchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import numpy.typing as npt

from zarr.abc.buffer import NDArrayLike
from zarr.abc.codec import Codec
from zarr.api.asynchronous import ArrayLike, PathLike
from zarr.core.array import (
Expand All @@ -25,7 +26,6 @@
ShardsLike,
)
from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
from zarr.core.buffer import NDArrayLike
from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
from zarr.core.common import (
JSON,
Expand Down
17 changes: 17 additions & 0 deletions src/zarr/buffer/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from zarr.abc.buffer import BufferPrototype
from zarr.buffer.cpu import numpy_buffer_prototype

__all__ = [
"default_buffer_prototype",
"numpy_buffer_prototype",
]


# The default buffer prototype used throughout the Zarr codebase.
def default_buffer_prototype() -> BufferPrototype:
from zarr.registry import (
get_buffer_class,
get_ndbuffer_class,
)

return BufferPrototype(buffer=get_buffer_class(), nd_buffer=get_ndbuffer_class())
24 changes: 13 additions & 11 deletions src/zarr/core/buffer/cpu.py → src/zarr/buffer/cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import numpy as np
import numpy.typing as npt

from zarr.core.buffer import core
import zarr.abc.buffer
from zarr.registry import (
register_buffer,
register_ndbuffer,
Expand All @@ -19,11 +19,11 @@
from collections.abc import Callable, Iterable
from typing import Self

from zarr.core.buffer.core import ArrayLike, NDArrayLike
from zarr.abc.buffer import ArrayLike, NDArrayLike
from zarr.core.common import BytesLike


class Buffer(core.Buffer):
class Buffer(zarr.abc.buffer.Buffer):
"""A flat contiguous memory block
We use Buffer throughout Zarr to represent a contiguous block of memory.
Expand Down Expand Up @@ -52,7 +52,7 @@ def create_zero_length(cls) -> Self:
return cls(np.array([], dtype="b"))

@classmethod
def from_buffer(cls, buffer: core.Buffer) -> Self:
def from_buffer(cls, buffer: zarr.abc.buffer.Buffer) -> Self:
"""Create a new buffer of an existing Buffer
This is useful if you want to ensure that an existing buffer is
Expand Down Expand Up @@ -107,7 +107,7 @@ def as_numpy_array(self) -> npt.NDArray[Any]:
"""
return np.asanyarray(self._data)

def __add__(self, other: core.Buffer) -> Self:
def __add__(self, other: zarr.abc.buffer.Buffer) -> Self:
"""Concatenate two buffers"""

other_array = other.as_array_like()
Expand All @@ -117,7 +117,7 @@ def __add__(self, other: core.Buffer) -> Self:
)


class NDBuffer(core.NDBuffer):
class NDBuffer(zarr.abc.buffer.NDBuffer):
"""An n-dimensional memory block
We use NDBuffer throughout Zarr to represent a n-dimensional memory block.
Expand Down Expand Up @@ -186,8 +186,10 @@ def __setitem__(self, key: Any, value: Any) -> None:


def as_numpy_array_wrapper(
func: Callable[[npt.NDArray[Any]], bytes], buf: core.Buffer, prototype: core.BufferPrototype
) -> core.Buffer:
func: Callable[[npt.NDArray[Any]], bytes],
buf: zarr.abc.buffer.Buffer,
prototype: zarr.abc.buffer.BufferPrototype,
) -> zarr.abc.buffer.Buffer:
"""Converts the input of `func` to a numpy array and the output back to `Buffer`.
This function is useful when calling a `func` that only support host memory such
Expand All @@ -214,13 +216,13 @@ def as_numpy_array_wrapper(


# CPU buffer prototype using numpy arrays
buffer_prototype = core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
buffer_prototype = zarr.abc.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
# default_buffer_prototype = buffer_prototype


# The numpy prototype used for E.g. when reading the shard index
def numpy_buffer_prototype() -> core.BufferPrototype:
return core.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
def numpy_buffer_prototype() -> zarr.abc.buffer.BufferPrototype:
return zarr.abc.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)


register_buffer(Buffer)
Expand Down
19 changes: 9 additions & 10 deletions src/zarr/core/buffer/gpu.py → src/zarr/buffer/gpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,7 @@
import numpy as np
import numpy.typing as npt

from zarr.core.buffer import core
from zarr.core.buffer.core import ArrayLike, BufferPrototype, NDArrayLike
import zarr.abc.buffer

if TYPE_CHECKING:
from collections.abc import Iterable
Expand All @@ -26,7 +25,7 @@
cp = None


class Buffer(core.Buffer):
class Buffer(zarr.abc.buffer.Buffer):
"""A flat contiguous memory block on the GPU
We use Buffer throughout Zarr to represent a contiguous block of memory.
Expand All @@ -47,7 +46,7 @@ class Buffer(core.Buffer):
array-like object that must be 1-dim, contiguous, and byte dtype.
"""

def __init__(self, array_like: ArrayLike) -> None:
def __init__(self, array_like: zarr.abc.buffer.ArrayLike) -> None:
if cp is None:
raise ImportError(
"Cannot use zarr.buffer.gpu.Buffer without cupy. Please install cupy."
Expand Down Expand Up @@ -83,7 +82,7 @@ def create_zero_length(cls) -> Self:
return cls(cp.array([], dtype="b"))

@classmethod
def from_buffer(cls, buffer: core.Buffer) -> Self:
def from_buffer(cls, buffer: zarr.abc.buffer.Buffer) -> Self:
"""Create an GPU Buffer given an arbitrary Buffer
This will try to be zero-copy if `buffer` is already on the
GPU and will trigger a copy if not.
Expand All @@ -101,7 +100,7 @@ def from_bytes(cls, bytes_like: BytesLike) -> Self:
def as_numpy_array(self) -> npt.NDArray[Any]:
return cast(npt.NDArray[Any], cp.asnumpy(self._data))

def __add__(self, other: core.Buffer) -> Self:
def __add__(self, other: zarr.abc.buffer.Buffer) -> Self:
other_array = other.as_array_like()
assert other_array.dtype == np.dtype("b")
gpu_other = Buffer(other_array)
Expand All @@ -111,7 +110,7 @@ def __add__(self, other: core.Buffer) -> Self:
)


class NDBuffer(core.NDBuffer):
class NDBuffer(zarr.abc.buffer.NDBuffer):
"""A n-dimensional memory block on the GPU
We use NDBuffer throughout Zarr to represent a n-dimensional memory block.
Expand All @@ -136,7 +135,7 @@ class NDBuffer(core.NDBuffer):
ndarray-like object that is convertible to a regular Numpy array.
"""

def __init__(self, array: NDArrayLike) -> None:
def __init__(self, array: zarr.abc.buffer.NDArrayLike) -> None:
if cp is None:
raise ImportError(
"Cannot use zarr.buffer.gpu.NDBuffer without cupy. Please install cupy."
Expand Down Expand Up @@ -208,10 +207,10 @@ def __getitem__(self, key: Any) -> Self:
def __setitem__(self, key: Any, value: Any) -> None:
if isinstance(value, NDBuffer):
value = value._data
elif isinstance(value, core.NDBuffer):
elif isinstance(value, zarr.abc.buffer.NDBuffer):
gpu_value = NDBuffer(value.as_ndarray_like())
value = gpu_value._data
self._data.__setitem__(key, value)


buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
buffer_prototype = zarr.abc.buffer.BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
4 changes: 2 additions & 2 deletions src/zarr/codecs/blosc.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,15 @@
from numcodecs.blosc import Blosc

from zarr.abc.codec import BytesBytesCodec
from zarr.core.buffer.cpu import as_numpy_array_wrapper
from zarr.buffer.cpu import as_numpy_array_wrapper
from zarr.core.common import JSON, parse_enum, parse_named_configuration
from zarr.registry import register_codec

if TYPE_CHECKING:
from typing import Self

from zarr.abc.buffer import Buffer
from zarr.core.array_spec import ArraySpec
from zarr.core.buffer import Buffer


class BloscShuffle(Enum):
Expand Down
2 changes: 1 addition & 1 deletion src/zarr/codecs/bytes.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@

import numpy as np

from zarr.abc.buffer import Buffer, NDArrayLike, NDBuffer
from zarr.abc.codec import ArrayBytesCodec
from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
from zarr.core.common import JSON, parse_enum, parse_named_configuration
from zarr.registry import register_codec

Expand Down
2 changes: 1 addition & 1 deletion src/zarr/codecs/crc32c_.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
if TYPE_CHECKING:
from typing import Self

from zarr.abc.buffer import Buffer
from zarr.core.array_spec import ArraySpec
from zarr.core.buffer import Buffer


@dataclass(frozen=True)
Expand Down
4 changes: 2 additions & 2 deletions src/zarr/codecs/gzip.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,15 @@
from numcodecs.gzip import GZip

from zarr.abc.codec import BytesBytesCodec
from zarr.core.buffer.cpu import as_numpy_array_wrapper
from zarr.buffer.cpu import as_numpy_array_wrapper
from zarr.core.common import JSON, parse_named_configuration
from zarr.registry import register_codec

if TYPE_CHECKING:
from typing import Self

from zarr.abc.buffer import Buffer
from zarr.core.array_spec import ArraySpec
from zarr.core.buffer import Buffer


def parse_gzip_level(data: JSON) -> int:
Expand Down
Loading

0 comments on commit 97ddaf8

Please sign in to comment.