feat(core): support fetching entities and entity identifiers for multiple operations at once (#1228)

* feat(core): support fetching entities and entity identifiers for multiple operations at once

Signed-off-by: Alessandro Pomponio <alessandro.pomponio1@ibm.com>

* refactor(core): rename method

Signed-off-by: Alessandro Pomponio <alessandro.pomponio1@ibm.com>

---------

Signed-off-by: Alessandro Pomponio <alessandro.pomponio1@ibm.com>

Author: AlessandroPomponio

ado/core/discoveryspace/space.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import typing
+import warnings
 from collections.abc import Callable, Iterator
 from functools import wraps
 from typing import Any
@@ -623,28 +624,7 @@ def sampledEntities(self) -> list[Entity]:
         if not operation_ids:
             return []
 
-        # Optimize for single operation: use direct query (1 query instead of 2)
-        if len(operation_ids) == 1:
-            sampled_entities = self.sample_store.entities_in_operation(
-                operation_id=next(iter(operation_ids))
-            )
-        else:
-            # Multiple operations: get entity IDs first, then fetch entities
-            # This approach handles deduplication across operations naturally
-            sampled_entity_ids = set()
-            for operationid in operation_ids:
-                sampled_entity_ids.update(
-                    self.entity_identifiers_in_operation(operation_id=operationid)
-                )
-
-            if not sampled_entity_ids:
-                return []
-
-            # Efficiently fetch only the entities that were sampled in operations
-            # This avoids loading all entities from the store when we only need a subset
-            sampled_entities = self.sample_store.entities_with_identifiers(
-                sampled_entity_ids
-            )
+        sampled_entities = self.sample_store.entities_in_operations(operation_ids)
 
         # TODO: Consider removing isEntitySpace check
         # The additional check of isEntityInSpace should not be required if things are working correctly
@@ -1072,10 +1052,25 @@ def complete_measurement_request_with_results_timeseries(
         )
 
     @_perform_preflight_checks_for_sample_store_methods
-    def entity_identifiers_in_operation(self, operation_id: str) -> set[str]:
-        return self.sample_store.entity_identifiers_in_operation(
-            operation_id=operation_id
+    def entity_identifiers_in_operations(
+        self, operation_ids: str | set[str]
+    ) -> set[str]:
+        """Return entity identifiers sampled in the given operation(s)."""
+        return self.sample_store.entity_identifiers_in_operations(
+            operation_ids=operation_ids
+        )
+
+    @_perform_preflight_checks_for_sample_store_methods
+    def entity_identifiers_in_operation(
+        self, operation_ids: str | set[str]
+    ) -> set[str]:
+        """Deprecated: use entity_identifiers_in_operations instead."""
+        warnings.warn(
+            "entity_identifiers_in_operation is deprecated, use entity_identifiers_in_operations instead.",
+            DeprecationWarning,
+            stacklevel=2,
         )
+        return self.entity_identifiers_in_operations(operation_ids)
 
     @_perform_preflight_checks_for_sample_store_methods
     def experiments_in_operation(self, operation_id: str) -> list[Experiment]:

ado/core/samplestore/base.py

@@ -3,6 +3,7 @@
 
 import abc
 import typing
+import warnings
 from abc import ABC
 from typing import Annotated, Literal
 
@@ -491,8 +492,25 @@ def commit(self) -> None:  # pragma: nocover
         """Commits all the changes to the source"""
 
     @abc.abstractmethod
-    def entities_in_operation(self, operation_id: str) -> list[Entity]:
-        """Returns list of entities in the given operation."""
+    def entities_in_operations(self, operation_ids: str | set[str]) -> list[Entity]:
+        """Returns list of entities in the given operation(s).
+
+        Args:
+            operation_ids: A single operation identifier or a set of operation
+                identifiers to fetch entities for.
+
+        Returns:
+            List of Entity objects that were sampled in the specified operation(s).
+        """
+
+    def entities_in_operation(self, operation_ids: str | set[str]) -> list[Entity]:
+        """Deprecated: use entities_in_operations instead."""
+        warnings.warn(
+            "entities_in_operation is deprecated, use entities_in_operations instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.entities_in_operations(operation_ids)
 
     @abc.abstractmethod
     def operation_entity_statistics(self, operation_id: str) -> dict[str, int]:

ado/core/samplestore/sql.py

@@ -6,6 +6,7 @@
 import logging
 import typing
 import uuid
+import warnings
 from typing import TYPE_CHECKING, Annotated, Literal
 
 import pydantic
@@ -822,19 +823,23 @@ def entities_with_identifiers(
 
         return list(entities_dict.values())
 
-    def entities_in_operation(self, operation_id: str) -> list[Entity]:
-        """Get entities directly from a single operation in one query.
+    def entities_in_operations(self, operation_ids: str | set[str]) -> list[Entity]:
+        """Get entities directly from one or more operations in one query.
 
-        This method is optimized for the common case of fetching entities from
-        a single operation. It performs the entire operation in a single database
-        query, avoiding the need to first fetch entity IDs and then fetch entities.
+        This method fetches entities from one or more operations in a single
+        database query, avoiding the need to first fetch entity IDs and then
+        fetch entities.
 
         Args:
-            operation_id: The operation identifier to fetch entities for
+            operation_ids: A single operation identifier or a set of operation
+                identifiers to fetch entities for.
 
         Returns:
-            List of Entity objects that were sampled in the specified operation
+            List of Entity objects that were sampled in the specified operation(s)
         """
+        if isinstance(operation_ids, str):
+            operation_ids = {operation_ids}
+
         query = sqlalchemy.text(f"""
             SELECT
                 ent.identifier,
@@ -844,16 +849,18 @@ def entities_in_operation(self, operation_id: str) -> list[Entity]:
             JOIN {self._tablename}_measurement_results res ON res.entity_id = ent.identifier
             JOIN {self._tablename}_measurement_requests_results reqres ON reqres.result_uid = res.uid
             JOIN {self._tablename}_measurement_requests req ON reqres.request_uid = req.uid
-            WHERE req.operation_id = :operation_id
+            WHERE req.operation_id IN :operation_ids
         """).bindparams(  # noqa: S608 - self._tablename is not untrusted
-            operation_id=operation_id
+            sqlalchemy.bindparam(
+                "operation_ids", value=list(operation_ids), expanding=True
+            )
         )
 
         try:
             with self.engine.begin() as connectable:
                 cur = connectable.execute(query)
         except SQLAlchemyError as error:
-            msg = f"Unable to fetch entities for operation {operation_id} from sample store {self._tablename}"
+            msg = f"Unable to fetch entities for operations {operation_ids} from sample store {self._tablename}"
             self.log.critical(f"{msg}. Error: {error}")
             raise SystemError(f"{msg}. Error: {error}") from error
 
@@ -903,6 +910,15 @@ def entities_in_operation(self, operation_id: str) -> list[Entity]:
 
         return list(entities_dict.values())
 
+    def entities_in_operation(self, operation_ids: str | set[str]) -> list[Entity]:
+        """Deprecated: use entities_in_operations instead."""
+        warnings.warn(
+            "entities_in_operation is deprecated, use entities_in_operations instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.entities_in_operations(operation_ids)
+
     @property
     def numberOfEntities(self) -> int:
 
@@ -2122,29 +2138,56 @@ def experiments_in_operation(self, operation_id: str) -> list[Experiment]:
             for e in cur
         ]
 
-    def entity_identifiers_in_operation(self, operation_id: str) -> set[str]:
+    def entity_identifiers_in_operations(
+        self, operation_ids: str | set[str]
+    ) -> set[str]:
+        """Get the set of entity identifiers sampled in one or more operations.
+
+        Args:
+            operation_ids: A single operation identifier or a set of operation
+                identifiers to look up entity identifiers for.
+
+        Returns:
+            Set of entity identifier strings across all specified operations.
+        """
+        if isinstance(operation_ids, str):
+            operation_ids = {operation_ids}
+
         try:
             with self.engine.begin() as connectable:
                 query = sqlalchemy.text(f"""
                     SELECT DISTINCT(res.entity_id)
                     FROM (
                         SELECT *
                         FROM {self._tablename}_measurement_requests
-                        WHERE operation_id = :operation_id
+                        WHERE operation_id IN :operation_ids
                     ) req
                     JOIN {self._tablename}_measurement_requests_results reqres ON reqres.request_uid = req.uid
                     JOIN {self._tablename}_measurement_results res ON reqres.result_uid = res.uid
                     """).bindparams(  # noqa: S608 - self._tablename is not untrusted
-                    operation_id=operation_id
+                    sqlalchemy.bindparam(
+                        "operation_ids", value=list(operation_ids), expanding=True
+                    )
                 )
                 cur = connectable.execute(query)
         except SQLAlchemyError as error:
-            msg = f"Unable to get the entity ids for operation {operation_id}"
+            msg = f"Unable to get the entity ids for operations {operation_ids}"
             self.log.critical(f"{msg}. Error: {error}")
             raise SystemError(f"{msg}. Error: {error}") from error
 
         return {ident[0] for ident in cur}
 
+    def entity_identifiers_in_operation(
+        self, operation_ids: str | set[str]
+    ) -> set[str]:
+        """Deprecated: use entity_identifiers_in_operations instead."""
+        warnings.warn(
+            "entity_identifiers_in_operation is deprecated, use entity_identifiers_in_operations instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.entity_identifiers_in_operations(operation_ids)
+
     def complete_measurement_request_with_results_timeseries(
         self,
         operation_id: str,