Fix fill type issue for the Structured Types

Author: dmitriyrepin

src/mdio/constants.py

@@ -45,14 +45,14 @@
     ScalarType.FLOAT16: np_nan,
     ScalarType.FLOAT32: np_nan,
     ScalarType.FLOAT64: np_nan,
-    ScalarType.UINT8: 2**8 - 1,  # Max value for uint8
-    ScalarType.UINT16: 2**16 - 1,  # Max value for uint16
-    ScalarType.UINT32: 2**32 - 1,  # Max value for uint32
-    ScalarType.UINT64: 2**64 - 1,  # Max value for uint64
-    ScalarType.INT8: 2**7 - 1,  # Max value for int8
-    ScalarType.INT16: 2**15 - 1,  # Max value for int16
-    ScalarType.INT32: 2**31 - 1,  # Max value for int32
-    ScalarType.INT64: 2**63 - 1,  # Max value for int64
+    ScalarType.UINT8: UINT8_MAX,
+    ScalarType.UINT16: UINT16_MAX,
+    ScalarType.UINT32: UINT32_MAX,
+    ScalarType.UINT64: UINT64_MAX,
+    ScalarType.INT8: INT8_MAX,
+    ScalarType.INT16: INT16_MAX,
+    ScalarType.INT32: INT32_MAX,
+    ScalarType.INT64: INT64_MAX,
     ScalarType.COMPLEX64: complex(np_nan, np_nan),
     ScalarType.COMPLEX128: complex(np_nan, np_nan),
     ScalarType.COMPLEX256: complex(np_nan, np_nan),

src/mdio/schemas/v1/dataset_serializer.py

@@ -2,11 +2,12 @@
 
 from collections.abc import Mapping
 
-from dask import array as dask_array
+import numpy as np
 from numcodecs import Blosc as nc_Blosc
 from numpy import dtype as np_dtype
 from xarray import DataArray as xr_DataArray
 from xarray import Dataset as xr_Dataset
+from zarr import zeros as zarr_zeros
 from zarr.core.chunk_key_encodings import V2ChunkKeyEncoding
 
 try:
@@ -90,15 +91,14 @@ def _get_coord_names(var: Variable) -> list[str]:
     return coord_names
 
 
-def _get_np_datatype(var: Variable) -> np_dtype:
+def _get_np_datatype(data_type: ScalarType | StructuredType) -> np_dtype:
     """Get the numpy dtype for a variable."""
-    data_type = var.data_type
     if isinstance(data_type, ScalarType):
         return np_dtype(data_type.value)
     if isinstance(data_type, StructuredType):
         return np_dtype([(f.name, f.format.value) for f in data_type.fields])
-    err = f"Unsupported data type: {type(data_type)} in variable {var.name}"
-    raise TypeError(err)
+    msg = f"Expected ScalarType or StructuredType, got {type(data_type).__name__}"
+    raise ValueError(msg)
 
 
 def _get_zarr_shape(var: Variable, all_named_dims: dict[str, NamedDimension]) -> tuple[int, ...]:
@@ -170,8 +170,9 @@ def _get_fill_value(data_type: ScalarType | StructuredType | str) -> any:
     if isinstance(data_type, ScalarType):
         return fill_value_map.get(data_type)
     if isinstance(data_type, StructuredType):
-        return tuple(fill_value_map.get(field.format) for field in data_type.fields)
-    if isinstance(data_type, str): 
+        d_type = _get_np_datatype(data_type)
+        return np.zeros((), dtype=d_type)
+    if isinstance(data_type, str):
         return ""
     # If we do not have a fill value for this type, use None
     return None
@@ -186,6 +187,10 @@ def to_xarray_dataset(mdio_ds: Dataset) -> xr_DataArray:  # noqa: PLR0912
     Args:
         mdio_ds: The source MDIO dataset to construct from.
 
+    Notes:
+        - We can't use Dask (e.g., dask_array.zeros) because of the problems with
+          structured type support. We will uze zarr.zeros instead
+
     Returns:
         The constructed dataset with proper MDIO structure and metadata.
     """
@@ -197,15 +202,16 @@ def to_xarray_dataset(mdio_ds: Dataset) -> xr_DataArray:  # noqa: PLR0912
     # First pass: Build all variables
     data_arrays: dict[str, xr_DataArray] = {}
     for v in mdio_ds.variables:
-        # Use dask array instead of numpy array for lazy evaluation
         shape = _get_zarr_shape(v, all_named_dims=all_named_dims)
-        dtype = _get_np_datatype(v)
+        dtype = _get_np_datatype(v.data_type)
         chunks = _get_zarr_chunks(v, all_named_dims=all_named_dims)
-        arr = dask_array.zeros(shape, dtype=dtype, chunks=chunks)
 
+        # Use zarr.zeros to create an empty array with the specified shape and dtype
+        # NOTE: zarr_format=2 is essential, to_zarr() will fail if zarr_format=2 is used
+        data = zarr_zeros(shape=shape, dtype=dtype, zarr_format=2)
         # Create a DataArray for the variable. We will set coords in the second pass
         dim_names = _get_dimension_names(v)
-        data_array = xr_DataArray(arr, dims=dim_names)
+        data_array = xr_DataArray(data, dims=dim_names)
 
         # Add array attributes
         if v.metadata is not None:
@@ -222,13 +228,8 @@ def to_xarray_dataset(mdio_ds: Dataset) -> xr_DataArray:  # noqa: PLR0912
         # Create a custom chunk key encoding with "/" as separator
         chunk_key_encoding = V2ChunkKeyEncoding(separator="/").to_dict()
         encoding = {
-            # Is this a bug in Zarr? For datatype:  
-            #   dtype([('cdp-x', '<i4'), ('cdp-y', '<i4'), ('elevation', '<f2'), ('some_scalar', '<f2')])
-            # I specify fill_value as 
-            #   (2147483647, 2147483647, nan, nan)
-            # But the fill_value stored in .zmetadata as
-            #    "fill_value": null
-            "fill_value": _get_fill_value(v.data_type),
+            # NOTE: See Zarr documentation on use of fill_value and _FillValue in Zarr v2 vs v3
+            "_FillValue": _get_fill_value(v.data_type),
             "chunks": chunks,
             "chunk_key_encoding": chunk_key_encoding,
             "compressor": _convert_compressor(v.compressor),

tests/unit/v1/test_dataset_serializer.py

@@ -5,9 +5,13 @@
 """Tests the schema v1 dataset_serializer public API."""
 
 import pytest
+from dask import array as dask_array
+from numpy import array as np_array
 from numpy import dtype as np_dtype
-from numpy import nan as np_nan
 from numpy import isnan as np_isnan
+from numpy import zeros as np_zeros
+from xarray import DataArray as xr_DataArray
+from zarr import zeros as zarr_zeros
 
 from mdio.constants import fill_value_map
 from mdio.schemas.chunk_grid import RegularChunkGrid
@@ -158,9 +162,7 @@ def test__get_np_datatype() -> None:
     ]
 
     for scalar_type, expected_numpy_type in scalar_type_tests:
-        variable = Variable(name="test_var", dimensions=[], data_type=scalar_type)
-
-        result = _get_np_datatype(variable)
+        result = _get_np_datatype(scalar_type)
         expected = np_dtype(expected_numpy_type)
 
         assert result == expected
@@ -176,12 +178,7 @@ def test__get_np_datatype() -> None:
         StructuredField(name="valid", format=ScalarType.BOOL),
     ]
     structured_multi = StructuredType(fields=multi_fields)
-
-    variable_multi_struct = Variable(
-        name="multi_struct_var", dimensions=[], data_type=structured_multi
-    )
-
-    result_multi = _get_np_datatype(variable_multi_struct)
+    result_multi = _get_np_datatype(structured_multi)
     expected_multi = np_dtype(
         [("x", "float64"), ("y", "float64"), ("z", "float64"), ("id", "int32"), ("valid", "bool")]
     )
@@ -254,7 +251,8 @@ def test__get_fill_value() -> None:
         ScalarType.INT32,
     ]
     for scalar_type in scalar_types:
-        assert fill_value_map[scalar_type] == _get_fill_value(scalar_type)
+        fill_value = _get_fill_value(scalar_type)
+        assert fill_value_map[scalar_type] == fill_value
 
     scalar_types = [
         ScalarType.COMPLEX64,
@@ -265,16 +263,23 @@ def test__get_fill_value() -> None:
         val = _get_fill_value(scalar_type)
         assert isinstance(val, complex)
         assert np_isnan(val.real)
-        assert np_isnan(val.imag)   
+        assert np_isnan(val.imag)
 
     # Test 2: StructuredType
     f1 = StructuredField(name="cdp-x", format=ScalarType.INT32)
     f2 = StructuredField(name="cdp-y", format=ScalarType.INT32)
     f3 = StructuredField(name="elevation", format=ScalarType.FLOAT16)
     f4 = StructuredField(name="some_scalar", format=ScalarType.FLOAT16)
     structured_type = StructuredType(fields=[f1, f2, f3, f4])
-    result_structured = _get_fill_value(structured_type)
-    assert result_structured == (2147483647, 2147483647, np_nan, np_nan)
+
+    expected = np_array(
+        (0, 0, 0.0, 0.0),
+        dtype=np_dtype(
+            [("cdp-x", "<i4"), ("cdp-y", "<i4"), ("elevation", "<f2"), ("some_scalar", "<f2")]
+        ),
+    )
+    result = _get_fill_value(structured_type)
+    assert expected == result
 
     # Test 3: String type - should return empty string
     result_string = _get_fill_value("string_type")
@@ -382,5 +387,107 @@ def test_seismic_poststack_3d_acceptance_to_xarray_dataset(tmp_path) -> None:  #
 
     xr_ds = to_xarray_dataset(dataset)
 
-    file_path = output_path(tmp_path, f"{xr_ds.attrs['name']}", debugging=True)
+    file_path = output_path(tmp_path, f"{xr_ds.attrs['name']}", debugging=False)
     to_zarr(xr_ds, file_path, mode="w")
+
+
+@pytest.mark.skip(reason="Issues serializing dask arrays of structured types to dask.")
+def test_to_zarr_dask(tmp_path) -> None:  # noqa: ANN001
+    """Test writing XArray dataset with data as dask array to Zar."""
+    # Create a data type and the fill value
+    dtype = np_dtype([("inline", "int32"), ("cdp-x", "float64")])
+    dtype_fill_value = np_zeros((), dtype=dtype)
+
+    # Use '_FillValue' instead of 'fill_value'
+    # 'fill_value' is not a valid encoding key in Zarr v2
+    my_attr_encoding = {
+        "_FillValue": dtype_fill_value,
+        "chunk_key_encoding": {"name": "v2", "separator": "/"},
+    }
+
+    # Create a dask array using the data type
+    # Do not specify encoding as the array attribute
+    data = dask_array.zeros((36,), dtype=dtype, chunks=(36,))
+    aa = xr_DataArray(name="myattr", data=data)
+
+    # Specify encoding per array
+    encoding = {"myattr": my_attr_encoding}
+    file_path = output_path(tmp_path, "to_zarr/zarr_dask", debugging=False)
+    aa.to_zarr(file_path, mode="w", zarr_format=2, encoding=encoding, compute=False)
+
+
+def test_to_zarr_zarr_zerros_1(tmp_path) -> None:  # noqa: ANN001
+    """Test writing XArray dataset with data as Zarr zero array to Zar.
+
+    Set encoding in as DataArray attributes
+    """
+    # Create a data type and the fill value
+    dtype = np_dtype([("inline", "int32"), ("cdp-x", "float64")])
+    dtype_fill_value = np_zeros((), dtype=dtype)
+
+    # Use '_FillValue' instead of 'fill_value'
+    # 'fill_value' is not a valid encoding key in Zarr v2
+    my_attr_encoding = {
+        "_FillValue": dtype_fill_value,
+        "chunk_key_encoding": {"name": "v2", "separator": "/"},
+    }
+
+    # Create a zarr array using the data type,
+    # Specify encoding as the array attribute
+    data = zarr_zeros((36, 36), dtype=dtype, zarr_format=2)
+    aa = xr_DataArray(name="myattr", data=data)
+    aa.encoding = my_attr_encoding
+
+    file_path = output_path(tmp_path, "to_zarr/zarr_zarr_zerros_1", debugging=False)
+    aa.to_zarr(file_path, mode="w", zarr_format=2, compute=False)
+
+
+def test_to_zarr_zarr_zerros_2(tmp_path) -> None:  # noqa: ANN001
+    """Test writing XArray dataset with data as Zarr zero array to Zar.
+
+    Set encoding in the to_zar method
+    """
+    # Create a data type and the fill value
+    dtype = np_dtype([("inline", "int32"), ("cdp-x", "float64")])
+    dtype_fill_value = np_zeros((), dtype=dtype)
+
+    # Use '_FillValue' instead of 'fill_value'
+    # 'fill_value' is not a valid encoding key in Zarr v2
+    my_attr_encoding = {
+        "_FillValue": dtype_fill_value,
+        "chunk_key_encoding": {"name": "v2", "separator": "/"},
+    }
+
+    # Create a zarr array using the data type,
+    # Do not specify encoding as the array attribute
+    data = zarr_zeros((36, 36), dtype=dtype, zarr_format=2)
+    aa = xr_DataArray(name="myattr", data=data)
+
+    file_path = output_path(tmp_path, "to_zarr/zarr_zarr_zerros_2", debugging=False)
+    # Specify encoding per array
+    encoding = {"myattr": my_attr_encoding}
+    aa.to_zarr(file_path, mode="w", zarr_format=2, encoding=encoding, compute=False)
+
+
+def test_to_zarr_np(tmp_path) -> None:  # noqa: ANN001
+    """Test writing XArray dataset with data as NumPy array to Zar."""
+    # Create a data type and the fill value
+    dtype = np_dtype([("inline", "int32"), ("cdp-x", "float64")])
+    dtype_fill_value = np_zeros((), dtype=dtype)
+
+    # Use '_FillValue' instead of 'fill_value'
+    # 'fill_value' is not a valid encoding key in Zarr v2
+    my_attr_encoding = {
+        "_FillValue": dtype_fill_value,
+        "chunk_key_encoding": {"name": "v2", "separator": "/"},
+    }
+
+    # Create a zarr array using the data type
+    # Do not specify encoding as the array attribute
+    data = np_zeros((36, 36), dtype=dtype)
+    aa = xr_DataArray(name="myattr", data=data)
+
+    file_path = output_path(tmp_path, "to_zarr/zarr_np", debugging=False)
+    # Specify encoding per array
+    encoding = {"myattr": my_attr_encoding}
+    aa.to_zarr(file_path, mode="w", zarr_format=2, encoding=encoding, compute=False)