Issue #197 Initial support for run_udf on DriverVectorCube

openeo_driver/datacube.py

@@ -16,6 +16,8 @@
 from pyproj import CRS
 import requests
 
+import openeo.udf
+import openeo.udf.run_code
 from openeo.metadata import CollectionMetadata
 from openeo.util import ensure_dir
 from openeo_driver.datastructs import SarBackscatterArgs, ResolutionMergeArgs, StacAsset
@@ -28,6 +30,21 @@
 log = logging.getLogger(__name__)
 
 
+class SupportsRunUdf(metaclass=abc.ABCMeta):
+    """
+    Interface for cube/result classes that (partially) support `run_udf`
+    """
+
+    @abc.abstractmethod
+    def supports_udf(self, udf: str, runtime: str = "Python") -> bool:
+        """Check if UDF code is supported."""
+        return False
+
+    @abc.abstractmethod
+    def run_udf(self, udf: str, runtime: str = "Python", context: Optional[dict] = None):
+        ...
+
+
 class DriverDataCube:
     """Base class for "driver" side raster data cubes."""
 
@@ -181,7 +198,7 @@ def __init__(self, message="Unspecified VectorCube error"):
         super(VectorCubeError, self).__init__(message=message)
 
 
-class DriverVectorCube:
+class DriverVectorCube(SupportsRunUdf):
     """
     Base class for driver-side 'vector cubes'
 
@@ -481,6 +498,25 @@ def buffer_points(self, distance: float = 10) -> "DriverVectorCube":
             ]
         )
 
+    def supports_udf(self, udf: str, runtime: str = "Python") -> bool:
+        udf_globals = openeo.udf.run_code.load_module_from_string(code=udf)
+        return any(
+            name in udf_globals
+            for name in [
+                "udf_apply_geojson_feature_collection",
+            ]
+        )
+
+    def run_udf(self, *, udf: str, runtime: str = "Python", context: Optional[dict] = None):
+        udf_globals = openeo.udf.run_code.load_module_from_string(code=udf)
+
+        if "udf_apply_geojson_feature_collection" in udf_globals:
+            callback = udf_globals["udf_apply_geojson_feature_collection"]
+            result = callback(self.to_geojson())
+            return DriverVectorCube.from_geojson(geojson=result)
+        else:
+            raise openeo.udf.OpenEoUdfException("No UDF found")
+
 
 class DriverMlModel:
     """Base class for driver-side 'ml-model' data structures"""
@@ -491,20 +527,3 @@ def get_model_metadata(self, directory: Union[str, Path]) -> Dict[str, Any]:
 
     def write_assets(self, directory: Union[str, Path]) -> Dict[str, StacAsset]:
         raise NotImplementedError
-
-
-class SupportsRunUdf(metaclass=abc.ABCMeta):
-    """
-    Interface for cube/result classes that (partially) support `run_udf`
-    """
-
-    @abc.abstractmethod
-    def supports_udf(self, udf: str, runtime: str = "Python") -> bool:
-        """Check if UDF code is supported."""
-        return False
-
-    @abc.abstractmethod
-    def run_udf(
-        self, udf: str, runtime: str = "Python", context: Optional[dict] = None
-    ):
-        ...

tests/test_vectorcube.py

@@ -1,3 +1,5 @@
+import textwrap
+
 import geopandas as gpd
 import numpy.testing
 import pyproj
@@ -446,3 +448,56 @@ def test_buffer_points(self):
                 ],
             }
         )
+
+    def test_run_udf_geojson(self, gdf):
+        vc = DriverVectorCube(gdf)
+        udf = textwrap.dedent(
+            """
+            def udf_apply_geojson_feature_collection(feature_collection: dict) -> dict:
+                features = []
+                for feature in feature_collection["features"]:
+                    # transpose and rescale vertices
+                    vertices = feature["geometry"]["coordinates"][0]
+                    vertices = [(y, 2 * x) for (x, y) in vertices]
+                    feature["geometry"]["coordinates"] = [vertices]
+
+                    # Manipulate properties
+                    feature["properties"]["hello"] = feature["properties"]["id"]
+                    feature["properties"]["vertices"] = len(vertices) - 1
+                    del feature["properties"]["pop"]
+
+                    features.append(feature)
+                return {"type": "FeatureCollection", "features": features}
+        """
+        )
+        vc2 = vc.run_udf(udf=udf)
+        assert isinstance(vc2, DriverVectorCube)
+        assert vc2.to_geojson() == DictSubSet(
+            {
+                "type": "FeatureCollection",
+                "features": [
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "id": "0",
+                            "geometry": {
+                                "type": "Polygon",
+                                "coordinates": (((1, 2), (1, 6), (3, 4), (1, 2)),),
+                            },
+                            "properties": {"hello": "first", "id": "first", "vertices": 3},
+                        }
+                    ),
+                    DictSubSet(
+                        {
+                            "type": "Feature",
+                            "id": "1",
+                            "geometry": {
+                                "type": "Polygon",
+                                "coordinates": (((2, 8), (4, 10), (4, 6), (2, 8)),),
+                            },
+                            "properties": {"hello": "second", "id": "second", "vertices": 3},
+                        },
+                    ),
+                ],
+            }
+        )