Zarr data types refactor compat round 2

conftest.py

@@ -358,11 +358,12 @@ def manifest_array(array_v3_metadata):
     def _manifest_array(
         shape: tuple = (5, 2),
         chunks: tuple = (5, 2),
+        data_type: np.dtype = np.dtype("int32"),
         codecs: list[dict] | None = [ARRAYBYTES_CODEC, ZLIB_CODEC],
         dimension_names: Iterable[str] | None = None,
     ):
         metadata = array_v3_metadata(
-            shape=shape, chunks=chunks, codecs=codecs, dimension_names=dimension_names
+            shape=shape, chunks=chunks, data_type=data_type, codecs=codecs, dimension_names=dimension_names
         )
         entries = _generate_chunk_entries(shape, chunks, _entry_from_chunk_key)
         chunkmanifest = ChunkManifest(entries=entries)

virtualizarr/manifests/array.py

@@ -84,11 +84,13 @@ def chunks(self) -> tuple[int, ...]:
         Individual chunk size by number of elements.
         """
         return self._metadata.chunks
-
+    
     @property
     def dtype(self) -> np.dtype:
-        dtype_str = self.metadata.data_type
-        return dtype_str.to_numpy()
+        """The native dtype of the data (typically a numpy dtype)"""
+        zdtype = self.metadata.data_type
+        dtype = zdtype.to_native_dtype()
+        return dtype.str
 
     @property
     def shape(self) -> tuple[int, ...]:

virtualizarr/manifests/array_api.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Callable, cast
+from typing import TYPE_CHECKING, Any, Callable, cast, Union
 
 import numpy as np
 
@@ -34,12 +34,17 @@ def decorator(func):
 
 
 @implements(np.result_type)
-def result_type(*arrays_and_dtypes) -> np.dtype:
+def result_type(*arrays_and_dtypes: Union["ManifestArray", np.dtype]) -> np.dtype:
     """Called by xarray to ensure all arguments to concat have the same dtype."""
-    first_dtype, *other_dtypes = (np.dtype(obj) for obj in arrays_and_dtypes)
+    from virtualizarr.manifests.array import ManifestArray
+
+    dtypes = (obj.dtype if isinstance(obj, ManifestArray) else np.dtype(obj) for obj in arrays_and_dtypes)
+    first_dtype, *other_dtypes = dtypes
+    unique_dtypes = set(dtypes)
     for other_dtype in other_dtypes:
         if other_dtype != first_dtype:
-            raise ValueError("dtypes not all consistent")
+            raise ValueError(f"Cannot combine arrays with inconsistent dtypes, but got {len(unique_dtypes)} distinct dtypes: {unique_dtypes}")
+
     return first_dtype
 
 

virtualizarr/manifests/utils.py

@@ -8,6 +8,7 @@
     parse_dimension_names,
     parse_shapelike,
 )
+from zarr.dtype import parse_data_type
 
 from virtualizarr.codecs import convert_to_codec_pipeline, get_codecs
 
@@ -75,15 +76,16 @@ def create_v3_array_metadata(
     ArrayV3Metadata
         A configured ArrayV3Metadata instance with standard defaults
     """
+    zdtype = parse_data_type(data_type, zarr_format=3)
     return ArrayV3Metadata(
         shape=shape,
-        data_type=data_type.name if hasattr(data_type, "name") else data_type,
+        data_type=zdtype,
         chunk_grid={
             "name": "regular",
             "configuration": {"chunk_shape": chunk_shape},
         },
         chunk_key_encoding=chunk_key_encoding,
-        fill_value=fill_value,
+        fill_value=zdtype.default_scalar() if fill_value is None else fill_value,
         codecs=convert_to_codec_pipeline(
             codecs=codecs or [],
             dtype=data_type,

virtualizarr/tests/test_manifests/test_array.py

@@ -9,7 +9,7 @@
 from virtualizarr.manifests import ChunkManifest, ManifestArray
 
 
-class TestManifestArray:
+class TestInit:
     def test_manifest_array(self, array_v3_metadata):
         chunks_dict = {
             "0.0.0": {"path": "s3://bucket/foo.nc", "offset": 100, "length": 100},
@@ -50,6 +50,23 @@ def test_manifest_array_dict_v3_metadata(self, array_v3_metadata):
         assert marr.ndim == 3
 
 
+class TestResultType:
+    def test_idempotent(self, manifest_array):
+        marr1 = manifest_array(shape=(), chunks=(), data_type=np.dtype("int32"))
+        marr2 = manifest_array(shape=(), chunks=(), data_type=np.dtype("int32"))
+
+        assert np.result_type(marr1) == marr1.dtype
+        assert np.result_type(marr1, marr1.dtype) == marr1.dtype
+        assert np.result_type(marr1, marr2) == marr1.dtype
+
+    def test_raises(self, manifest_array):
+        marr1 = manifest_array(shape=(), chunks=(), data_type=np.dtype("int32"))
+        marr2 = manifest_array(shape=(), chunks=(), data_type=np.dtype("int64"))
+
+        with pytest.raises(ValueError, match="inconsistent"):
+            np.result_type(marr1, marr2)
+
+
 class TestEquals:
     def test_equals(self, array_v3_metadata):
         chunks_dict = {

virtualizarr/tests/test_manifests/test_group.py

@@ -27,7 +27,7 @@ def test_manifest_repr(self, manifest_array):
         expected_repr = textwrap.dedent(
             """
             ManifestGroup(
-                arrays={'foo': ManifestArray<shape=(5, 2), dtype=int32, chunks=(5, 2)>},
+                arrays={'foo': ManifestArray<shape=(5, 2), dtype=<i4, chunks=(5, 2)>},
                 groups={},
                 metadata=GroupMetadata(attributes={}, zarr_format=3, consolidated_metadata=None, node_type='group'),
             )

virtualizarr/utils.py

@@ -150,7 +150,7 @@ def convert_v3_to_v2_metadata(
 
     v2_metadata = ArrayV2Metadata(
         shape=v3_metadata.shape,
-        dtype=v3_metadata.data_type.to_numpy(),
+        dtype=v3_metadata.data_type,
         chunks=v3_metadata.chunks,
         fill_value=fill_value or v3_metadata.fill_value,
         compressor=compressor_config,

virtualizarr/writers/icechunk.py

@@ -344,7 +344,7 @@ def write_virtual_variable_to_icechunk(
             name=name,
             shape=metadata.shape,
             chunks=metadata.chunks,
-            dtype=metadata.data_type.to_numpy(),
+            dtype=metadata.data_type,
             filters=filters,
             compressors=compressors,
             dimension_names=var.dims,