Add support for Zarr-Python v3 and Zarr format v3

pyproject.toml

@@ -23,7 +23,7 @@ zarr = [
     's3fs',
     'scipy',
     'xarray',
-    'zarr>=2,<3',
+    'zarr>=3,<4',
     'pyarrow',
     'psutil' # psutil is needed so large FITS images are not loaded into memory
 ]
@@ -38,7 +38,7 @@ test = [
     's3fs',
     'scipy',
     'xarray',
-    'zarr>=2,<3',
+    'zarr>=3,<4',
     'pyarrow',
     'psutil',
     'pytest',
@@ -53,7 +53,7 @@ casacore = [
     's3fs',
     'scipy',
     'xarray',
-    'zarr>=2,<3',
+    'zarr>=3,<4',
     'pyarrow',
     'psutil',
     'python_casacore>=3.6.1; sys_platform != "darwin"'
@@ -65,7 +65,7 @@ interactive = [
     's3fs',
     'scipy',
     'xarray',
-    'zarr>=2,<3',
+    'zarr>=3,<4',
     'pyarrow',
     'psutil',
     'matplotlib',
@@ -95,7 +95,7 @@ all = [
     's3fs',
     'scipy',
     'xarray',
-    'zarr>=2,<3',
+    'zarr>=3,<4',
     'pyarrow',
     'psutil',
     'pytest',

src/xradio/_utils/_casacore/casacore_from_casatools.py

@@ -781,7 +781,7 @@ def imageinfo(self) -> dict:
             if "imageinfo" in tb.keywordnames():
                 image_metadata = tb.getkeyword("imageinfo")
             else:
-                image_metadata = {"imagetype": "Intensity", "objectname": ""}
+                image_metadata = {"imagetype": "sky", "objectname": ""}
 
         image_metadata["imagetype"] = _validate_image_type(
             image_metadata.get("imagetype", "Intensity")

src/xradio/_utils/zarr/config.py

@@ -0,0 +1,4 @@
+# Added for the transition from zarr-python 2.x to 3.x. Zarr-Python v3 still supports
+# Zarr format v2 in addition to the new v3. This config option is set for the new/default
+# format v3 but can be changed to use format v2.
+ZARR_FORMAT = 3

src/xradio/image/_util/_zarr/xds_from_zarr.py

@@ -93,8 +93,8 @@ def _decode_sub_xdses(xarrayObj, top_dir: str, id_dict: dict) -> None:
                 ky = d.name[len(_top_level_sub_xds) :]
                 xarrayObj.attrs[ky] = _read_zarr(path, id_dict)
                 # TODO if attrs that are xdses have attrs that are xdses ...
-            else:
-                # descend into the directory
+            elif d.name != "c":
+                # descend into the directory (except for the Zarr v3 "c" dirs)
                 _decode_sub_xdses(xarrayObj[d.name], path, id_dict)
 
 

src/xradio/image/_util/_zarr/xds_to_zarr.py

@@ -5,6 +5,8 @@
 import os
 from .common import _np_types, _top_level_sub_xds
 
+from xradio._utils.zarr.config import ZARR_FORMAT
+
 
 def _write_zarr(xds: xr.Dataset, zarr_store: str):
     max_chunk_size = 0.95 * 2**30
@@ -25,7 +27,7 @@ def _write_zarr(xds: xr.Dataset, zarr_store: str):
                 )
     xds_copy = xds.copy(deep=True)
     sub_xds_dict = _encode(xds_copy, zarr_store)
-    z_obj = xds_copy.to_zarr(store=zarr_store, compute=True)
+    z_obj = xds_copy.to_zarr(store=zarr_store, compute=True, zarr_format=ZARR_FORMAT)
     if sub_xds_dict:
         _write_sub_xdses(sub_xds_dict)
 
@@ -62,4 +64,4 @@ def _encode_dict(my_dict: dict, top_path: str, sub_xds_dict) -> tuple:
 
 def _write_sub_xdses(sub_xds: dict):
     for k, v in sub_xds.items():
-        z_obj = v.to_zarr(store=k, compute=True)
+        z_obj = v.to_zarr(store=k, compute=True, zarr_format=ZARR_FORMAT)

src/xradio/image/_util/image_factory.py

@@ -627,7 +627,7 @@ def detect_store_type(store):
         elif os.path.isdir(store):
             if "table.info" in os.listdir(store):
                 store_type = "casa"
-            elif ".zattrs" in os.listdir(store):
+            elif ".zattrs" in os.listdir(store) or "zarr.json" in os.listdir(store):
                 store_type = "zarr"
             else:
                 xradio_logger().error("Unknown directory structure.")

src/xradio/measurement_set/_utils/_msv2/_tables/read.py

@@ -123,6 +123,34 @@ def convert_casacore_time_to_mjd(rawtimes: np.ndarray) -> np.ndarray:
     return rawtimes / SECS_IN_DAY
 
 
+def casacore_numpy_to_json_safe_type(value: object) -> object:
+    """
+    This funciton converts values with numpy types (commonly found in "raw" xds loaded with load_generic_table()) to
+    native types that can be safely written to JSON.
+
+    This is handy when loading values from table columns that one wants to put in attributes. Starting
+    with Zarr 3, numpy types are not converted/encoded before writing to JSON, and an exception is raised.
+
+    Parameters
+    ----------
+    value : object
+        A value or numpy array of values, of presumably numpy type
+
+    Returns
+    -------
+    object
+        The same value converted to a JSON safe type (for example int(a_numpy_int32_value))
+    """
+    if isinstance(value, np.ndarray):
+        return ",".join([scalar for scalar in value])
+    elif isinstance(value, np.integer):
+        return int(value)
+    elif isinstance(value, np.floating):
+        return float(value)
+    else:
+        return value
+
+
 def make_taql_where_between_min_max(
     min_max: Tuple[np.float64, np.float64],
     path: str,

src/xradio/measurement_set/_utils/_msv2/conversion.py

@@ -1,16 +1,16 @@
 from collections import deque
 import datetime
 import importlib
-import numcodecs
 import os
 import pathlib
 import time
+import traceback
 from typing import Callable, Dict, Union
 
 import dask.array as da
 import numpy as np
 import xarray as xr
-import traceback
+import zarr.codecs
 
 try:
     from casacore import tables
@@ -1012,7 +1012,7 @@ def convert_and_write_partition(
     ephemeris_interpolate: bool = False,
     phase_cal_interpolate: bool = False,
     sys_cal_interpolate: bool = False,
-    compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
+    compressor: zarr.abc.codec.BytesBytesCodec = zarr.codecs.ZstdCodec(level=2),
     add_reshaping_indices: bool = False,
     storage_backend="zarr",
     parallel_mode: str = "none",
@@ -1412,8 +1412,12 @@ def get_observation_info(in_file, observation_id, scan_intents):
             ms_xdt["/phased_array_xds"] = phased_array_xds
 
         if storage_backend == "zarr":
+            from xradio._utils.zarr.config import ZARR_FORMAT
+
             ms_xdt.to_zarr(
-                store=os.path.join(out_file, ms_v4_name), mode=persistence_mode
+                store=os.path.join(out_file, ms_v4_name),
+                mode=persistence_mode,
+                zarr_format=ZARR_FORMAT,
             )
         elif storage_backend == "netcdf":
             # xds.to_netcdf(path=file_name+"/MAIN", mode=mode) #Does not work

src/xradio/measurement_set/_utils/_msv2/msv4_info_dicts.py

@@ -1,6 +1,7 @@
 import re
 
 import numpy as np
+
 import xarray as xr
 
 try:
@@ -10,7 +11,11 @@
 
 
 from .subtables import subt_rename_ids
-from ._tables.read import load_generic_table, convert_casacore_time
+from ._tables.read import (
+    casacore_numpy_to_json_safe_type,
+    convert_casacore_time,
+    load_generic_table,
+)
 from xradio._utils.list_and_array import check_if_consistent
 from xradio._utils.logging import xradio_logger
 
@@ -191,7 +196,7 @@ def extract_optional_fields_asdm_asis_table(
 
     Returns:
     --------
-    info: dict
+    table_info: dict
         info dict with description from an ASDM_* subtable, ready
         for the MSv4 observation_info dict
     """
@@ -200,10 +205,7 @@ def extract_optional_fields_asdm_asis_table(
     for field_msv4, col_msv2 in optional_fields.items():
         if col_msv2 in asdm_asis_xds.data_vars:
             msv2_value = asdm_asis_xds[col_msv2].values[0]
-            if isinstance(msv2_value, np.ndarray):
-                table_info[field_msv4] = ",".join([log for log in msv2_value])
-            else:
-                table_info[field_msv4] = msv2_value
+            table_info[field_msv4] = casacore_numpy_to_json_safe_type(msv2_value)
 
     return table_info
 

src/xradio/measurement_set/_utils/_zarr/encoding.py

@@ -6,4 +6,4 @@ def add_encoding(xds, compressor, chunks=None):
 
     for da_name in list(xds.data_vars):
         da_chunks = [chunks[dim_name] for dim_name in xds[da_name].sizes]
-        xds[da_name].encoding = {"compressor": compressor, "chunks": da_chunks}
+        xds[da_name].encoding = {"compressors": (compressor,), "chunks": da_chunks}

src/xradio/measurement_set/convert_msv2_to_processing_set.py

@@ -1,8 +1,8 @@
-import numcodecs
 from typing import Dict, Union, Literal
 import time
 
 import dask
+import zarr.codecs
 
 from xradio.measurement_set._utils._msv2.partition_queries import (
     create_partitions,
@@ -11,6 +11,7 @@
     convert_and_write_partition,
     estimate_memory_and_cores_for_partitions,
 )
+from xradio._utils.zarr.config import ZARR_FORMAT
 from xradio._utils.logging import xradio_logger
 
 
@@ -64,7 +65,7 @@ def convert_msv2_to_processing_set(
     phase_cal_interpolate: bool = False,
     sys_cal_interpolate: bool = False,
     use_table_iter: bool = False,
-    compressor: numcodecs.abc.Codec = numcodecs.Zstd(level=2),
+    compressor: zarr.abc.codec.BytesBytesCodec = zarr.codecs.ZstdCodec(level=2),
     add_reshaping_indices: bool = False,
     storage_backend: Literal["zarr", "netcdf"] = "zarr",
     parallel_mode: Literal["none", "partition", "time"] = "none",
@@ -125,7 +126,7 @@ def convert_msv2_to_processing_set(
     if not str(out_file).endswith("ps.zarr"):
         out_file += ".ps.zarr"
 
-    ps_dt.to_zarr(store=out_file, mode=persistence_mode)
+    ps_dt.to_zarr(store=out_file, mode=persistence_mode, zarr_format=ZARR_FORMAT)
 
     # Check `parallel_mode` is valid
     try:
@@ -226,6 +227,7 @@ def convert_msv2_to_processing_set(
 
     import zarr
 
-    root_group = zarr.open(out_file, mode="r+")  # Open in read/write mode
+    # Open in read/write mode
+    root_group = zarr.open(out_file, mode="r+", zarr_format=ZARR_FORMAT)
     root_group.attrs["type"] = "processing_set"  # Replace
-    zarr.convenience.consolidate_metadata(root_group.store)
+    zarr.consolidate_metadata(root_group.store)

src/xradio/measurement_set/processing_set_xdt.py

@@ -213,7 +213,7 @@ def _summary(self, data_group_name: str = None):
         from astropy.coordinates import SkyCoord
         import astropy.units as u
 
-        for key, value in self._xdt.items():
+        for key, value in sorted(self._xdt.items()):
             partition_info = value.xr_ms.get_partition_info()
             observation_info = value.observation_info
 

tests/stakeholder/test_measure_set_stakeholder.py

@@ -24,6 +24,7 @@
     ProcessingSetXdt,
 )
 from xradio.schema.check import check_datatree
+from xradio._utils.zarr.config import ZARR_FORMAT
 
 # relative_tolerance = 10 ** (-12)
 relative_tolerance = 10 ** (-5)
@@ -364,7 +365,8 @@ def base_test(
         else:
             ps_copy_name = str(ps_name) + "_copy"
 
-        ps_lazy_xdt.to_zarr(ps_copy_name)  # Test writing to disk.
+        # Test writing to disk.
+        ps_lazy_xdt.to_zarr(ps_copy_name, zarr_format=ZARR_FORMAT)
 
         ms_xdt_name = list(ps_lazy_xdt.keys())[0]
         ms_xds = ps_lazy_xdt[ms_xdt_name].ds

tests/unit/measurement_set/_utils/_msv2/test_conversion.py

@@ -426,7 +426,7 @@ def test_convert_and_write_partition_with_antenna1(ms_minimal_required):
         )
 
         # Will need a SD-like test ms. Otherwise the partition is empty:
-        with pytest.raises(FileNotFoundError, match="No such file or directory"):
+        with pytest.raises(FileNotFoundError, match="does not exist"):
             msv4_xds = xr.open_dataset(
                 out_name
                 + "/"

tests/unit/measurement_set/_utils/_zarr/test_encoding.py

@@ -20,7 +20,7 @@ def test_add_encoding_wo_chunks():
 
     add_encoding(xds, single_encoding)
     assert xds
-    assert xds.da.encoding == {"compressor": single_encoding, "chunks": [3]}
+    assert xds.da.encoding == {"compressors": (single_encoding,), "chunks": [3]}
 
 
 def test_add_encoding_with_wrong_chunks():
@@ -38,7 +38,7 @@ def test_add_encoding_with_wrong_chunks():
     chunks_size = 1
     add_encoding(xds, single_encoding, chunks={"zz_not_there": chunks_size})
     assert xds
-    assert xds.da.encoding == {"compressor": single_encoding, "chunks": [3]}
+    assert xds.da.encoding == {"compressors": (single_encoding,), "chunks": [3]}
 
 
 def test_add_encoding_with_chunks():
@@ -56,4 +56,7 @@ def test_add_encoding_with_chunks():
     chunks_size = 1
     add_encoding(xds, single_encoding, chunks={"x": chunks_size})
     assert xds
-    assert xds.da.encoding == {"compressor": single_encoding, "chunks": [chunks_size]}
+    assert xds.da.encoding == {
+        "compressors": (single_encoding,),
+        "chunks": [chunks_size],
+    }