Issue #114/#211/#197 vector cube apply_dimension: put results in cu…

…be (too)
Open-EO · Aug 8, 2023 · c862e1b · c862e1b
1 parent 463afba
commit c862e1b
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 10 deletions.
diff --git a/openeo_driver/_version.py b/openeo_driver/_version.py
@@ -1 +1 @@
-__version__ = "0.61.1a1"
+__version__ = "0.61.2a1"
diff --git a/openeo_driver/datacube.py b/openeo_driver/datacube.py
@@ -676,7 +676,7 @@ def apply_dimension(
                     raise ValueError(
                         f"UDF should return single feature collection but got {result_features and len(result_features)}"
                     )
-                return DriverVectorCube(geometries=result_features[0].data)
+                return DriverVectorCube.from_geodataframe(result_features[0].data)
 
         raise FeatureUnsupportedException(
             message=f"DriverVectorCube.apply_dimension with {dimension=} and {bool(single_run_udf)=}"

diff --git a/tests/test_vectorcube.py b/tests/test_vectorcube.py
@@ -10,7 +10,7 @@
 
 from openeo_driver.datacube import DriverVectorCube
 from openeo_driver.errors import OpenEOApiException
-from openeo_driver.testing import ApproxGeometry, DictSubSet, IsNan
+from openeo_driver.testing import ApproxGeometry, DictSubSet, IsNan, ApproxGeoJSONByBounds
 from openeo_driver.util.geometry import as_geojson_feature_collection
 from openeo_driver.utils import EvalEnv
 
@@ -808,7 +808,7 @@ def test_apply_dimension_run_udf_change_geometry(self, gdf, backend_implementati
         udf = textwrap.dedent(
             """
             from openeo.udf import UdfData, FeatureCollection
-            def process_geometries(udf_data: UdfData) -> UdfData:
+            def process_vector_cube(udf_data: UdfData) -> UdfData:
                 [feature_collection] = udf_data.get_feature_collection_list()
                 gdf = feature_collection.data
                 gdf["geometry"] = gdf["geometry"].buffer(distance=1, resolution=2)
@@ -850,3 +850,62 @@ def process_geometries(udf_data: UdfData) -> UdfData:
                 ],
             }
         )
+
+    @pytest.mark.parametrize("dimension", ["bands", "properties"])
+    def test_apply_dimension_run_udf_add_properties(self, gdf, backend_implementation, dimension):
+        vc = DriverVectorCube.from_geodataframe(gdf, dimension_name=dimension)
+        udf = textwrap.dedent(
+            """
+            from openeo.udf import UdfData, FeatureCollection
+            def process_vector_cube(udf_data: UdfData) -> UdfData:
+                [feature_collection] = udf_data.get_feature_collection_list()
+                gdf = feature_collection.data
+                gdf["popone"] = gdf["pop"] + 1
+                gdf["poppop"] = gdf["pop"] ** 2
+                udf_data.set_feature_collection_list([
+                    FeatureCollection(id="_", data=gdf),
+                ])
+            """
+        )
+        callback = {
+            "runudf1": {
+                "process_id": "run_udf",
+                "arguments": {"data": {"from_parameter": "data"}, "udf": udf, "runtime": "Python"},
+                "result": True,
+            }
+        }
+        env = EvalEnv({"backend_implementation": backend_implementation})
+        result = vc.apply_dimension(process=callback, dimension=dimension, env=env)
+        assert isinstance(result, DriverVectorCube)
+        assert result.to_internal_json() == {
+            "geometries": {
+                "type": "FeatureCollection",
+                "features": [
+                    {
+                        "type": "Feature",
+                        "geometry": ApproxGeoJSONByBounds(1, 1, 3, 3, types=["Polygon"], abs=0.01),
+                        "id": "0",
+                        "properties": {"id": "first", "pop": 1234, "popone": 1235, "poppop": 1522756},
+                        "bbox": pytest.approx((1, 1, 3, 3), abs=0.01),
+                    },
+                    {
+                        "type": "Feature",
+                        "geometry": ApproxGeoJSONByBounds(3, 2, 5, 4, types=["Polygon"], abs=0.01),
+                        "id": "1",
+                        "properties": {"id": "second", "pop": 5678, "popone": 5679, "poppop": 32239684},
+                        "bbox": pytest.approx((3, 2, 5, 4), abs=0.01),
+                    },
+                ],
+                "bbox": pytest.approx((1, 1, 5, 4), abs=0.01),
+            },
+            "cube": {
+                "name": None,
+                "dims": ("geometries", "properties"),
+                "coords": {
+                    "geometries": {"attrs": {}, "data": [0, 1], "dims": ("geometries",)},
+                    "properties": {"attrs": {}, "data": ["pop", "popone", "poppop"], "dims": ("properties",)},
+                },
+                "data": [[1234, 1235, 1522756], [5678, 5679, 32239684]],
+                "attrs": {},
+            },
+        }
diff --git a/tests/test_views_execute.py b/tests/test_views_execute.py
@@ -3763,7 +3763,7 @@ def test_apply_dimension_run_udf_change_geometry(self, api100, dimension):
                     "process": self._build_run_udf_callback(
                         """
                         from openeo.udf import UdfData, FeatureCollection
-                        def process_geometries(udf_data: UdfData) -> UdfData:
+                        def process_vector_cube(udf_data: UdfData) -> UdfData:
                             [feature_collection] = udf_data.get_feature_collection_list()
                             gdf = feature_collection.data
                             gdf["geometry"] = gdf["geometry"].buffer(distance=1, resolution=2)
@@ -3827,7 +3827,7 @@ def test_apply_dimension_run_udf_filter_on_geometries(self, api100, dimension):
                         """
                         from openeo.udf import UdfData, FeatureCollection
                         import shapely.geometry
-                        def process_geometries(udf_data: UdfData) -> UdfData:
+                        def process_vector_cube(udf_data: UdfData) -> UdfData:
                             [feature_collection] = udf_data.get_feature_collection_list()
                             gdf = feature_collection.data
                             to_intersect = shapely.geometry.box(4, 3, 8, 4)
@@ -3896,8 +3896,7 @@ def test_apply_dimension_run_udf_filter_on_properties(self, api100, dimension):
                     "process": self._build_run_udf_callback(
                         """
                         from openeo.udf import UdfData, FeatureCollection
-                        import shapely.geometry
-                        def process_geometries(udf_data: UdfData) -> UdfData:
+                        def process_vector_cube(udf_data: UdfData) -> UdfData:
                             [feature_collection] = udf_data.get_feature_collection_list()
                             gdf = feature_collection.data
                             gdf = gdf[gdf["pop"] > 500]
@@ -3959,8 +3958,7 @@ def test_apply_dimension_run_udf_add_properties(self, api100, dimension):
                     "process": self._build_run_udf_callback(
                         """
                         from openeo.udf import UdfData, FeatureCollection
-                        import shapely.geometry
-                        def process_geometries(udf_data: UdfData) -> UdfData:
+                        def process_vector_cube(udf_data: UdfData) -> UdfData:
                             [feature_collection] = udf_data.get_feature_collection_list()
                             gdf = feature_collection.data
                             gdf["poppop"] = gdf["pop"] ** 2