ENH: Add data type to sparse

CHANGELOG.md

@@ -3,6 +3,12 @@ Write the date in place of the "Unreleased" in the case a new version is release
 
 # Changelog
 
+## v0.1.0b9 (Unreleased)
+
+### Added
+
+- `dtype` property for sparse arrays in `COOAdapter` and `COOStructure`.
+
 ## v0.1.0b8 (2024-09-06)
 
 ### Fixed

tiled/_tests/test_protocols.py

@@ -247,9 +247,13 @@ def test_sparseadapter_protocol(mocker: MockFixture) -> None:
     mock_call4 = mocker.patch.object(CustomSparseAdapter, "specs")
     mock_call5 = mocker.patch.object(CustomSparseAdapter, "metadata")
 
-    structure = COOStructure(shape=(2 * 5,), chunks=((5, 5),))
-
     array = numpy.random.rand(2, 512, 512)
+
+    structure = COOStructure(
+        shape=(2 * 5,),
+        chunks=((5, 5),),
+        data_type=BuiltinDtype.from_numpy_dtype(array.dtype),
+    )
     blocks: Dict[Tuple[int, ...], Tuple[NDArray[Any], Any]] = {(1,): (array, (1,))}
     metadata: JSON = {"foo": "bar"}
     anyslice = (1, 1, 1)

tiled/_tests/test_writing.py

@@ -26,6 +26,7 @@
 from ..mimetypes import PARQUET_MIMETYPE
 from ..queries import Key
 from ..server.app import build_app
+from ..structures.array import BuiltinDtype
 from ..structures.core import Spec, StructureFamily
 from ..structures.data_source import DataSource
 from ..structures.sparse import COOStructure
@@ -245,7 +246,13 @@ def test_write_sparse_chunked(tree):
                 "sparse",
                 [
                     DataSource(
-                        structure=COOStructure(shape=(2 * N,), chunks=((N, N),)),
+                        structure=COOStructure(
+                            shape=(2 * N,),
+                            chunks=((N, N),),
+                            data_type=BuiltinDtype.from_numpy_dtype(
+                                numpy.dtype("float64")
+                            ),
+                        ),
                         structure_family="sparse",
                     )
                 ],

tiled/adapters/sparse.py

@@ -6,6 +6,7 @@
 import sparse
 from numpy._typing import NDArray
 
+from ..structures.array import BuiltinDtype
 from ..structures.core import Spec, StructureFamily
 from ..structures.sparse import COOStructure
 from .array import slice_and_shape_from_block_and_chunks
@@ -49,6 +50,7 @@ def from_arrays(
             dims=dims,
             shape=shape,
             chunks=tuple((dim,) for dim in shape),
+            data_type=BuiltinDtype.from_numpy_dtype(data.dtype),
             resizable=False,
         )
         return cls(
@@ -133,6 +135,7 @@ def from_global_ref(
             dims=dims,
             shape=shape,
             chunks=chunks,
+            data_type=BuiltinDtype.from_numpy_dtype(data.dtype),
             resizable=False,
         )
         return cls(

tiled/client/container.py

@@ -886,13 +886,15 @@ def write_sparse(
         >>> x.write_block(coords=[[2, 4]], data=[3.1, 2.8], block=(0,))
         >>> x.write_block(coords=[[0, 1]], data=[6.7, 1.2], block=(1,))
         """
+        from ..structures.array import BuiltinDtype
         from ..structures.sparse import COOStructure
 
         structure = COOStructure(
             shape=shape,
             # This method only supports single-chunk COO arrays.
             chunks=tuple((dim,) for dim in shape),
             dims=dims,
+            data_type=BuiltinDtype.from_numpy_dtype(data.dtype),
         )
         client = self.new(
             StructureFamily.sparse,

tiled/server/pydantic_sparse.py

@@ -2,12 +2,14 @@
 
 import pydantic
 
+from ..structures.array import BuiltinDtype, StructDtype
 from ..structures.sparse import SparseLayout
 
 
 class COOStructure(pydantic.BaseModel):
     shape: Tuple[int, ...]  # tuple of ints like (3, 3)
     chunks: Tuple[Tuple[int, ...], ...]  # tuple-of-tuples-of-ints like ((3,), (3,))
+    data_type: Optional[Union[BuiltinDtype, StructDtype]] = None
     dims: Optional[Tuple[str, ...]] = None  # None or tuple of names like ("x", "y")
     resizable: Union[bool, Tuple[bool, ...]] = False
     layout: SparseLayout = SparseLayout.COO

tiled/structures/sparse.py

@@ -2,6 +2,8 @@
 from dataclasses import dataclass
 from typing import Optional, Tuple, Union
 
+from .array import BuiltinDtype, StructDtype
+
 
 class SparseLayout(str, enum.Enum):
     # Only COO is currently supported, but this lays a path
@@ -13,14 +15,21 @@ class SparseLayout(str, enum.Enum):
 class COOStructure:
     chunks: Tuple[Tuple[int, ...], ...]  # tuple-of-tuples-of-ints like ((3,), (3,))
     shape: Tuple[int, ...]  # tuple of ints like (3, 3)
+    data_type: Optional[Union[BuiltinDtype, StructDtype]] = None
     dims: Optional[Tuple[str, ...]] = None  # None or tuple of names like ("x", "y")
     resizable: Union[bool, Tuple[bool, ...]] = False
     layout: SparseLayout = SparseLayout.COO
     # TODO Include fill_value?
 
     @classmethod
     def from_json(cls, structure):
+        data_type = structure.get("data_type", None)
+        if data_type is not None and "fields" in data_type:
+            data_type = StructDtype.from_json(data_type)
+        else:
+            data_type = BuiltinDtype.from_json(data_type)
         return cls(
+            data_type=data_type,
             chunks=tuple(map(tuple, structure["chunks"])),
             shape=tuple(structure["shape"]),
             dims=structure["dims"],