Skip to content

Commit

Permalink
Issue #197 Initial support for run_udf on DriverVectorCube
Browse files Browse the repository at this point in the history
  • Loading branch information
soxofaan committed Jun 6, 2023
1 parent fdf43b0 commit 97b8e58
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 18 deletions.
55 changes: 37 additions & 18 deletions openeo_driver/datacube.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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."""

Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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"""
Expand All @@ -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
):
...
55 changes: 55 additions & 0 deletions tests/test_vectorcube.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import textwrap

import geopandas as gpd
import numpy.testing
import pyproj
Expand Down Expand Up @@ -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},
},
),
],
}
)

0 comments on commit 97b8e58

Please sign in to comment.