Merge pull request #324 from Open-EO/323-dry_run-of-run_udf-process-o…

…nly-produces-jsonresult

ensure run_udf returns a DryRunDataCube if required #323

CHANGELOG.md

@@ -23,6 +23,7 @@ and start a new "In Progress" section above it.
 
 # 0.116.0
 - Propagate alternate `href`s of job result assets ([Open-EO/openeo-geopyspark-driver#883](https://github.com/Open-EO/openeo-geopyspark-driver/issues/883))
+- Ensure that a top level UDF can return a DriverVectorCube. Previously it only returned a JSONResult ([#323](https://github.com/Open-EO/openeo-python-driver/issues/323))
 
 # 0.115.0
 - Support pointing `href` of job result asset to workspace URI ([Open-EO/openeo-geopyspark-driver#883](https://github.com/Open-EO/openeo-geopyspark-driver/issues/883))

openeo_driver/ProcessGraphDeserializer.py

@@ -1516,7 +1516,7 @@ def run_udf(args: dict, env: EvalEnv):
         # Note: Other data types do execute the UDF during the dry-run.
         # E.g. A DelayedVector (when the user directly provides geometries as input).
         # This way a weak_spatial_extent can be calculated from the UDF's output.
-        return JSONResult({})
+        return data.run_udf()
 
     if isinstance(data, SupportsRunUdf) and data.supports_udf(udf=udf, runtime=runtime):
         _log.info(f"run_udf: data of type {type(data)} has direct run_udf support")

openeo_driver/_version.py

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

openeo_driver/dry_run.py

@@ -627,6 +627,8 @@ def raster_to_vector(self):
 
         return self._process(operation="raster_to_vector", arguments={},metadata=CollectionMetadata(metadata={}, dimensions=dimensions))
 
+    def run_udf(self):
+        return self._process(operation="run_udf", arguments={})
 
     def resample_cube_spatial(self, target: 'DryRunDataCube', method: str = 'near') -> 'DryRunDataCube':
         cube = self._process("process_type", [ProcessType.FOCAL_SPACE])

tests/test_views_execute.py

@@ -21,7 +21,7 @@
 from openeo_driver.datacube import DriverDataCube, DriverVectorCube
 from openeo_driver.datastructs import ResolutionMergeArgs, SarBackscatterArgs
 from openeo_driver.delayed_vector import DelayedVector
-from openeo_driver.dry_run import ProcessType
+from openeo_driver.dry_run import ProcessType, DryRunDataCube, DryRunDataTracer
 from openeo_driver.dummy import dummy_backend
 from openeo_driver.dummy.dummy_backend import DummyVisitor
 from openeo_driver.errors import (
@@ -1564,6 +1564,55 @@ def test_run_udf_on_list(api, udf_code):
     assert resp.json == [1, 4, 9, 25, 64]
 
 
+@pytest.mark.parametrize(
+    "udf_code",
+    [
+        """
+        from openeo.udf import UdfData, StructuredData
+        from geopandas import GeoDataFrame
+        from shapely.geometry import Point
+        def transform(data: UdfData) -> UdfData:
+            data.set_feature_collection_list([FeatureCollection("t", GeoDataFrame([{"geometry": Point(0.0, 0.1)}]))])
+            return data
+    """,
+    ],
+)
+def test_run_udf_on_aggregate_spatial(api, udf_code):
+    udf_code = textwrap.dedent(udf_code)
+    process_graph = {
+        "lc": {"process_id": "load_collection", "arguments": {"id": "S2_FOOBAR", "bands": ["B02", "B03", "B04"]}},
+        "ag": {
+            "process_id": "aggregate_spatial",
+            "arguments": {
+                "data": {"from_node": "lc"},
+                "geometries": {"type": "Point", "coordinates": [5.5, 51.5]},
+                "reducer": {
+                    "process_graph": {
+                        "mean": {
+                            "process_id": "mean",
+                            "arguments": {"data": {"from_parameter": "data"}},
+                            "result": True,
+                        }
+                    }
+                },
+            },
+        },
+        "udf": {
+            "process_id": "run_udf",
+            "arguments": {"data": {"from_node": "ag"}, "udf": udf_code, "runtime": "Python"},
+        },
+        "sr": {
+            "process_id": "save_result",
+            "arguments": {"data": {"from_node": "udf"}, "format": "GeoJSON"},
+            "result": True,
+        },
+    }
+    resp = api.check_result(process_graph)
+    assert resp.json["type"] == "FeatureCollection"
+    assert len(resp.json["features"]) == 1
+    assert resp.json["features"][0]["geometry"]["coordinates"] == [0.0, 0.1]
+
+
 @pytest.mark.parametrize(["runtime", "version", "failure"], [
     ("Python", None, None),
     ("pYthOn", None, None),