Rebase Ryan's work in zarr-developers#2031

Author: tomwhite

.pre-commit-config.yaml

@@ -36,6 +36,7 @@ repos:
           - typing_extensions
           - universal-pathlib
           - obstore>=0.5.1
+          - pyarrow
           # Tests
           - pytest
           - hypothesis

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     'numcodecs[crc32c]>=0.14',
     'typing_extensions>=4.9',
     'donfig>=0.8',
+    'pyarrow',
 ]
 
 dynamic = [

src/zarr/codecs/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from zarr.codecs.arrow import ArrowRecordBatchCodec
 from zarr.codecs.blosc import BloscCname, BloscCodec, BloscShuffle
 from zarr.codecs.bytes import BytesCodec, Endian
 from zarr.codecs.crc32c_ import Crc32cCodec
@@ -34,6 +35,7 @@
 from zarr.registry import register_codec
 
 __all__ = [
+    "ArrowRecordBatchCodec",
     "BloscCname",
     "BloscCodec",
     "BloscShuffle",
@@ -49,6 +51,7 @@
     "ZstdCodec",
 ]
 
+register_codec("arrow", ArrowRecordBatchCodec)
 register_codec("blosc", BloscCodec)
 register_codec("bytes", BytesCodec)
 

src/zarr/codecs/arrow.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+
+from zarr.abc.codec import ArrayBytesCodec
+from zarr.core.array_spec import ArraySpec
+from zarr.core.buffer import Buffer, NDBuffer
+from zarr.core.common import JSON, parse_named_configuration
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+CHUNK_FIELD_NAME = "zarr_chunk"
+
+
+@dataclass(frozen=True)
+class ArrowRecordBatchCodec(ArrayBytesCodec):
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def from_dict(cls, data: dict[str, JSON]) -> Self:
+        _, configuration_parsed = parse_named_configuration(
+            data, "arrow", require_configuration=False
+        )
+        configuration_parsed = configuration_parsed or {}
+        return cls(**configuration_parsed)
+
+    def to_dict(self) -> dict[str, JSON]:
+        return {"name": "arrow"}
+
+    def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
+        return self
+
+    async def _decode_single(
+        self,
+        chunk_bytes: Buffer,
+        chunk_spec: ArraySpec,
+    ) -> NDBuffer:
+        assert isinstance(chunk_bytes, Buffer)
+
+        # TODO: make this compatible with buffer prototype
+        arrow_buffer = pa.py_buffer(chunk_bytes.to_bytes())
+        with pa.ipc.open_stream(arrow_buffer) as reader:
+            batches = [b for b in reader]
+        assert len(batches) == 1
+        arrow_array = batches[0][CHUNK_FIELD_NAME]
+        chunk_array = chunk_spec.prototype.nd_buffer.from_ndarray_like(
+            arrow_array.to_numpy(zero_copy_only=False)
+        )
+
+        # ensure correct chunk shape
+        if chunk_array.shape != chunk_spec.shape:
+            chunk_array = chunk_array.reshape(
+                chunk_spec.shape,
+            )
+        return chunk_array
+
+    async def _encode_single(
+        self,
+        chunk_array: NDBuffer,
+        chunk_spec: ArraySpec,
+    ) -> Buffer | None:
+        assert isinstance(chunk_array, NDBuffer)
+        arrow_array = pa.array(chunk_array.as_ndarray_like().ravel())
+        rb = pa.record_batch([arrow_array], names=[CHUNK_FIELD_NAME])
+        # TODO: allocate buffer differently
+        sink = pa.BufferOutputStream()
+        with pa.ipc.new_stream(sink, rb.schema) as writer:
+            writer.write_batch(rb)
+        return chunk_spec.prototype.buffer.from_bytes(memoryview(sink.getvalue()))
+
+    def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int:
+        raise ValueError("Don't know how to compute encoded size!")

tests/test_codecs/test_arrow.py

@@ -0,0 +1,57 @@
+import numpy as np
+import pytest
+
+import zarr
+from zarr.abc.store import Store
+from zarr.codecs import ArrowRecordBatchCodec
+from zarr.storage import StorePath
+
+
+@pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        "uint8",
+        "uint16",
+        "uint32",
+        "uint64",
+        "int8",
+        "int16",
+        "int32",
+        "int64",
+        "float32",
+        "float64",
+    ],
+)
+def test_arrow_standard_dtypes(store: Store, dtype) -> None:
+    data = np.arange(0, 256, dtype=dtype).reshape((16, 16))
+
+    a = zarr.create_array(
+        StorePath(store, path="arrow"),
+        shape=data.shape,
+        chunks=(16, 16),
+        dtype=data.dtype,
+        fill_value=0,
+        serializer=ArrowRecordBatchCodec(),
+    )
+
+    a[:, :] = data
+    assert np.array_equal(data, a[:, :])
+
+
+@pytest.mark.parametrize("store", ("local", "memory"), indirect=["store"])
+def test_arrow_vlen_string(store: Store) -> None:
+    strings = ["hello", "world", "this", "is", "a", "test"]
+    data = np.array(strings).reshape((2, 3))
+
+    a = zarr.create_array(
+        StorePath(store, path="arrow"),
+        shape=data.shape,
+        chunks=data.shape,
+        dtype=data.dtype,
+        fill_value=0,
+        serializer=ArrowRecordBatchCodec(),
+    )
+
+    a[:, :] = data
+    assert np.array_equal(data, a[:, :])