STACAPIJobDatabase.get_by_status() sorts df on index #707 (#734)

* STACAPIJobDatabase.get_by_status() sorts df on index #707

* STACAPIJobDatabase: item ID is no longer derived from pandas.DataFrame index, but rather from item_id column; which is created if it doesn't exist #707

* created STACAPIJobDatabase._normalize_df() method #707

* update changelog for #707

Author: VincentVerelst

CHANGELOG.md

@@ -14,7 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Removed
 
 ### Fixed
-
+- `STACAPIJobDatabase.get_by_status()` now always returns a `pandas.DataFrame` with an index compatible with `MultiBackendJobManager`. ([#707](https://github.com/Open-EO/openeo-python-client/issues/707))
 
 ## [0.39.1] - 2025-02-26
 

openeo/extra/job_management/stac_job_db.py

@@ -1,7 +1,7 @@
 import concurrent.futures
 import datetime
 import logging
-from typing import Iterable, List
+from typing import Iterable, List, Optional
 
 import geopandas as gpd
 import numpy as np
@@ -53,6 +53,17 @@ def __init__(
 
     def exists(self) -> bool:
         return any(c.id == self.collection_id for c in self.client.get_collections())
+    
+    def _normalize_df(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Normalize the given dataframe to be compatible with :py:class:`MultiBackendJobManager`
+        by adding the default columns and setting the index.
+        """
+        df = MultiBackendJobManager._normalize_df(df)
+        # If the user doesn't specify the item_id column, we will use the index.
+        if "item_id" not in df.columns:
+            df = df.reset_index(names=["item_id"])
+        return df
 
     def initialize_from_df(self, df: pd.DataFrame, *, on_exists: str = "error"):
         """
@@ -83,15 +94,15 @@ def initialize_from_df(self, df: pd.DataFrame, *, on_exists: str = "error"):
                 raise FileExistsError(f"Job database {self!r} already exists.")
             elif on_exists == "append":
                 existing_df = self.get_by_status([])
-                df = MultiBackendJobManager._normalize_df(df)
+                df = self._normalize_df(df)
                 df = pd.concat([existing_df, df], ignore_index=True).replace({np.nan: None})
                 self.persist(df)
                 return self
 
             else:
                 raise ValueError(f"Invalid on_exists={on_exists!r}")
 
-        df = MultiBackendJobManager._normalize_df(df)
+        df = self._normalize_df(df)
         self.persist(df)
         # Return self to allow chaining with constructor.
         return self
@@ -105,7 +116,6 @@ def series_from(self, item: pystac.Item) -> pd.Series:
         """
         item_dict = item.to_dict()
         item_id = item_dict["id"]
-        dt = item_dict["properties"]["datetime"]
 
         return pd.Series(item_dict["properties"], name=item_id)
 
@@ -118,6 +128,7 @@ def item_from(self, series: pd.Series) -> pystac.Item:
         :return: pystac.Item
         """
         series_dict = series.to_dict()
+        item_id = series_dict.pop("item_id")
         item_dict = {}
         item_dict.setdefault("stac_version", pystac.get_stac_version())
         item_dict.setdefault("type", "Feature")
@@ -139,7 +150,7 @@ def item_from(self, series: pd.Series) -> pystac.Item:
             item_dict["geometry"] = None
 
         # from_dict handles associating any Links and Assets with the Item
-        item_dict["id"] = series.name
+        item_dict["id"] = item_id
         item = pystac.Item.from_dict(item_dict)
         if self.has_geometry:
             item.bbox = shape(series[self.geometry_column]).bounds
@@ -151,13 +162,13 @@ def count_by_status(self, statuses: Iterable[str] = ()) -> dict:
         if isinstance(statuses, str):
             statuses = {statuses}
         statuses = set(statuses)
-        items = self.get_by_status(statuses, max=200)
+        items = self.get_by_status(statuses)
         if items is None:
             return {k: 0 for k in statuses}
         else:
             return items["status"].value_counts().to_dict()
 
-    def get_by_status(self, statuses: Iterable[str], max=None) -> pd.DataFrame:
+    def get_by_status(self, statuses: Iterable[str], max: Optional[int] = None) -> pd.DataFrame:
         if isinstance(statuses, str):
             statuses = {statuses}
         statuses = set(statuses)
@@ -172,10 +183,10 @@ def get_by_status(self, statuses: Iterable[str], max=None) -> pd.DataFrame:
 
         series = [self.series_from(item) for item in search_results.items()]
 
-        df = pd.DataFrame(series)
+        df = pd.DataFrame(series).reset_index(names=["item_id"])
         if len(series) == 0:
             # TODO: What if default columns are overwritten by the user?
-            df = MultiBackendJobManager._normalize_df(
+            df = self._normalize_df(
                 df
             )  # Even for an empty dataframe the default columns are required
         return df

tests/extra/job_management/test_stac_job_db.py

@@ -69,6 +69,7 @@ def dummy_dataframe() -> pd.DataFrame:
 def normalized_dummy_dataframe() -> pd.DataFrame:
     return pd.DataFrame(
         {
+            "item_id": [0],
             "no": [1],
             "geometry": [2],
             "here": [3],
@@ -87,13 +88,14 @@ def normalized_dummy_dataframe() -> pd.DataFrame:
 
 @pytest.fixture
 def another_dummy_dataframe() -> pd.DataFrame:
-    return pd.DataFrame({"no": [4], "geometry": [5], "here": [6]})
+    return pd.DataFrame({"item_id": [1], "no": [4], "geometry": [5], "here": [6]})
 
 
 @pytest.fixture
 def normalized_merged_dummy_dataframe() -> pd.DataFrame:
     return pd.DataFrame(
         {
+            "item_id": [0, 1],
             "no": [1, 4],
             "geometry": [2, 5],
             "here": [3, 6],
@@ -126,6 +128,7 @@ def dummy_geodataframe() -> gpd.GeoDataFrame:
 def normalized_dummy_geodataframe() -> pd.DataFrame:
     return pd.DataFrame(
         {
+            "item_id": [0],
             "there": [1],
             "is": [2],
             "geometry": [{"type": "Point", "coordinates": (1.0, 1.0)}],
@@ -174,13 +177,20 @@ def dummy_stac_item_geometry() -> pystac.Item:
 
 @pytest.fixture
 def dummy_series() -> pd.Series:
-    return pd.Series({"datetime": pystac.utils.datetime_to_str(FAKE_NOW), "some_property": "value"}, name="test")
+    return pd.Series(
+        {"item_id": "test", "datetime": pystac.utils.datetime_to_str(FAKE_NOW), "some_property": "value"}, name="test"
+    )
 
 
+@pytest.fixture
+def dummy_series_no_item_id() -> pd.Series:
+    return pd.Series({"datetime": pystac.utils.datetime_to_str(FAKE_NOW), "some_property": "value"}, name="test")
+
 @pytest.fixture
 def dummy_series_geometry() -> pd.Series:
     return pd.Series(
         {
+            "item_id": "test",
             "datetime": pystac.utils.datetime_to_str(FAKE_NOW),
             "some_property": "value",
             "geometry": {
@@ -203,6 +213,7 @@ def patch_datetime_now():
 def bulk_dataframe():
     return pd.DataFrame(
         {
+            "item_id": [f"test-{i}" for i in range(10)],
             "some_property": [f"value-{i}" for i in range(10)],
         },
         index=[i for i in range(10)],
@@ -259,8 +270,8 @@ def test_initialize_from_df_with_geometry(
         assert job_db_not_exists.has_geometry == True
         assert job_db_not_exists.geometry_column == "geometry"
 
-    def test_series_from(self, job_db_exists, dummy_series, dummy_stac_item):
-        pdt.assert_series_equal(job_db_exists.series_from(dummy_stac_item), dummy_series)
+    def test_series_from(self, job_db_exists, dummy_series_no_item_id, dummy_stac_item):
+        pdt.assert_series_equal(job_db_exists.series_from(dummy_stac_item), dummy_series_no_item_id)
 
     def test_item_from(self, patch_datetime_now, job_db_exists, dummy_series, dummy_stac_item):
         item = job_db_exists.item_from(dummy_series)
@@ -298,10 +309,11 @@ def test_get_by_status_result(self, job_db_exists):
             df,
             pd.DataFrame(
                 {
+                    "item_id": ["test"],
                     "datetime": [pystac.utils.datetime_to_str(FAKE_NOW)],
                     "some_property": ["value"],
                 },
-                index=["test"],
+                index=[0],
             ),
         )
 
@@ -326,7 +338,6 @@ def handle_row(series):
         mock_requests_post.reason = "OK"
 
         job_db_exists.persist(bulk_dataframe)
-        # job_db_exists._upload_items_bulk(collection_id=job_db_exists.collection_id, items=items)
 
         mock_requests_post.assert_called_once()