Add new comparable blueprint abstract method

src/brad/planner/beam/query_based_candidate.py

@@ -30,6 +30,7 @@
     compute_single_aurora_table_cost,
     compute_single_table_movement_time_and_cost,
 )
+from brad.planner.workload import Workload
 from brad.planner.workload.query import Query
 from brad.routing.abstract_policy import FullRoutingPolicy
 from brad.routing.cached import CachedLocationPolicy
@@ -659,6 +660,18 @@ def get_aurora_provisioning(self) -> Provisioning:
     def get_redshift_provisioning(self) -> Provisioning:
         return self.redshift_provisioning
 
+    def get_routing_decisions(self) -> npt.NDArray:
+        engine_query = np.array(
+            [
+                (Workload.EngineLatencyIndex[engine], idx)
+                for engine, qidxs in self.query_locations.items()
+                for idx in qidxs
+            ]
+        )
+        si = np.argsort(engine_query[:, 1])
+        sorted_queries = engine_query[si]
+        return sorted_queries[:, 0]
+
     def get_predicted_analytical_latencies(self) -> npt.NDArray:
         relevant = []
         relevant.append(self.scaled_query_latencies[Engine.Aurora])

src/brad/planner/beam/table_based_candidate.py

@@ -30,6 +30,7 @@
     compute_single_aurora_table_cost,
     compute_single_table_movement_time_and_cost,
 )
+from brad.planner.workload import Workload
 from brad.routing.router import Router
 
 logger = logging.getLogger(__name__)
@@ -617,6 +618,18 @@ def get_aurora_provisioning(self) -> Provisioning:
     def get_redshift_provisioning(self) -> Provisioning:
         return self.redshift_provisioning
 
+    def get_routing_decisions(self) -> npt.NDArray:
+        engine_query = np.array(
+            [
+                (Workload.EngineLatencyIndex[engine], idx)
+                for engine, qidxs in self.query_locations.items()
+                for idx in qidxs
+            ]
+        )
+        si = np.argsort(engine_query[:, 1])
+        sorted_queries = engine_query[si]
+        return sorted_queries[:, 0]
+
     def get_predicted_analytical_latencies(self) -> npt.NDArray:
         relevant = []
         relevant.append(self.scaled_query_latencies[Engine.Aurora])

src/brad/planner/compare/blueprint.py

@@ -22,6 +22,15 @@ def get_aurora_provisioning(self) -> Provisioning:
     def get_redshift_provisioning(self) -> Provisioning:
         raise NotImplementedError
 
+    def get_routing_decisions(self) -> npt.NDArray:
+        """
+        The query routing decisions for each analytical query in the workload.
+        Use `Workload.EngineLatencyIndex` to map engines to numeric indices.
+
+        Should return an NDArray of shape (N,).
+        """
+        raise NotImplementedError
+
     # Predicted performance.
 
     def get_predicted_analytical_latencies(self) -> npt.NDArray: