Handle timedelta serde

Author: moradology

src/zarr/core/metadata/v2.py

@@ -310,6 +310,13 @@ def _parse_structured_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
         raise ValueError(f"Fill_value {fill_value} is not valid for dtype {dtype}.") from e
 
 
+def _parse_timedelta(td: dict[str, Any]) -> np.timedelta64:
+    if td["value"] is None:
+        return np.timedelta64("NaT")
+    else:
+        return np.timedelta64(int(td["value"]), td["unit"])
+
+
 def parse_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
     """
     Parse a potential fill value into a value that is compatible with the provided dtype.
@@ -329,13 +336,18 @@ def parse_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
     if fill_value is None or dtype.hasobject:
         # Pass through None or if dtype is object
         pass
-    elif dtype.kind in "M":
+    elif dtype.kind == "M":  # datetime
         # Check for both string "NaT" and the int64 representation of NaT
         if fill_value == "NaT" or fill_value == np.iinfo(np.int64).min:
             fill_value = dtype.type("NaT")
         else:
             fill_value = np.array(fill_value, dtype=dtype)[()]
         # Fall through for non-NaT datetime/timedelta values (handled below)
+    elif dtype.kind == "m":  # timedelta
+        if isinstance(fill_value, dict):
+            return _parse_timedelta(fill_value)
+        else:  # if raw value is passed rather than unit-based serialization
+            return np.timedelta64(fill_value)
     elif dtype.fields is not None:
         # the dtype is structured (has multiple fields), so the fill_value might be a
         # compound value (e.g., a tuple or dict) that needs field-wise processing.
@@ -373,6 +385,14 @@ def parse_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> Any:
     return fill_value
 
 
+def _serialize_timedelta(td: np.timedelta64) -> JSON:
+    if np.isnat(td):
+        return {"value": None, "unit": None}
+    else:
+        val, unit = int(td.astype(int)), td.dtype.name.split("[")[-1][:-1]
+        return {"value": val, "unit": unit}
+
+
 def _serialize_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> JSON:
     serialized: JSON
 
@@ -383,7 +403,9 @@ def _serialize_fill_value(fill_value: Any, dtype: np.dtype[Any]) -> JSON:
         # that mypy isn't aware of. The fact that we have S or V dtype here
         # means we should have a bytes-type fill_value.
         serialized = base64.standard_b64encode(cast(bytes, fill_value)).decode("ascii")
-    elif isinstance(fill_value, np.datetime64):
+    elif dtype.kind == "m":
+        serialized = _serialize_timedelta(fill_value)
+    elif dtype.kind == "M":
         serialized = np.datetime_as_string(fill_value)
     elif isinstance(fill_value, numbers.Integral):
         serialized = int(fill_value)
@@ -423,7 +445,9 @@ def _default_fill_value(dtype: np.dtype[Any]) -> Any:
         return b""
     elif dtype.kind in "UO":
         return ""
-    elif dtype.kind in "Mm":
+    elif dtype.kind in "m":  # timedelta64
+        return np.timedelta64("NaT")
+    elif dtype.kind in "M":  # datetime64
         return dtype.type("nat")
     elif dtype.kind == "V":
         if dtype.fields is not None:

src/zarr/testing/strategies.py

@@ -64,7 +64,7 @@ def v2_dtypes() -> st.SearchStrategy[np.dtype]:
         | npst.byte_string_dtypes(endianness="=")
         | npst.unicode_string_dtypes(endianness="=")
         | npst.datetime64_dtypes(endianness="=")
-        # | npst.timedelta64_dtypes()
+        | npst.timedelta64_dtypes(endianness="=")
     )
 
 

tests/test_metadata/test_v2.py

@@ -340,6 +340,29 @@ def test_parse_v2_fill_value_non_nat(dtype_str: str) -> None:
     assert meta.fill_value.dtype == dtype
 
 
+@pytest.mark.parametrize("dtype_str", ["timedelta64[ms]"])
+@pytest.mark.parametrize("fill_value_raw", [12345, np.timedelta64(12345, "ms")])  # Example values
+def test_parse_v2_fill_value_non_nat_raw_timedelta(dtype_str: str, fill_value_raw: Any) -> None:
+    dtype = np.dtype(dtype_str)
+    expected_value = (
+        np.timedelta64(fill_value_raw, "ms") if isinstance(fill_value_raw, int) else fill_value_raw
+    )
+
+    metadata_dict = {
+        "zarr_format": 2,
+        "shape": (10,),
+        "chunks": (5,),
+        "dtype": dtype.str,
+        "compressor": None,
+        "filters": None,
+        "fill_value": fill_value_raw,
+        "order": "C",
+    }
+    meta = ArrayV2Metadata.from_dict(metadata_dict)
+    assert isinstance(meta.fill_value, np.timedelta64)
+    assert meta.fill_value == expected_value
+
+
 def test_from_dict_extra_fields() -> None:
     data = {
         "_nczarr_array": {"dimrefs": ["/dim1", "/dim2"], "storage": "chunked"},