PR #200 change back to apply_dimension with dimension=bands

Open-EO · Aug 3, 2023 · 4f97c75 · 4f97c75
1 parent eb03aa8
commit 4f97c75
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/openeo_driver/datacube.py b/openeo_driver/datacube.py
@@ -216,6 +216,7 @@ class DriverVectorCube:
     These components are "joined" on the GeoPandas dataframe's index and DataArray first dimension
     """
     DIM_GEOMETRIES = "geometries"
+    DIM_BANDS = "bands"
     FLATTEN_PREFIX = "vc"
 
     def __init__(
@@ -519,7 +520,7 @@ def apply_dimension(
 
         if single_run_udf:
             # Process with single "run_udf" node
-            if self._cube is None and dimension == self.DIM_GEOMETRIES and target_dimension is None:
+            if dimension == self.DIM_BANDS and target_dimension is None:
                 log.warning(
                     f"Using experimental feature: DriverVectorCube.apply_dimension along dim {dimension} and empty cube"
                 )
@@ -535,14 +536,18 @@ def apply_dimension(
                 result_data = env.backend_implementation.processing.run_udf(udf=single_run_udf.udf, data=udf_data)
                 log.info(f"[run_udf] UDF resulted in {result_data!r}")
 
-                if isinstance(result_data, openeo.udf.UdfData):
-                    result_features = result_data.get_feature_collection_list()
-                    if result_features and len(result_features) == 1:
-                        return DriverVectorCube(geometries=result_features[0].data)
-                raise ValueError(f"Could not handle UDF result: {result_data}")
-
-        raise FeatureUnsupportedException()
-
+                if not isinstance(result_data, openeo.udf.UdfData):
+                    raise ValueError(f"UDF should return UdfData, but got {type(result_data)}")
+                result_features = result_data.get_feature_collection_list()
+                if not (result_features and len(result_features) == 1):
+                    raise ValueError(
+                        f"UDF should return single feature collection but got {result_features and len(result_features)}"
+                    )
+                return DriverVectorCube(geometries=result_features[0].data)
+
+        raise FeatureUnsupportedException(
+            message=f"DriverVectorCube.apply_dimension with {dimension=} and {bool(single_run_udf)=}"
+        )
 
 
 class DriverMlModel:

diff --git a/tests/test_vectorcube.py b/tests/test_vectorcube.py
@@ -454,7 +454,7 @@ def test_buffer_points(self):
             }
         )
 
-    def test_apply_dimension_run_udf(self, vc, backend_implementation):
+    def test_apply_dimension_run_udf_change_geometry(self, vc, backend_implementation):
         udf = textwrap.dedent(
             """
             from openeo.udf import UdfData, FeatureCollection
@@ -475,7 +475,7 @@ def process_geometries(udf_data: UdfData) -> UdfData:
             }
         }
         env = EvalEnv({"backend_implementation": backend_implementation})
-        result = vc.apply_dimension(process=callback, dimension="geometries", env=env)
+        result = vc.apply_dimension(process=callback, dimension="bands", env=env)
         assert isinstance(result, DriverVectorCube)
         feature_collection = result.to_geojson()
         assert feature_collection == DictSubSet(