Write _ARRAY_DIMENSIONS array attribute

Author: sjperkins

daskms/experimental/zarr/__init__.py

@@ -10,6 +10,8 @@
 import numpy as np
 import warnings
 
+ARRAY_DIMENSION = "_ARRAY_DIMENSIONS"
+
 from daskms.constants import DASKMS_PARTITION_KEY
 from daskms.dataset import Dataset, Variable
 from daskms.dataset_schema import DatasetSchema, encode_type, decode_type, decode_attr
@@ -111,6 +113,8 @@ def create_array(ds_group, column, column_schema, schema_chunks, coordinate=Fals
         exact=True,
     )
 
+    array.attrs[ARRAY_DIMENSION] = column_schema.dims
+
     array.attrs[DASKMS_ATTR_KEY] = {
         **column_schema.attrs,
         "dims": column_schema.dims,

daskms/experimental/zarr/tests/test_zarr.py

@@ -235,6 +235,7 @@ def test_multiprocess_create(ms, tmp_path_factory):
 @pytest.mark.skipif(xarray is None, reason="depends on xarray")
 def test_xarray_to_zarr(ms, tmp_path_factory):
     store = tmp_path_factory.mktemp("zarr_store")
+
     datasets = xds_from_ms(ms)
 
     for i, ds in enumerate(datasets):
@@ -407,3 +408,26 @@ def test_xds_from_zarr_assert_on_empty_store(tmp_path_factory, ms):
 
     with pytest.raises(UnknownStoreTypeError, match="Unable to infer table type"):
         xds_from_zarr(path)
+
+
+@pytest.mark.skipif(xarray is None, reason="depends on xarray")
+def test_xarray_reading_daskms_written_dataset(ms, tmp_path_factory):
+    store = tmp_path_factory.mktemp("zarr_store")
+
+    datasets = xds_from_ms(ms)
+
+    for i, ds in enumerate(datasets):
+        chunks = ds.chunks
+        row = sum(chunks["row"])
+        chan = sum(chunks["chan"])
+        corr = sum(chunks["corr"])
+
+        datasets[i] = ds.assign_coords(
+            row=np.arange(row), chan=np.arange(chan), corr=np.arange(corr)
+        )
+
+    path = store / "test.zarr"
+    dask.compute(xds_to_zarr(datasets, path, consolidated=True))
+
+    ds = xarray.open_zarr(path / "MAIN" / "MAIN_0")
+    assert ds == datasets[0]