apply_neighborhood: support changing band names via apply_metadata (#1158)

* apply_neighborhood: support changing band names via apply_metadata
#1155

* apply_dimension: support changing band names via apply_metadata
#1155

* apply_dimension: support changing band names via apply_metadata
#1155

* apply_metadata only update metadata in more specific case
#1155

Author: jdries

CHANGELOG.md

@@ -12,8 +12,10 @@ without compromising stable operations.
 
 <!-- start-of-changelog -->
 
+
 ## In progress: 0.66.0
 
+- `apply_neighborhood`/`apply_dimension`: support changing band names via apply_metadata ([#1155](https://github.com/Open-EO/openeo-geopyspark-driver/issues/1155))
 
 ## 0.65.0
 

openeogeotrellis/geopysparkdatacube.py

@@ -630,7 +630,10 @@ def partitionByKey(spatialkey):
             _log.info(f"apply_neighborhood created datacube {metadata}")
             return gps.TiledRasterLayer.from_numpy_rdd(gps.LayerType.SPACETIME, numpy_rdd, metadata)
 
-        return self.apply_to_levels(partial(rdd_function, self.metadata))
+        updated_cube_metadata = self.metadata
+        if ("apply_metadata" in udf_code and runtime.lower() is not "python-jep"):
+            updated_cube_metadata = self.apply_metadata(udf_code, udf_context)
+        return self.apply_to_levels(partial(rdd_function, self.metadata), updated_cube_metadata)
 
     @callsite
     def chunk_polygon(
@@ -828,8 +831,11 @@ def tile_function(metadata: Metadata,
                 metadata = GeopysparkDataCube._transform_metadata(rdd.layer_metadata, cellType=CellType.FLOAT32)
                 return gps.TiledRasterLayer.from_numpy_rdd(rdd.layer_type, numpy_rdd, metadata)
 
+        updated_cube_metadata = self.metadata
+        if ("apply_metadata" in udf_code):
+            updated_cube_metadata = self.apply_metadata(udf_code, context)
         # Apply the UDF to every tile for every zoom level of the pyramid.
-        return self.apply_to_levels(partial(rdd_function, self.metadata))
+        return self.apply_to_levels(partial(rdd_function, self.metadata),updated_cube_metadata)
 
     def aggregate_time(self, temporal_window, aggregationfunction) -> Series :
         #group keys
@@ -1141,6 +1147,34 @@ def apply_kernel(self, kernel: np.ndarray, factor=1, border=0, replace_invalid=0
                 lambda rdd, level: pysc._jvm.org.openeo.geotrellis.OpenEOProcesses().apply_kernel_spatial(rdd,geotrellis_tile))
         return result_collection
 
+
+    def apply_metadata(self,udf_code,context):
+
+        pysc = gps.get_spark_context()
+        metadata = self.metadata
+        _log.info(f"run_udf: detected use of apply_metadata to transform: {self.metadata}")
+        def get_metadata(x):
+            from openeo.udf.run_code import load_module_from_string
+            module = load_module_from_string(udf_code)
+            functions = list([(k, v) for (k, v) in module.items() if callable(v) and k == "apply_metadata"])
+            if len(functions)>=1:
+                apply_metadata_func = functions[0][1]
+                transformed_metadata = apply_metadata_func(metadata,context)
+                return transformed_metadata
+            else:
+                raise ValueError("run_udf: apply_metadata function not found in the provided code.")
+        metadata_list = pysc.parallelize([0]).map(get_metadata).collect()
+        result_metadata: GeopysparkCubeMetadata = metadata_list[0]
+
+        _log.info(f"run_udf: apply_metadata resulted in {result_metadata}")
+        if not result_metadata.has_band_dimension():
+            raise ValueError(f"run_udf: apply_metadata function should not remove the band dimension, received metadata: {result_metadata}.")
+        if not isinstance(result_metadata, GeopysparkCubeMetadata):
+            raise ValueError(f"run_udf: apply_metadata function should retain the type of the input metadata object, received: {result_metadata}.")
+
+        return result_metadata
+
+
     @callsite
     def apply_neighborhood(
         self, process: dict, *, size: List[dict], overlap: List[dict], context: Optional[dict] = None, env: EvalEnv

tests/conftest.py

@@ -482,3 +482,33 @@ def unload_dummy_packages():
         if package in sys.modules:
             del sys.modules[package]
     importlib.invalidate_caches()
+
+
+@pytest.fixture
+def identity_udf_rename_bands():
+    udf_code = """
+from openeo.metadata import CollectionMetadata
+from xarray import DataArray
+
+def apply_metadata(metadata: CollectionMetadata, context: dict) -> CollectionMetadata:
+     return metadata.rename_labels(
+         dimension="bands",
+         target=["computed_band_1", "computed_band_2"]
+     )
+
+def apply_datacube(cube: DataArray, context: dict) -> DataArray:
+    return cube
+"""
+    udf_process = {
+        "udf_process": {
+            "process_id": "run_udf",
+            "arguments": {
+                "data": {
+                    "from_parameter": "data"
+                },
+                "udf": udf_code
+            },
+            "result": True
+        },
+    }
+    return udf_process

tests/test_apply_dimension.py

@@ -66,6 +66,17 @@ def test_apply_dimension_bands_udf(imagecollection_with_two_bands_and_three_date
     assert_array_almost_equal(input, subresult)
 
 
+def test_apply_metadata(imagecollection_with_two_bands_and_three_dates, identity_udf_rename_bands):
+
+    result = imagecollection_with_two_bands_and_three_dates.apply_dimension(
+        process=identity_udf_rename_bands, dimension="bands", target_dimension=None, context={}, env=EvalEnv()
+    )
+
+    assert result.metadata.band_names == ["computed_band_1", "computed_band_2"]
+    result_xarray = result._to_xarray()
+    assert list(result_xarray.bands.values) == ["computed_band_1", "computed_band_2"]
+
+
 def test_apply_dimension_invalid_dimension(imagecollection_with_two_bands_and_three_dates,udf_noop):
     the_date = datetime.datetime(2017, 9, 25, 11, 37)
     with pytest.raises(FeatureUnsupportedException):

tests/test_apply_neighborhood.py

@@ -62,6 +62,23 @@ def test_apply_neighborhood_overlap_udf(imagecollection_with_two_bands_and_three
     assert_array_almost_equal(input, subresult)
 
 
+def test_apply_metadata(imagecollection_with_two_bands_and_three_dates, identity_udf_rename_bands):
+
+    the_date = datetime.datetime(2017, 9, 25, 11, 37)
+
+    result = imagecollection_with_two_bands_and_three_dates.apply_neighborhood(
+        process=identity_udf_rename_bands,
+        size=[{'dimension': 'x', 'unit': 'px', 'value': 32}, {'dimension': 'y', 'unit': 'px', 'value': 32}],
+        overlap=[{'dimension': 'x', 'unit': 'px', 'value': 8}, {'dimension': 'y', 'unit': 'px', 'value': 8}],
+        context={},
+        env=EvalEnv()
+    )
+    assert result.metadata.band_names == ["computed_band_1", "computed_band_2"]
+    result_xarray = result._to_xarray()
+    assert list(result_xarray.bands.values) == ["computed_band_1", "computed_band_2"]
+
+
+
 def test_apply_neighborhood_on_timeseries(imagecollection_with_two_bands_and_three_dates):
     the_date = datetime.datetime(2017, 9, 25, 11, 37)
     graph = {