labelled arrays

openeo_processes_dask/process_implementations/arrays.py

@@ -7,6 +7,7 @@
 import pandas as pd
 import xarray as xr
 from numpy.typing import ArrayLike
+from openeo_pg_parser_networkx.pg_schema import DateTime
 from xarray.core.duck_array_ops import isnull, notnull
 
 from openeo_processes_dask.process_implementations.cubes.utils import _is_dask_array
@@ -42,6 +43,8 @@
     label: Optional[str] = None,
     return_nodata: Optional[bool] = False,
     axis=None,
+    context=None,
+    dim_labels=None,
 ):
     if index is None and label is None:
         raise ArrayElementParameterMissing(
@@ -54,14 +57,20 @@
         )
 
     if label is not None:
-        raise NotImplementedError(
-            "labelled arrays are currently not implemented. Please use index instead."
-        )
+        if isinstance(label, DateTime):
+            label = label.to_numpy()
+        (index,) = np.where(dim_labels == label)
+        if len(index) == 0:
+            index = None
+        else:
+            index = index[0]
 
     try:
         if index is not None:
             element = np.take(data, index, axis=axis)
             return element
+        else:
+            raise IndexError
     except IndexError:
         if return_nodata:
             logger.warning(

openeo_processes_dask/process_implementations/core.py

@@ -68,7 +68,13 @@ def wrapper(
             else:
                 resolved_kwargs[k] = arg
 
-        special_args = ["axis", "keepdims", "source_transposed_axis"]
+        special_args = [
+            "axis",
+            "keepdims",
+            "source_transposed_axis",
+            "context",
+            "dim_labels",
+        ]
         # Remove 'axis' and keepdims parameter if not expected in function signature.
         for arg in special_args:
             if arg not in inspect.signature(f).parameters:

openeo_processes_dask/process_implementations/cubes/reduce.py

@@ -21,15 +21,16 @@ def reduce_dimension(
             f"Provided dimension ({dimension}) not found in data.dims: {data.dims}"
         )
 
-    positional_parameters = {"data": 0}
-    named_parameters = {"context": context}
+    dim_labels = data[dimension].values
 
+    positional_parameters = {"data": 0}
     reduced_data = data.reduce(
         reducer,
         dim=dimension,
         keep_attrs=True,
         positional_parameters=positional_parameters,
-        named_parameters=named_parameters,
+        context=context,
+        dim_labels=dim_labels,
     )
 
     # Preset

tests/test_arrays.py

@@ -48,6 +48,24 @@ def test_array_element(
 
     xr.testing.assert_equal(output_cube, input_cube.isel({"bands": 1}, drop=True))
 
+    # Use a label
+    _process = partial(
+        process_registry["array_element"].implementation,
+        label="B02",
+        data=ParameterReference(from_parameter="data"),
+    )
+
+    output_cube = reduce_dimension(data=input_cube, reducer=_process, dimension="bands")
+
+    general_output_checks(
+        input_cube=input_cube,
+        output_cube=output_cube,
+        verify_attrs=False,
+        verify_crs=True,
+    )
+
+    xr.testing.assert_equal(output_cube, input_cube.loc[{"bands": "B02"}].drop("bands"))
+
     # When the index is out of range, we expect an ArrayElementNotAvailable exception to be thrown
     _process_not_available = partial(
         process_registry["array_element"].implementation,