Repository: ecmwf/earthkit-data
Pull request #514 (2024-11-07 11:29:26)

Make earthkit attributes optional in Xarray engine
This PR adds the `add_earthkit_attrs` kwarg to `to_xarray()`.

```rst
* add_earthkit_attrs: bool, None
   If True, add earthkit specific attributes to the dataset. Its default value
   (None) expands to True unless the ``profile`` overwrites it.
```

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/514?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 91.24%. Comparing base [(`eb0e16a`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/eb0e16aec0cc8542fa5fa8cbfc814e6e75bdd0eb?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf) to head [(`fa327ec`)](https://app.codecov.io/gh/ecmwf/earthkit-data/commit/fa327ec98287964e87ee0c51a09bb302ae65a253?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff            @@
##           develop     #514   +/-   ##
========================================
  Coverage    91.23%   91.24%           
========================================
  Files          147      147           
  Lines        10689    10698    +9     
  Branches       507      508    +1     
========================================
+ Hits          9752     9761    +9     
  Misses         774      774           
  Partials       163      163           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/ecmwf/earthkit-data/pull/514?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ecmwf).


Patch:
diff --git a/src/earthkit/data/readers/grib/xarray.py b/src/earthkit/data/readers/grib/xarray.py
index b906651..b02bb80 100644
--- a/src/earthkit/data/readers/grib/xarray.py
+++ b/src/earthkit/data/readers/grib/xarray.py
@@ -252,6 +252,9 @@ class XarrayMixIn:
                 default value (None) expands to [] unless the ``profile`` overwrites it.
             * coord_attrs: dict, None
                 To be documented. Default is None.
+            * add_earthkit_attrs: bool, None
+                If True, add earthkit specific attributes to the dataset. Its default value
+                (None) expands to True unless the ``profile`` overwrites it.
             * rename_attrs: dict, None
                 A dictionary of attribute to rename. Default is None.
             * remapping: dict, None
diff --git a/src/earthkit/data/utils/xarray/builder.py b/src/earthkit/data/utils/xarray/builder.py
index 5faff33..c146ad5 100644
--- a/src/earthkit/data/utils/xarray/builder.py
+++ b/src/earthkit/data/utils/xarray/builder.py
@@ -65,11 +65,13 @@ class VariableBuilder:
         self.tensor = tensor
         self.remapping = remapping
 
-    def build(self):
-        attrs = {
-            "message": self.tensor.source[0].metadata().override()._handle.get_buffer(),
-        }
-        self._attrs["_earthkit"] = attrs
+    def build(self, add_earthkit_attrs=True):
+        if add_earthkit_attrs:
+            attrs = {
+                "message": self.tensor.source[0].metadata().override()._handle.get_buffer(),
+            }
+            self._attrs["_earthkit"] = attrs
+
         self._attrs.update(self.fixed_local_attrs)
         data = self.data_maker(self.tensor, self.var_dims, self.name)
         return xarray.Variable(self.var_dims, data, attrs=self._attrs)
@@ -115,7 +117,8 @@ class VariableBuilder:
         # TODO: do we need a strict mode here? The extra cost has to be justified
         if keys_strict:
             assert strict
-            res.update(self.tensor.source.unique_values(keys_strict))
+            v, _ = self.tensor.source.unique_values(keys_strict)
+            res.update(v)
 
         self._attrs = res
         self._attrs.update(fixed_attrs)
@@ -299,7 +302,10 @@ class BackendDataBuilder(metaclass=ABCMeta):
         xr_coords = self.coords()
 
         # build variables
-        xr_vars = {self.profile.rename_variable(k): v.build() for k, v in var_builders.items()}
+        xr_vars = {
+            self.profile.rename_variable(k): v.build(add_earthkit_attrs=self.profile.add_earthkit_attrs)
+            for k, v in var_builders.items()
+        }
 
         # build dataset
         dataset = xarray.Dataset(xr_vars, coords=xr_coords, attrs=xr_attrs)
@@ -377,7 +383,9 @@ class BackendDataBuilder(metaclass=ABCMeta):
         # LOG.debug(f"{name=} {dims=}")
 
         # First check if the dims/coords are consistent with the tensors of the previous variables
-        vals, component_vals = ds.unique_values([d.key for d in dims], component=True)
+        vals, component_vals = ds.unique_values(
+            [d.key for d in dims], component=self.profile.add_earthkit_attrs
+        )
 
         # LOG.debug(f"unique_values={vals}")
         # LOG.debug(f"ensure_dims={self.profile.dims.ensure_dims}")
@@ -526,7 +534,9 @@ class DatasetBuilder:
         # LOG.debug(f"{profile.index_keys=}")
 
         # create a new fieldlist for optimised access to unique values
-        ds_xr = XArrayInputFieldList(ds, keys=profile.index_keys, remapping=remapping)
+        ds_xr = XArrayInputFieldList(
+            ds, keys=profile.index_keys, remapping=remapping, component=profile.add_earthkit_attrs
+        )
         # LOG.debug(f"{ds.db=}")
 
         # LOG.debug(f"before update: {profile.dim_keys=}")
@@ -609,7 +619,7 @@ class SplitDatasetBuilder(DatasetBuilder):
         # LOG.debug(f"split_dims={self.split_dims}")
         ds_xr = XArrayInputFieldList(self.ds, keys=self.profile.index_keys, remapping=remapping)
 
-        vals = ds_xr.unique_values(*keys)
+        vals, _ = ds_xr.unique_values(*keys)
 
         return ds_xr, vals
 
diff --git a/src/earthkit/data/utils/xarray/coord.py b/src/earthkit/data/utils/xarray/coord.py
index 3bbd043..1268e01 100644
--- a/src/earthkit/data/utils/xarray/coord.py
+++ b/src/earthkit/data/utils/xarray/coord.py
@@ -77,7 +77,7 @@ class Coord:
 
     def attrs(self, name, profile):
         attrs = profile.attrs.coord_attrs.get(name, {})
-        if self.component:
+        if profile.add_earthkit_attrs and self.component:
             attrs["_earthkit"] = {"keys": self.component[0], "values": self.component[1]}
         return attrs
 
diff --git a/src/earthkit/data/utils/xarray/defaults.yaml b/src/earthkit/data/utils/xarray/defaults.yaml
index 7498101..c4c4d1c 100644
--- a/src/earthkit/data/utils/xarray/defaults.yaml
+++ b/src/earthkit/data/utils/xarray/defaults.yaml
@@ -20,6 +20,7 @@ attrs_mode: "fixed"
 attrs: []
 variable_attrs: []
 global_attrs: []
+add_earthkit_attrs: true
 rename_attrs:
   cfName: standard_name
   name: long_name
diff --git a/src/earthkit/data/utils/xarray/engine.py b/src/earthkit/data/utils/xarray/engine.py
index 2bc152b..68cd548 100644
--- a/src/earthkit/data/utils/xarray/engine.py
+++ b/src/earthkit/data/utils/xarray/engine.py
@@ -43,6 +43,7 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
         variable_attrs=None,
         global_attrs=None,
         coord_attrs=None,
+        add_earthkit_attrs=None,
         rename_attrs=None,
         remapping=None,
         flatten_values=None,
@@ -210,6 +211,9 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
             default value (None) expands to [] unless the ``profile`` overwrites it.
         coord_attrs: dict, None
             To be documented. Default is None.
+        add_earthkit_attrs: bool, None
+            If True, add earthkit specific attributes to the dataset. Its default value
+            (None) expands to True unless the ``profile`` overwrites it.
         rename_attrs: dict, None
             A dictionary of attribute to rename. Default is None.
         remapping: dict, None
@@ -265,6 +269,7 @@ class EarthkitBackendEntrypoint(BackendEntrypoint):
                 variable_attrs=variable_attrs,
                 global_attrs=global_attrs,
                 coord_attrs=coord_attrs,
+                add_earthkit_attrs=add_earthkit_attrs,
                 rename_attrs=rename_attrs,
                 add_valid_time_coord=add_valid_time_coord,
                 add_geo_coords=add_geo_coords,
diff --git a/src/earthkit/data/utils/xarray/fieldlist.py b/src/earthkit/data/utils/xarray/fieldlist.py
index 3d44e82..c4186da 100644
--- a/src/earthkit/data/utils/xarray/fieldlist.py
+++ b/src/earthkit/data/utils/xarray/fieldlist.py
@@ -48,6 +48,7 @@ class IndexSelection(Selection):
 
 class IndexDB:
     def __init__(self, index, component):
+        # print(f"IndexDB: {index=}, {component=}")
         self._index = index if index is not None else dict()
         self._component = component if component is not None else dict()
 
@@ -56,7 +57,7 @@ class IndexDB:
         if key not in self._index:
             # # LOG.debug(f"Key={key} not found in IndexDB")
             if maker is not None:
-                self._index[key] = maker(key)[key]
+                self._index[key] = maker(key)[0][key]
             else:
                 raise KeyError(f"Could not find index for {key=}")
         return self._index[key]
@@ -102,7 +103,7 @@ class IndexDB:
 
 
 class XArrayInputFieldList(FieldList):
-    def __init__(self, fieldlist, keys=None, db=None, remapping=None, scan_only=False):
+    def __init__(self, fieldlist, keys=None, db=None, remapping=None, scan_only=False, component=True):
         super().__init__()
         self.ds = fieldlist
 
@@ -114,7 +115,7 @@ class XArrayInputFieldList(FieldList):
         if db is not None:
             self.db = db
         elif keys:
-            self.db = IndexDB(*self.unique_values(keys, component=True))
+            self.db = IndexDB(*self.unique_values(keys, component=component))
 
         assert self.db
 
@@ -225,7 +226,7 @@ class XArrayInputFieldList(FieldList):
         if component:
             return indices, components
         else:
-            return indices
+            return indices, None
 
 
 class ReleasableField:
diff --git a/src/earthkit/data/utils/xarray/profile.py b/src/earthkit/data/utils/xarray/profile.py
index 5d88d29..5a5bbe6 100644
--- a/src/earthkit/data/utils/xarray/profile.py
+++ b/src/earthkit/data/utils/xarray/profile.py
@@ -146,6 +146,8 @@ class Profile:
             kwargs.pop("rename_attrs"),
         )
 
+        self.add_earthkit_attrs = kwargs.pop("add_earthkit_attrs")
+
         # generic
         self.decode_times = kwargs.pop("decode_times")
         self.decode_timedelta = kwargs.pop("decode_timedelta")
diff --git a/src/earthkit/data/utils/xarray/splitter.py b/src/earthkit/data/utils/xarray/splitter.py
index f1ed9ab..5a5d77b 100644
--- a/src/earthkit/data/utils/xarray/splitter.py
+++ b/src/earthkit/data/utils/xarray/splitter.py
@@ -18,7 +18,7 @@ LOG = logging.getLogger(__name__)
 class Splitter(metaclass=ABCMeta):
     @staticmethod
     def grids(ds):
-        v = ds.unique_values(["md5GridSection"])
+        v, _ = ds.unique_values(["md5GridSection"])
         return v["md5GridSection"]
 
     @abstractmethod


Test patch:
diff --git a/tests/xr_engine/test_xr_engine.py b/tests/xr_engine/test_xr_engine.py
index f68ecc2..40735b6 100644
--- a/tests/xr_engine/test_xr_engine.py
+++ b/tests/xr_engine/test_xr_engine.py
@@ -505,3 +505,23 @@ def test_xr_engine_single_field():
     for k, v in coords_ref_full.items():
         assert np.allclose(ds.coords[k].values, v)
     assert [v for v in ds.data_vars] == data_vars
+
+
+@pytest.mark.cache
+@pytest.mark.parametrize("add", [False, True])
+def test_xr_engine_add_earthkit_attrs(add):
+    ds_ek = from_source("url", earthkit_remote_test_data_file("test-data/xr_engine/level/pl.grib"))
+    ds_ek = ds_ek[0]
+
+    ds = ds_ek.to_xarray(
+        time_dim_mode="raw",
+        decode_times=False,
+        decode_timedelta=False,
+        add_valid_time_coord=False,
+        add_earthkit_attrs=add,
+    )
+
+    if add:
+        assert "_earthkit" in ds["t"].attrs
+    else:
+        assert "_earthkit" not in ds["t"].attrs