Repository: ecmwf/earthkit-data
Pull request #199 (2023-09-20 14:52:28)

Implement lazy fieldlist scanning for NetCDF data
The scanning for fields in a NetCDF fieldlist is now delayed until the fields are actually needed.

```python
# fields are not scanned from the input data when creating ds
ds = from_source("file", earthkit_examples_file("test.nc"))

# field scanning starts when calling len()
print(len(ds))
```


Patch:
diff --git a/earthkit/data/readers/netcdf.py b/earthkit/data/readers/netcdf.py
index 90cd068..5bf7333 100644
--- a/earthkit/data/readers/netcdf.py
+++ b/earthkit/data/readers/netcdf.py
@@ -144,6 +144,122 @@ class DataSet:
         return self._bbox[(lat, lon)]
 
 
+def get_fields_from_ds(
+    ds,
+    field_type=None,
+    check_only=False,
+):  # noqa C901
+    # Select only geographical variables
+    has_lat = False
+    has_lon = False
+
+    fields = []
+
+    skip = set()
+
+    for name in ds.data_vars:
+        v = ds[name]
+        skip.update(getattr(v, "coordinates", "").split(" "))
+        skip.update(getattr(v, "bounds", "").split(" "))
+        skip.update(getattr(v, "grid_mapping", "").split(" "))
+
+    for name in ds.data_vars:
+        if name in skip:
+            continue
+
+        v = ds[name]
+
+        coordinates = []
+
+        # self.log.info('Scanning file: %s var=%s coords=%s', self.path, name, v.coords)
+
+        info = [value for value in v.coords if value not in v.dims]
+        non_dim_coords = {}
+        for coord in v.coords:
+            if coord not in v.dims:
+                non_dim_coords[coord] = ds[coord].values
+                continue
+
+            c = ds[coord]
+
+            # self.log.info("COORD %s %s %s %s", coord, type(coord), hasattr(c, 'calendar'), c)
+
+            standard_name = getattr(c, "standard_name", "")
+            axis = getattr(c, "axis", "")
+            long_name = getattr(c, "long_name", "")
+
+            use = False
+
+            if (
+                standard_name.lower() in GEOGRAPHIC_COORDS["x"]
+                or (long_name == "longitude")
+                or (axis == "X")
+            ):
+                has_lon = True
+                use = True
+
+            if (
+                standard_name.lower() in GEOGRAPHIC_COORDS["y"]
+                or (long_name == "latitude")
+                or (axis == "Y")
+            ):
+                has_lat = True
+                use = True
+
+            # print(f"  standard_name={standard_name}")
+
+            # Of course, not every one sets the standard_name
+            if (
+                standard_name in ["time", "forecast_reference_time"]
+                or long_name in ["time"]
+                or axis == "T"
+            ):
+                # we might not be able to convert time to datetime
+                try:
+                    coordinates.append(TimeCoordinate(c, coord in info))
+                    use = True
+                except ValueError:
+                    break
+
+            # TODO: Support other level types
+            if standard_name in [
+                "air_pressure",
+                "model_level_number",
+                "altitude",
+            ] or long_name in [
+                "pressure_level"
+            ]:  # or axis == 'Z':
+                coordinates.append(LevelCoordinate(c, coord in info))
+                use = True
+
+            if axis in ("X", "Y"):
+                use = True
+
+            if not use:
+                coordinates.append(OtherCoordinate(c, coord in info))
+
+        if not (has_lat and has_lon):
+            # self.log.info("NetCDFReader: skip %s (Not a 2 field)", name)
+            continue
+
+        for values in product(*[c.values for c in coordinates]):
+            slices = []
+            for value, coordinate in zip(values, coordinates):
+                slices.append(coordinate.make_slice(value))
+
+            if check_only:
+                return True
+
+            fields.append(field_type(ds, name, slices, non_dim_coords))
+
+    # if not fields:
+    #     raise Exception("NetCDFReader no 2D fields found in %s" % (self.path,))
+
+    if check_only:
+        return False
+    return fields
+
+
 class XArrayFieldGeography(Geography):
     def __init__(self, metadata, da, ds, variable):
         self.metadata = metadata
@@ -253,7 +369,6 @@ class XArrayMetadata(RawMetadata):
 class XArrayField(Field):
     def __init__(self, ds, variable, slices, non_dim_coords):
         super().__init__()
-
         self._ds = ds
         self._da = ds[variable]
 
@@ -343,118 +458,20 @@ class XArrayFieldListCore(FieldList):
             self._scan()
         return self._fields
 
+    def has_fields(self):
+        if self._fields is None:
+            return get_fields_from_ds(
+                DataSet(self.ds), field_type=self.FIELD_TYPE, check_only=True
+            )
+        else:
+            return len(self._fields)
+
     def _scan(self):
         if self._fields is None:
             self._fields = self._get_fields()
 
     def _get_fields(self):
-        return self._get_fields_from_ds(DataSet(self.ds))
-
-    def _get_fields_from_ds(self, ds):  # noqa C901
-        # Select only geographical variables
-        has_lat = False
-        has_lon = False
-
-        fields = []
-
-        skip = set()
-
-        for name in ds.data_vars:
-            v = ds[name]
-            skip.update(getattr(v, "coordinates", "").split(" "))
-            skip.update(getattr(v, "bounds", "").split(" "))
-            skip.update(getattr(v, "grid_mapping", "").split(" "))
-
-        for name in ds.data_vars:
-            if name in skip:
-                continue
-
-            v = ds[name]
-
-            coordinates = []
-
-            # self.log.info('Scanning file: %s var=%s coords=%s', self.path, name, v.coords)
-
-            info = [value for value in v.coords if value not in v.dims]
-            non_dim_coords = {}
-            for coord in v.coords:
-                if coord not in v.dims:
-                    non_dim_coords[coord] = ds[coord].values
-                    continue
-
-                c = ds[coord]
-
-                # self.log.info("COORD %s %s %s %s", coord, type(coord), hasattr(c, 'calendar'), c)
-
-                standard_name = getattr(c, "standard_name", "")
-                axis = getattr(c, "axis", "")
-                long_name = getattr(c, "long_name", "")
-
-                use = False
-
-                if (
-                    standard_name.lower() in GEOGRAPHIC_COORDS["x"]
-                    or (long_name == "longitude")
-                    or (axis == "X")
-                ):
-                    has_lon = True
-                    use = True
-
-                if (
-                    standard_name.lower() in GEOGRAPHIC_COORDS["y"]
-                    or (long_name == "latitude")
-                    or (axis == "Y")
-                ):
-                    has_lat = True
-                    use = True
-
-                # print(f"  standard_name={standard_name}")
-
-                # Of course, not every one sets the standard_name
-                if (
-                    standard_name in ["time", "forecast_reference_time"]
-                    or long_name in ["time"]
-                    or axis == "T"
-                ):
-                    # we might not be able to convert time to datetime
-                    try:
-                        coordinates.append(TimeCoordinate(c, coord in info))
-                        use = True
-                    except ValueError:
-                        break
-
-                # TODO: Support other level types
-                if standard_name in [
-                    "air_pressure",
-                    "model_level_number",
-                    "altitude",
-                ] or long_name in [
-                    "pressure_level"
-                ]:  # or axis == 'Z':
-                    coordinates.append(LevelCoordinate(c, coord in info))
-                    use = True
-
-                if axis in ("X", "Y"):
-                    use = True
-
-                if not use:
-                    coordinates.append(OtherCoordinate(c, coord in info))
-
-            if not (has_lat and has_lon):
-                # self.log.info("NetCDFReader: skip %s (Not a 2 field)", name)
-                continue
-
-            for values in product(*[c.values for c in coordinates]):
-                slices = []
-                for value, coordinate in zip(values, coordinates):
-                    slices.append(coordinate.make_slice(value))
-
-                fields.append(self.FIELD_TYPE(ds, name, slices, non_dim_coords))
-
-        # if not fields:
-        #     raise Exception("NetCDFReader no 2D fields found in %s" % (self.path,))
-
-        return fields
+        return get_fields_from_ds(DataSet(self.ds), field_type=self.FIELD_TYPE)
 
     def to_pandas(self):
         return self.to_xarray().to_pandas()
@@ -524,8 +541,9 @@ class NetCDFField(XArrayField):
 
 
 class NetCDFFieldList(XArrayFieldListCore):
+    FIELD_TYPE = NetCDFField
+
     def __init__(self, path, *args, **kwargs):
-        self.FIELD_TYPE = NetCDFField
         self.path = path
         # self._fields = None
         super().__init__(None, *args, **kwargs)
@@ -536,7 +554,20 @@ class NetCDFFieldList(XArrayFieldListCore):
         with closing(
             xr.open_mfdataset(self.path, combine="by_coords")
         ) as ds:  # or nested
-            return self._get_fields_from_ds(DataSet(ds))
+            return get_fields_from_ds(DataSet(ds), field_type=self.FIELD_TYPE)
+
+    def has_fields(self):
+        if self._fields is None:
+            import xarray as xr
+
+            with closing(
+                xr.open_mfdataset(self.path, combine="by_coords")
+            ) as ds:  # or nested
+                return get_fields_from_ds(
+                    DataSet(ds), field_type=self.FIELD_TYPE, check_only=True
+                )
+        else:
+            return len(self._fields)
 
     @classmethod
     def merge(cls, sources):
@@ -670,7 +701,7 @@ def _match_magic(magic, deeper_check):
 def reader(source, path, magic=None, deeper_check=False):
     if _match_magic(magic, deeper_check):
         fs = NetCDFFieldListReader(source, path)
-        if len(fs) > 0:
+        if fs.has_fields():
             return fs
         else:
             return NetCDFReader(source, path)
diff --git a/earthkit/data/wrappers/xarray.py b/earthkit/data/wrappers/xarray.py
index 4d088ea..189e7c6 100644
--- a/earthkit/data/wrappers/xarray.py
+++ b/earthkit/data/wrappers/xarray.py
@@ -145,7 +145,7 @@ def wrapper(data, *args, **kwargs):
 
     if ds is not None:
         fs = netcdf.XArrayFieldList(ds, **kwargs)
-        if len(fs) > 0:
+        if fs.has_fields():
             return fs
         else:
             return XArrayDatasetWrapper(ds, *args, **kwargs)


Test patch:
diff --git a/tests/readers/test_netcdf_reader.py b/tests/readers/test_netcdf_reader.py
index 9011383..1ebe5d1 100644
--- a/tests/readers/test_netcdf_reader.py
+++ b/tests/readers/test_netcdf_reader.py
@@ -248,6 +248,14 @@ def test_netcdf_non_fieldlist():
     ek_ch4_l2.to_xarray()
 
 
+@pytest.mark.no_eccodes
+def test_netcdf_lazy_fieldlist_scan():
+    ds = from_source("file", earthkit_examples_file("test.nc"))
+    assert ds._fields is None
+    assert len(ds) == 2
+    assert len(ds._fields) == 2
+
+
 if __name__ == "__main__":
     from earthkit.data.testing import main
 
diff --git a/tests/wrappers/test_xarray.py b/tests/wrappers/test_xarray.py
index fc8e739..3088132 100644
--- a/tests/wrappers/test_xarray.py
+++ b/tests/wrappers/test_xarray.py
@@ -9,15 +9,14 @@
 # nor does it submit to any jurisdiction.
 #
 
-
-import logging
+import pytest
 
 from earthkit.data import from_object, wrappers
+from earthkit.data.testing import earthkit_examples_file
 from earthkit.data.wrappers import xarray as xr_wrapper
 
-LOG = logging.getLogger(__name__)
-
 
+@pytest.mark.no_eccodes
 def test_dataset_wrapper():
     import xarray as xr
 
@@ -29,6 +28,7 @@ def test_dataset_wrapper():
     assert isinstance(_wrapper, xr_wrapper.XArrayDatasetWrapper)
 
 
+@pytest.mark.no_eccodes
 def test_dataarray_wrapper():
     import xarray as xr
 
@@ -38,3 +38,13 @@ def test_dataarray_wrapper():
     assert isinstance(_wrapper, xr_wrapper.XArrayDataArrayWrapper)
     _wrapper = from_object(xr.DataArray())
     assert isinstance(_wrapper, xr_wrapper.XArrayDataArrayWrapper)
+
+
+@pytest.mark.no_eccodes
+def test_xarray_lazy_fieldlist_scan():
+    import xarray as xr
+
+    ds = from_object(xr.open_dataset(earthkit_examples_file("test.nc")))
+    assert ds._fields is None
+    assert len(ds) == 2
+    assert len(ds._fields) == 2