Merge branch 'issue394-load_stac-resample'

Author: soxofaan

CHANGELOG.md

@@ -23,6 +23,8 @@ and start a new "In Progress" section above it.
 
 - Introduce `asset_url` option to allow backend implementations to have custom code for retrieving assets. Default
   behavior remains unchanged.
+- Improve data cube dimension detection in `load_stac` dry-run ([#394](https://github.com/Open-EO/openeo-python-driver/issues/394))
+
 
 ## 0.133.0
 

openeo_driver/_version.py

@@ -1 +1 @@
-__version__ = "0.134.0a1"
+__version__ = "0.134.0a2"

openeo_driver/datacube.py

@@ -18,7 +18,7 @@
 import shapely.ops
 import xarray
 from geopandas import GeoDataFrame, GeoSeries
-from openeo.metadata import CollectionMetadata
+from openeo.metadata import CollectionMetadata, CubeMetadata
 from openeo.util import ensure_dir, str_truncate
 from pyproj import CRS
 
@@ -55,10 +55,12 @@ def run_udf(self, udf: str, *, runtime: str = "Python", context: Optional[dict]
 class DriverDataCube:
     """Base class for "driver" side raster data cubes."""
 
-    def __init__(self, metadata: CollectionMetadata = None):
-        self.metadata = (
-            metadata if isinstance(metadata, CollectionMetadata) else CollectionMetadata(metadata=metadata or {})
-        )
+    def __init__(self, metadata: Optional[CubeMetadata] = None):
+        if isinstance(metadata, dict):
+            # TODO: remove this security net once we're sure it's not necessary anymore
+            log.warning("DriverDataCube: deprecated dict-based metadata usage")
+            metadata = CollectionMetadata(metadata=metadata)
+        self.metadata: CubeMetadata = metadata or CubeMetadata()
 
     def __eq__(self, o: object) -> bool:
         if o.__class__ == self.__class__:

openeo_driver/dry_run.py

@@ -71,6 +71,7 @@
     GeometryBufferer,
     geojson_to_geometry, reproject_geometry,
 )
+import openeo_driver.stac.datacube
 from openeo_driver.utils import EvalEnv, to_hashable
 
 _log = logging.getLogger(__name__)
@@ -279,9 +280,10 @@ def load_collection(
         self, collection_id: str, arguments: dict, metadata: dict = None, env: EvalEnv = EvalEnv()
     ) -> "DryRunDataCube":
         """Create a DryRunDataCube from a `load_collection` process."""
-        # TODO #275 avoid VITO/Terrascope specific handling here?
+        metadata = CollectionMetadata(metadata=metadata)
         properties = {
-            **CollectionMetadata(metadata).get("_vito", "properties", default={}),
+            # TODO #275 avoid VITO/Terrascope specific handling here?
+            **metadata.get("_vito", "properties", default={}),
             **arguments.get("properties", {}),
         }
 
@@ -339,15 +341,20 @@ def load_stac(self, url: str, arguments: dict, env: EvalEnv = EvalEnv()) -> "Dry
         trace = DataSource.load_stac(url=url, properties=properties, bands=arguments.get("bands", []), env=env)
         self.add_trace(trace)
 
-        metadata = CollectionMetadata(
-            {},
-            dimensions=[
-                SpatialDimension(name="x", extent=[]),
-                SpatialDimension(name="y", extent=[]),
-                TemporalDimension(name="t", extent=[]),
-                BandDimension(name="bands", bands=[Band("unknown")]),
-            ],
-        )
+        try:
+            metadata = openeo_driver.stac.datacube.stac_to_cube_metadata(stac_ref=url)
+        except Exception as e:
+            _log.exception(
+                f"Dry-run load_stac: failed to parse cube metadata from {url!r} ({e!r}). Falling back in generic metadata"
+            )
+            metadata = CubeMetadata(
+                dimensions=[
+                    SpatialDimension(name="x", extent=[]),
+                    SpatialDimension(name="y", extent=[]),
+                    TemporalDimension(name="t", extent=[]),
+                    BandDimension(name="bands", bands=[Band("unknown")]),
+                ]
+            )
 
         cube = DryRunDataCube(traces=[trace], data_tracer=self, metadata=metadata)
         if "temporal_extent" in arguments:
@@ -534,7 +541,9 @@ class DryRunDataCube(DriverDataCube):
     estimate memory/cpu usage, ...
     """
 
-    def __init__(self, traces: List[DataTraceBase], data_tracer: DryRunDataTracer, metadata: CubeMetadata = None):
+    def __init__(
+        self, traces: List[DataTraceBase], data_tracer: DryRunDataTracer, metadata: Optional[CubeMetadata] = None
+    ):
         super(DryRunDataCube, self).__init__(metadata=metadata)
         self._traces = traces or []
         self._data_tracer = data_tracer

openeo_driver/dummy/dummy_backend.py

@@ -26,6 +26,7 @@
     CollectionMetadata,
     SpatialDimension,
     TemporalDimension,
+    CubeMetadata,
 )
 from openeo.util import rfc3339
 from openeo.utils.version import ComparableVersion
@@ -202,7 +203,7 @@ def mock_side_effect(fun):
 
 class DummyDataCube(DriverDataCube):
 
-    def __init__(self, metadata: CollectionMetadata = None):
+    def __init__(self, metadata: Optional[CubeMetadata] = None):
         super(DummyDataCube, self).__init__(metadata=metadata)
 
         # TODO #47: remove this non-standard process?
@@ -227,15 +228,15 @@ def __init__(self, metadata: CollectionMetadata = None):
     def reduce_dimension(
         self, reducer, *, dimension: str, context: Optional[dict] = None, env: EvalEnv
     ) -> "DummyDataCube":
-        return DummyDataCube(self.metadata.reduce_dimension(dimension_name=dimension))
+        return DummyDataCube(metadata=self.metadata.reduce_dimension(dimension_name=dimension))
 
     @mock_side_effect
     def add_dimension(self, name: str, label, type: str = "other") -> 'DummyDataCube':
-        return DummyDataCube(self.metadata.add_dimension(name=name, label=label, type=type))
+        return DummyDataCube(metadata=self.metadata.add_dimension(name=name, label=label, type=type))
 
     @mock_side_effect
     def drop_dimension(self, name: str) -> 'DriverDataCube':
-        return DummyDataCube(self.metadata.drop_dimension(name=name))
+        return DummyDataCube(metadata=self.metadata.drop_dimension(name=name))
 
     @mock_side_effect
     def dimension_labels(self, dimension: str) -> 'DriverDataCube':
@@ -1068,8 +1069,7 @@ def load_disk_data(
             self, format: str, glob_pattern: str, options: dict, load_params: LoadParameters, env: EvalEnv
     ) -> DummyDataCube:
         _register_load_collection_call(glob_pattern, load_params)
-        metadata = CollectionMetadata(
-            {},
+        metadata = CubeMetadata(
             dimensions=[
                 SpatialDimension(name="x", extent=[]),
                 SpatialDimension(name="y", extent=[]),

openeo_driver/stac/__init__.py

@@ -0,0 +1,84 @@
+"""
+Utilities for working with STAC datacube extension
+
+Based on pystac, with adaption layer to compensate for incompatibilities
+between available pystac version and actual STAC metadata.
+"""
+# TODO: move this functionality to the openeo-python-client for better reuse and client-server consistency?
+
+import logging
+from typing import Dict, Union
+from pathlib import Path
+
+import pystac
+import pystac.extensions.datacube
+
+import openeo.metadata
+
+_log = logging.getLogger(__name__)
+
+StacRef = Union[pystac.STACObject, str, Path]
+
+
+def as_stac_object(stac_ref: StacRef) -> pystac.STACObject:
+    if isinstance(stac_ref, pystac.STACObject):
+        return stac_ref
+    else:
+        return pystac.read_file(stac_ref)
+
+
+def _get_dimensions(stac_ref: StacRef) -> Dict[str, pystac.extensions.datacube.Dimension]:
+    stac_obj: pystac.STACObject = as_stac_object(stac_ref)
+    # TODO #396 update this to new pystac extension API
+    if pystac.extensions.datacube.DatacubeExtension.has_extension(stac_obj):
+        cube = pystac.extensions.datacube.DatacubeExtension.ext(stac_obj)
+        dimensions = cube.dimensions
+    elif any(e.startswith("https://stac-extensions.github.io/datacube/") for e in stac_obj.stac_extensions):
+        # TODO #370/#396 as we're currently stuck on an old pystac version that
+        #       doesn't support current versions of the datacube extension,
+        #       we need workarounds like this
+        _log.warning("Forcing pystac datacube extension on possibly unsupported metadata")
+        pystac.extensions.datacube.DatacubeExtension.add_to(stac_obj)
+        cube = pystac.extensions.datacube.DatacubeExtension.ext(stac_obj)
+        dimensions = cube.dimensions
+    else:
+        raise ValueError(f"No datacube extension found in STAC object {stac_ref=}")
+    return dimensions
+
+
+def stac_to_cube_metadata(stac_ref: StacRef) -> openeo.metadata.CubeMetadata:
+    """
+    Parse STAC metadata and convert it to :py:class:`openeo.metadata.CubeMetadata`.
+    """
+    # Dimensions as pystac objects
+    pystac_dimensions = _get_dimensions(stac_ref=stac_ref)
+
+    # Convert to openeo.metadata-style objects
+    dimensions = []
+    for name, dim in pystac_dimensions.items():
+        if isinstance(dim, pystac.extensions.datacube.HorizontalSpatialDimension):
+            dimensions.append(
+                openeo.metadata.SpatialDimension(
+                    name=name,
+                    extent=dim.extent,
+                    crs=dim.reference_system,
+                    step=dim.step,
+                )
+            )
+        elif isinstance(dim, pystac.extensions.datacube.TemporalDimension):
+            dimensions.append(
+                openeo.metadata.TemporalDimension(
+                    name=name,
+                    extent=dim.extent,
+                )
+            )
+        elif isinstance(dim, pystac.extensions.datacube.AdditionalDimension) and dim.dim_type == "bands":
+            dimensions.append(
+                openeo.metadata.BandDimension(
+                    name=name,
+                    bands=[openeo.metadata.Band(name=b) for b in dim.values],
+                )
+            )
+        else:
+            _log.info("Ignoring dimension %s of type %s", name, type(dim))
+    return openeo.metadata.CubeMetadata(dimensions=dimensions)

openeo_driver/stac/datacube.py

@@ -73,7 +73,7 @@
         "reretry~=0.11.8",
         "markdown>3.4",
         "traceback-with-variables==2.0.4",
-        "pystac~=1.8.0",
+        "pystac~=1.8.0",  # TODO #370/#396 bump to more recent pystac version
     ],
     extras_require={
         "dev": tests_require,

setup.py

@@ -0,0 +1,88 @@
+import json
+
+import pytest
+from openeo.testing.stac import StacDummyBuilder
+
+from openeo_driver.stac.datacube import stac_to_cube_metadata
+
+
+def test_stac_to_cube_metadata_basic(tmp_path):
+    path = tmp_path / "collection.json"
+    path.write_text(
+        json.dumps(
+            StacDummyBuilder.collection(
+                cube_dimensions={
+                    "x": {"type": "spatial", "axis": "x", "extent": [-180, 180], "reference_system": 4326},
+                    "y": {"type": "spatial", "axis": "y", "extent": [-56, 83], "reference_system": 4326},
+                    "t": {
+                        "type": "temporal",
+                        "extent": ["2015-06-23T00:00:00Z", "2019-07-10T13:44:56Z"],
+                        "step": "P5D",
+                    },
+                    "bands": {"type": "bands", "values": ["B1", "B2"]},
+                },
+            )
+        )
+    )
+
+    metadata = stac_to_cube_metadata(stac_ref=path)
+
+    [dim_x, dim_y] = metadata.spatial_dimensions
+    assert (dim_x.name, dim_x.extent, dim_x.crs, dim_x.step) == ("x", [-180, 180], 4326, None)
+    assert (dim_y.name, dim_y.extent, dim_y.crs, dim_y.step) == ("y", [-56, 83], 4326, None)
+
+    dim_t = metadata.temporal_dimension
+    assert (dim_t.name, dim_t.extent) == ("t", ["2015-06-23T00:00:00Z", "2019-07-10T13:44:56Z"])
+
+    dim_bands = metadata.band_dimension
+    assert (dim_bands.name, dim_bands.band_names) == ("bands", ["B1", "B2"])
+
+
+@pytest.mark.parametrize(
+    ["datacube_extension"],
+    [
+        ("https://stac-extensions.github.io/datacube/v2.0.0/schema.json",),
+        ("https://stac-extensions.github.io/datacube/v2.2.0/schema.json",),
+    ],
+)
+def test_stac_to_cube_metadata_version_compatibility(tmp_path, datacube_extension):
+    path = tmp_path / "collection.json"
+    path.write_text(
+        json.dumps(
+            {
+                "type": "Collection",
+                "stac_version": "1.0.0",
+                "stac_extensions": [datacube_extension],
+                "id": "collection123",
+                "description": "Collection 123",
+                "license": "proprietary",
+                "extent": {
+                    "spatial": {"bbox": [[3, 4, 5, 6]]},
+                    "temporal": {"interval": [["2024-01-01", "2024-05-05"]]},
+                },
+                "cube:dimensions": {
+                    "x": {"type": "spatial", "axis": "x", "extent": [-180, 180], "reference_system": 4326},
+                    "y": {"type": "spatial", "axis": "y", "extent": [-56, 83], "reference_system": 4326},
+                    "t": {
+                        "type": "temporal",
+                        "extent": ["2015-06-23T00:00:00Z", "2019-07-10T13:44:56Z"],
+                        "step": "P5D",
+                    },
+                    "bands": {"type": "bands", "values": ["B1", "B2"]},
+                },
+                "links": [],
+            }
+        )
+    )
+
+    metadata = stac_to_cube_metadata(stac_ref=path)
+
+    [dim_x, dim_y] = metadata.spatial_dimensions
+    assert (dim_x.name, dim_x.extent, dim_x.crs, dim_x.step) == ("x", [-180, 180], 4326, None)
+    assert (dim_y.name, dim_y.extent, dim_y.crs, dim_y.step) == ("y", [-56, 83], 4326, None)
+
+    dim_t = metadata.temporal_dimension
+    assert (dim_t.name, dim_t.extent) == ("t", ["2015-06-23T00:00:00Z", "2019-07-10T13:44:56Z"])
+
+    dim_bands = metadata.band_dimension
+    assert (dim_bands.name, dim_bands.band_names) == ("bands", ["B1", "B2"])