STACAPIJobDatabase: faster counting, paging and existence check (#972)

CHANGELOG.md

@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- `STACAPIJobDatabase`: improve performance by using `numberMatched` counts, a larger search page size and a direct collection lookup ([#972](https://github.com/Open-EO/openeo-python-client/issues/972))
+
 ### Removed
 
 ### Fixed

openeo/extra/job_management/_manager.py

@@ -384,7 +384,8 @@ def run_loop():
                 stats["run_jobs loop"] += 1
 
                 # Show current stats and sleep
-                _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}")
+                if _log.isEnabledFor(logging.INFO):
+                    _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}")
                 for _ in range(int(max(1, self.poll_sleep))):
                     time.sleep(1)
                     if self._stop_thread:
@@ -526,7 +527,8 @@ def run_jobs(
             stats["run_jobs loop"] += 1
 
             # Show current stats and sleep
-            _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}")
+            if _log.isEnabledFor(logging.INFO):
+                _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}")
             time.sleep(self.poll_sleep)
             stats["sleep"] += 1
 

openeo/extra/job_management/stac_job_db.py

@@ -3,6 +3,7 @@
 import concurrent.futures
 import datetime
 import logging
+import warnings
 from typing import Iterable, List, Optional, Union
 
 import geopandas as gpd
@@ -17,6 +18,7 @@
 try:
     # Optional dependecy
     import pystac_client
+    import pystac_client.exceptions
 except ImportError:
     pass
 
@@ -58,9 +60,18 @@ def __init__(
         self.geometry_column = geometry_column
         self.base_url = stac_root_url
         self.bulk_size = 500
+        # Page size to request from the STAC API when searching items (servers may cap this to a lower maximum).
+        self.search_page_size = 1000
 
     def exists(self) -> bool:
-        return any(c.id == self.collection_id for c in self.client.get_collections())
+        # Look up the collection directly instead of paging through all collections of the STAC API.
+        try:
+            return self.client.get_collection(self.collection_id) is not None
+        except pystac_client.exceptions.APIError as e:
+            # Older pystac_client versions raise on a missing collection instead of returning None.
+            if getattr(e, "status_code", None) == 404:
+                return False
+            raise
 
     def _normalize_df(self, df: pd.DataFrame) -> pd.DataFrame:
         """
@@ -167,12 +178,32 @@ def count_by_status(self, statuses: Iterable[str] = ()) -> dict:
         if isinstance(statuses, str):
             statuses = {statuses}
         statuses = set(statuses)
+        if statuses:
+            # Ask the STAC API for the number of matching items, instead of downloading all of them.
+            counts = {s: self._count_matched(status_filter=self._status_filter([s])) for s in statuses}
+            if all(c is not None for c in counts.values()):
+                return counts
         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()
 
+    @staticmethod
+    def _status_filter(statuses: Iterable[str]) -> Optional[str]:
+        return " OR ".join([f"\"properties.status\"='{s}'" for s in statuses]) or None
+
+    def _count_matched(self, status_filter: Optional[str]) -> Optional[int]:
+        """
+        Get the number of matching items as reported by the STAC API (``numberMatched`` or ``context.matched``),
+        or ``None`` if the STAC API does not report it.
+        """
+        search = self.client.search(method="GET", collections=[self.collection_id], filter=status_filter, limit=1)
+        with warnings.catch_warnings():
+            # pystac_client warns when the count is not supported, we handle that by falling back.
+            warnings.simplefilter("ignore")
+            return search.matched()
+
     def _search_result_to_df(self, search_result: pystac_client.ItemSearch) -> pd.DataFrame:
         """Build a DataFrame from a STAC ItemSearch result."""
         series = [self.series_from(item) for item in search_result.items()]
@@ -184,12 +215,12 @@ def get_by_status(self, statuses: Iterable[str], max: Optional[int] = None) -> p
             statuses = {statuses}
         statuses = set(statuses)
 
-        status_filter = " OR ".join([f"\"properties.status\"='{s}'" for s in statuses]) if statuses else None
         search_results = self.client.search(
             method="GET",
             collections=[self.collection_id],
-            filter=status_filter,
+            filter=self._status_filter(statuses),
             max_items=max,
+            limit=self.search_page_size,
         )
 
         df = self._search_result_to_df(search_results)
@@ -204,6 +235,7 @@ def get_by_indices(self, indices: Iterable[Union[int, str]]) -> pd.DataFrame:
             method="GET",
             collections=[self.collection_id],
             ids=[str(i) for i in indices],
+            limit=self.search_page_size,
         )
         df = self._search_result_to_df(search_results)
         return df

tests/extra/job_management/test_stac_job_db.py

@@ -27,6 +27,9 @@ def mock_pystac_client():
         MagicMock(id="collection-1"),
         MagicMock(id="collection-2"),
     ]
+    mock_client.get_collection.side_effect = lambda collection_id: (
+        MagicMock(id=collection_id) if collection_id in {"collection-1", "collection-2"} else None
+    )
 
     with patch("pystac_client.Client.open", return_value=mock_client):
         yield mock_client
@@ -340,17 +343,56 @@ def test_count_by_status(self, mock_get_by_status, job_db_exists):
         )
         assert job_db_exists.count_by_status() == {"not_started": 1}
 
+    def test_exists_does_not_list_all_collections(self, job_db_exists):
+        assert job_db_exists.exists() is True
+        job_db_exists.client.get_collection.assert_called_once_with("collection-1")
+        job_db_exists.client.get_collections.assert_not_called()
+
+    def test_count_by_status_uses_number_matched(self, requests_mock):
+        dummy_stac_api = DummyStacApi(root_url="http://stacapi.test", requests_mock=requests_mock)
+        for i, status in enumerate(["not_started", "running", "not_started", "finished"]):
+            dummy_stac_api.predefine_item(
+                collection_id="collection-123",
+                item=_pystac_item(id=f"item-{i}", properties={"status": status}),
+            )
+        job_db = STACAPIJobDatabase(collection_id="collection-123", stac_root_url="http://stacapi.test")
+
+        counts = job_db.count_by_status(["not_started", "running", "queued"])
+
+        assert counts == {"not_started": 2, "running": 1, "queued": 0}
+        searches = [r for r in requests_mock.request_history if r.path == "/search"]
+        # One search per status, each only asking for a single item: no full download of the job database.
+        assert len(searches) == 3
+        assert all(r.qs["limit"] == ["1"] for r in searches)
+
+    def test_count_by_status_falls_back_without_number_matched(self, requests_mock):
+        dummy_stac_api = DummyStacApi(
+            root_url="http://stacapi.test", requests_mock=requests_mock, support_number_matched=False
+        )
+        for i, status in enumerate(["not_started", "running", "not_started"]):
+            dummy_stac_api.predefine_item(
+                collection_id="collection-123",
+                item=_pystac_item(id=f"item-{i}", properties={"status": status}),
+            )
+        job_db = STACAPIJobDatabase(collection_id="collection-123", stac_root_url="http://stacapi.test")
+
+        assert job_db.count_by_status(["not_started", "running"]) == {"not_started": 2, "running": 1}
+
     def test_get_by_status_no_filter(self, job_db_exists):
         job_db_exists.get_by_status(())
 
         job_db_exists.client.search.assert_called_once_with(
-            method="GET", collections=["collection-1"], filter=None, max_items=None
+            method="GET", collections=["collection-1"], filter=None, max_items=None, limit=1000
         )
 
     def test_get_by_status_with_filter(self, job_db_exists):
         job_db_exists.get_by_status(["not_started"])
         job_db_exists.client.search.assert_called_once_with(
-            method="GET", collections=["collection-1"], filter="\"properties.status\"='not_started'", max_items=None
+            method="GET",
+            collections=["collection-1"],
+            filter="\"properties.status\"='not_started'",
+            max_items=None,
+            limit=1000,
         )
 
     def test_get_by_status_result(self, requests_mock):
@@ -491,14 +533,16 @@ def sleep_mock():
 class DummyStacApi:
     """Minimal dummy implementation of a STAC API for testing purposes."""
 
-    def __init__(self, *, root_url: str = "http://stacapi.test", requests_mock):
+    def __init__(self, *, root_url: str = "http://stacapi.test", requests_mock, support_number_matched: bool = True):
         self.root_url = root_url.rstrip("/")
         self._requests_mock = requests_mock
+        self._support_number_matched = support_number_matched
 
         requests_mock.get(f"{self.root_url}/", json=self._get_root())
         self.collections: List[dict] = []
         requests_mock.get(f"{self.root_url}/collections", json=self._get_collections)
         requests_mock.post(f"{self.root_url}/collections", json=self._post_collections)
+        requests_mock.get(re.compile(rf"{re.escape(self.root_url)}/collections/[^/]+$"), json=self._get_collection)
 
         self.items: Dict[str, Dict[str, Any]] = collections.defaultdict(dict)
         requests_mock.post(
@@ -527,6 +571,15 @@ def _get_collections(self, request, context):
         """Handler of `GET /collections` requests."""
         return {"collections": self.collections}
 
+    def _get_collection(self, request, context):
+        """Handler of `GET /collections/{collection_id}` requests."""
+        collection_id = request.path.split("/")[-1]
+        for collection in self.collections:
+            if collection["id"] == collection_id:
+                return collection
+        context.status_code = 404
+        return {"code": "NotFound", "description": f"Collection {collection_id} not found"}
+
     def _post_collections(self, request, context):
         """Handler of `POST /collections` requests."""
         post_data = request.json()
@@ -581,11 +634,14 @@ def _get_search(self, request, context):
             statuses = set(re.findall(r"\"properties\.status\"='(\w+)'", property_filter))
             items = [i for i in items if i.get("properties", {}).get("status") in statuses]
 
-        return {
+        result = {
             "type": "FeatureCollection",
             "features": items,
             "links": [],
         }
+        if self._support_number_matched:
+            result["numberMatched"] = len(items)
+        return result
 
 
 def test_run_jobs_basic(tmp_path, dummy_backend_foo, requests_mock, sleep_mock):