Reuse existing aggregation when resolving summary vector blob id

libs/services/src/ri_cloud_services/sumo_access/summary_access.py

@@ -6,7 +6,9 @@
 
 from __future__ import annotations
 
-from fmu.sumo.explorer.objects import Table
+import asyncio
+
+from fmu.sumo.explorer.objects import SearchContext, Table
 
 from ri_cloud_services.service_exceptions import (
     InvalidDataError,
@@ -21,6 +23,10 @@
 # out when listing available vectors.
 _SUMMARY_METADATA_COLUMNS = {"DATE", "REAL", "ENSEMBLE", "ITER"}
 
+# Aggregation operation used for summary vectors. Must be the same when looking for an existing
+# aggregation and when triggering one, or the lookup never matches what was produced.
+_AGGREGATION_OPERATION = "collection"
+
 
 class SummaryAccess:
     """Access summary (timeseries) data for a given Sumo case + ensemble."""
@@ -71,8 +77,7 @@ async def get_available_vectors_async(self) -> list[str]:
     async def get_vector_blob_id_async(self, vector_name: str) -> str:
         """Get the blob ID for the given summary vector.
 
-        The temporary solution is not optimized, so we trigger aggregation to ensure the blob ID is available,
-        this triggers an aggregation
+        Aggregation is triggered on Sumo if no usable aggregation exists yet.
 
         Returns the raw Azure blob ID. The caller should authenticate using
         OAuth Bearer token (same token used for Sumo API access).
@@ -86,21 +91,41 @@ async def get_vector_blob_id_async(self, vector_name: str) -> str:
     async def _get_vector_agg_table(self, vector_name: str) -> Table:
         """Get the aggregated table for the given summary vector.
 
-        The temporary solution is not optimized, so we trigger aggregation to ensure the aggregated table is available,
-        this triggers an aggregation
+        Reuses an existing aggregation when there is one, and falls back to triggering an
+        aggregation on Sumo. Triggering costs several seconds, so the fast path matters.
 
         Returns the aggregated table object. The caller should authenticate using
         OAuth Bearer token (same token used for Sumo API access).
         """
         case = get_case_by_uuid(self._access_token, self._case_uuid)
 
-        sc_per_real_tables = case.tables.filter(
-            ensemble=self._ensemble_name,
+        sc_tables_basis = case.tables.filter(
             column=vector_name,
+            ensemble=self._ensemble_name,
             standard_result="simulationtimeseries",  # TODO: Use standard_result type from fmu-data-io?
-            realization=True,
         )
 
+        # Look for an existing aggregation. Note that this filter must not carry realization=True:
+        # an object cannot be both a realization and an aggregation, so such a filter never matches.
+        # SearchContext.aggregation_async() probes on the context it is called on, which is why
+        # calling it on the per-realization context below always ends up re-triggering aggregation.
+        sc_existing_agg_tables = sc_tables_basis.filter(aggregation=_AGGREGATION_OPERATION)
+        existing_agg_table_count = await sc_existing_agg_tables.length_async()
+        if existing_agg_table_count > 1:
+            raise MultipleDataMatchesError(
+                f"Multiple existing aggregation tables found for vector '{vector_name}' in "
+                f"case='{self._case_uuid}', ensemble='{self._ensemble_name}'",
+                Service.SUMO,
+            )
+        if existing_agg_table_count == 1:
+            existing_agg_table = await sc_existing_agg_tables.single_async
+            if isinstance(existing_agg_table, Table) and await self._is_agg_valid_for_reals_async(
+                existing_agg_table, sc_tables_basis
+            ):
+                return existing_agg_table
+
+        sc_per_real_tables = sc_tables_basis.filter(realization=True)
+
         table_names = await sc_per_real_tables.names_async
         num_tables = len(table_names)
         if num_tables == 0:
@@ -115,7 +140,7 @@ async def _get_vector_agg_table(self, vector_name: str) -> Table:
             )
 
         # Trigger aggregation if not existing
-        agg_table = await sc_per_real_tables.aggregation_async(operation="collection", column=vector_name)
+        agg_table = await sc_per_real_tables.aggregation_async(operation=_AGGREGATION_OPERATION, column=vector_name)
 
         if not isinstance(agg_table, Table):
             raise InvalidDataError(
@@ -124,3 +149,37 @@ async def _get_vector_agg_table(self, vector_name: str) -> Table:
             )
 
         return agg_table
+
+    @staticmethod
+    async def _is_agg_valid_for_reals_async(agg_table: Table, sc_tables_basis: SearchContext) -> bool:
+        """Tell whether an existing aggregation still covers all realizations.
+
+        Realizations can be added after an aggregation was made, which leaves the aggregation
+        holding a subset of the data. This is the same check SearchContext.aggregation_async()
+        applies before reusing an aggregation: the realizations that existed when the aggregation
+        was made must be exactly the ones it recorded, and no realization may have been added since.
+        """
+        try:
+            recorded_realization_ids = agg_table.metadata["fmu"]["aggregation"]["realization_ids"]
+            aggregation_timestamp = agg_table.metadata["_sumo"]["timestamp"]
+        except KeyError:
+            return False
+
+        sc_real_tables = sc_tables_basis.filter(realization=True)
+
+        # Neither query depends on the other's result, so run them concurrently rather than
+        # paying for two sequential round-trips on every call.
+        async with asyncio.TaskGroup() as tg:
+            older_ids_task = tg.create_task(
+                sc_real_tables.filter(
+                    complex={"range": {"_sumo.timestamp": {"lt": aggregation_timestamp}}}
+                ).realizationids_async
+            )
+            current_count_task = tg.create_task(sc_real_tables.length_async())
+
+        realization_ids = older_ids_task.result()
+
+        if set(realization_ids) != set(recorded_realization_ids):
+            return False
+
+        return len(realization_ids) == current_count_task.result()