Merge pull request #118 from EOPF-Explorer/fix-nan

Fix JSON compliance for NaN values in zarr attributes

Author: d-v-b

src/eopf_geozarr/conversion/fs_utils.py

@@ -3,7 +3,7 @@
 import json
 import os
 from collections.abc import Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse
 
 import s3fs
@@ -13,6 +13,92 @@
 
 from eopf_geozarr.types import S3Credentials, S3FsOptions
 
+if TYPE_CHECKING:
+    import xarray as xr
+
+
+def replace_json_invalid_floats(obj: object) -> object:
+    """
+    Recursively replace NaN and Infinity float values in a JSON-like object
+    with their string representations, to make them JSON-compliant.
+
+    Parameters
+    ----------
+    obj : object
+        The JSON-like object to process
+
+    Returns
+    -------
+    object
+        The processed object with NaN and Infinity replaced
+    """
+    if isinstance(obj, float):
+        if obj != obj:  # NaN check
+            return "NaN"
+        if obj == float("inf"):
+            return "Infinity"
+        if obj == float("-inf"):
+            return "-Infinity"
+        return obj
+    if isinstance(obj, dict):
+        return {k: replace_json_invalid_floats(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [replace_json_invalid_floats(item) for item in obj]
+    if isinstance(obj, tuple):
+        return tuple(replace_json_invalid_floats(item) for item in obj)
+
+    return obj
+
+
+class NanCompatibleJSONEncoder(json.JSONEncoder):
+    """
+    Custom JSON encoder that converts NaN, Inf, -Inf values to JSON-safe equivalents
+    to ensure valid JSON output.
+    """
+
+    def encode(self, obj: Any) -> str:
+        """
+        Encode object to JSON string, converting NaN values to "NaN".
+        """
+
+        converted_obj = replace_json_invalid_floats(obj)
+        return super().encode(converted_obj)
+
+
+def sanitize_dataset_attributes(ds: "xr.Dataset") -> "xr.Dataset":
+    """
+    Sanitize all NaN values in a dataset's attributes, variable attributes,
+    and coordinate attributes, converting them to "NaN" strings.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset to sanitize
+
+    Returns
+    -------
+    xr.Dataset
+        Dataset with sanitized attributes
+    """
+
+    # Create a copy to avoid modifying the original
+    ds_clean = ds.copy()
+
+    # Sanitize dataset attributes
+    ds_clean.attrs = replace_json_invalid_floats(ds_clean.attrs)
+
+    # Sanitize variable attributes
+    for var_name in ds_clean.data_vars:
+        var = ds_clean[var_name]
+        var.attrs = replace_json_invalid_floats(var.attrs)
+
+    # Sanitize coordinate attributes
+    for coord_name in ds_clean.coords:
+        coord = ds_clean[coord_name]
+        coord.attrs = replace_json_invalid_floats(coord.attrs)
+
+    return ds_clean
+
 
 def normalize_s3_path(s3_path: str) -> str:
     """
@@ -210,7 +296,7 @@ def write_s3_json_metadata(s3_path: str, metadata: Mapping[str, Any], **s3_kwarg
     fs = s3fs.S3FileSystem(**s3_config)
 
     # Write JSON content
-    json_content = json.dumps(metadata, indent=2)
+    json_content = json.dumps(metadata, indent=2, cls=NanCompatibleJSONEncoder)
     with fs.open(s3_path, "w") as f:
         f.write(json_content)
 
@@ -426,7 +512,7 @@ def write_json_metadata(path: str, metadata: dict[str, Any], **kwargs: Any) -> N
             fs.makedirs(parent_dir, exist_ok=True)
 
     # Write JSON content using fsspec
-    json_content = json.dumps(metadata, indent=2)
+    json_content = json.dumps(metadata, indent=2, cls=NanCompatibleJSONEncoder)
     with fs.open(path, "w") as f:
         f.write(json_content)
 

src/eopf_geozarr/conversion/geozarr.py

@@ -44,6 +44,7 @@
 )
 
 from . import fs_utils, utils
+from .fs_utils import sanitize_dataset_attributes
 from .sentinel1_reprojection import reproject_sentinel1_with_gcps
 
 log = structlog.get_logger()
@@ -350,6 +351,9 @@ def iterative_copy(
         if node.data_vars:
             log.info("Writing %s with data variables to GeoZarr DataTree", current_group_path)
 
+            # Sanitize NaN values in attributes before writing
+            ds = sanitize_dataset_attributes(ds)
+
             # Set up encoding
             encoding = _create_encoding(ds, compressor, spatial_chunk)
 
@@ -718,6 +722,9 @@ def create_geozarr_compliant_multiscales(
         if not fs_utils.is_s3_path(overview_path):
             os.makedirs(os.path.dirname(overview_path), exist_ok=True)
 
+        # Sanitize NaN values in overview dataset attributes
+        overview_ds = sanitize_dataset_attributes(overview_ds)
+
         # Write the overview dataset
         overview_group = f"{group_name}/{level}"
         # When sharding enabled, let Dask rechunk to shard boundaries
@@ -1170,6 +1177,9 @@ def cleanup_prefix(prefix: str) -> None:
                 else:
                     single_var_ds[var] = single_var_ds[var].chunk()
 
+                # Sanitize NaN values in single variable dataset attributes
+                single_var_ds = sanitize_dataset_attributes(single_var_ds)
+
                 single_var_ds.to_zarr(
                     output_path,
                     group=group_name,

src/eopf_geozarr/s2_optimization/s2_multiscale.py

@@ -16,6 +16,7 @@
 from pydantic.experimental.missing_sentinel import MISSING
 from pyproj import CRS
 
+from eopf_geozarr.conversion.fs_utils import sanitize_dataset_attributes
 from eopf_geozarr.conversion.geozarr import (
     _create_tile_matrix_limits,
     create_native_crs_tile_matrix_set,
@@ -873,6 +874,9 @@ def stream_write_dataset(
     if "/measurements/" in path or "/quality/" in path:
         write_geo_metadata(dataset, crs=crs)
 
+    # Sanitize NaN values in dataset attributes before writing
+    dataset = sanitize_dataset_attributes(dataset)
+
     # Write with streaming computation and progress tracking
     # The to_zarr operation will trigger all lazy computations
     write_job = dataset.to_zarr(