Ensure common aggregations run before other aggregations

pipeline/workflow/aggregation-helper/aggregation/common.py

@@ -41,3 +41,40 @@ def _escape_sql_literal(val: str) -> str:
     - Single quotes (') are escaped to '' to prevent terminating Spanner string.
     """
     return val.replace('\\', '\\\\\\\\').replace('"', '\\"').replace("'", "''")
+
+
+# Execution priority ranks for calculation types within a stage.
+# Deterministically resolves execution order when multiple calculation types exist in the same stage.
+#
+# Tier 0 (0-2): Graph Topology & Schema Prerequisite Generation
+#   - LINKED_EDGES (0): Computes transitive graph closures (linkedContainedInPlace, linkedMemberOf, linkedMember).
+#     Must run first as place rollups rely on containment graph edges to determine target places.
+#   - STAT_VAR_GROUPS (1): Constructs StatVarGroup nodes and hierarchy edges (specializationOf, memberOf).
+#     Queries specializationOf and curated memberOf edges.
+#   - PROVENANCE_SUMMARY (2): Generates summary statistics for observation tables in KeyValueStore.
+#     Summary over observation data and place type edges.
+#
+# Tier 1 (10-15): Data Rollups, Series, & Derived Calculations
+#   - PLACE_AGGREGATION (10): Primary spatial rollup aggregating raw observations up geographic containment trees.
+#   - STAT_VAR_AGGREGATION (11): Sums component variables into ancestor variables across spatial rollups.
+#   - ENTITY_AGGREGATION (12): Aggregates event occurrences across spatial/temporal bounds.
+#   - STAT_VAR_SERIES_AGGREGATION (13): Aggregates time series across temporal granularities (e.g. Daily -> Monthly).
+#   - STAT_VAR_CALCULATION (14): Computes derived metrics/formulas (e.g. per-capita ratios like Income / Population).
+#     Must run after rollups so numerators and denominators exist across all spatial and variable levels.
+#   - SUPER_ENUM_AGGREGATION (15): Aggregates across enumerated concept categories.
+#
+# Tier 2 (20): Global Post-Processing
+#   - EMBEDDING_GENERATION (20): Computes vector embeddings across finalized node properties.
+CALCULATION_TYPE_PRIORITY = {
+    "LINKED_EDGES": 0,
+    "STAT_VAR_GROUPS": 1,
+    "PROVENANCE_SUMMARY": 2,
+    "PLACE_AGGREGATION": 10,
+    "STAT_VAR_AGGREGATION": 11,
+    "ENTITY_AGGREGATION": 12,
+    "STAT_VAR_SERIES_AGGREGATION": 13,
+    "STAT_VAR_CALCULATION": 14,
+    "SUPER_ENUM_AGGREGATION": 15,
+    "EMBEDDING_GENERATION": 20,
+}
+

pipeline/workflow/aggregation-helper/aggregation/configs/common.yaml

@@ -2,21 +2,21 @@ calculations:
   # Generates linkedContainedInPlace, linkedMemberOf, etc.
   - name: "Global: ContainedInPlace & MemberOf Graph Linked Edges"
     type: LINKED_EDGES
-    stage: 1
+    stage: 0
     input_imports:
       - "*"
 
   # Generates summary statistics in the KeyValueStore table
   - name: "Global: KeyValueStore Provenance & Lineage Summary"
     type: PROVENANCE_SUMMARY
-    stage: 1
+    stage: 0
     input_imports:
       - "*"
 
   # Generates the Statistical Variable hierarchy/verticals
   - name: "Global: StatVar Group Hierarchy & Verticals"
     type: STAT_VAR_GROUPS
-    stage: 1
+    stage: 0
     input_imports:
       - "*"
 

pipeline/workflow/aggregation-helper/aggregation/orchestrator.py

@@ -33,6 +33,7 @@
 from .stat_var_series_aggregator import StatVarSeriesAggregator
 from .entity_aggregation_generator import EntityAggregationGenerator, EntityAggregationConfig
 from .super_enum_aggregation_generator import SuperEnumAggregationGenerator
+from .common import CALCULATION_TYPE_PRIORITY
 from .validator import validate_config
 from .deleter import AggregationDeleter
 
@@ -146,6 +147,15 @@ def __init__(
         else:
             self.calculations = validate_config(target_config, schema_file_path)
 
+        # Deterministically sort calculations by stage and calculation priority tier
+        self.calculations.sort(
+            key=lambda c: (
+                c.get("stage", 1),
+                CALCULATION_TYPE_PRIORITY.get(c.get("type", ""), 99)
+            )
+        )
+
+
     def run(self, active_imports: Optional[List[str]] = None, dry_run: bool = True, skip_deletions: bool = False) -> AggregationRunResult:
         """Executes aggregations independently for each active import.
 

pipeline/workflow/aggregation-helper/aggregation/orchestrator_test.py

@@ -21,9 +21,9 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-sys.path.append(os.path.dirname(os.path.dirname(__file__)))
+from aggregation.common import CALCULATION_TYPE_PRIORITY
+from aggregation.orchestrator import AggregationOrchestrator, CalculationType
 
-from aggregation import AggregationOrchestrator
 
 VALID_CONFIG_YAML = textwrap.dedent("""\
     calculations:
@@ -426,5 +426,83 @@ def test_run_global_calcs_failure(self, mock_embedding_gen, mock_executor_cls):
         self.assertEqual(result.import_results["GLOBAL"].error_message, "Vertex AI quota exceeded")
 
 
+ORDERING_CONFIG_YAML = textwrap.dedent("""\
+    calculations:
+      - type: PLACE_AGGREGATION
+        input_imports: ["TestImport"]
+        output_import: "TestImport_Agg"
+        stage: 1
+        place_aggregation:
+          from_place_types: County
+          to_place_types: State
+
+      - type: PROVENANCE_SUMMARY
+        input_imports: ["*"]
+        stage: 0
+
+      - type: LINKED_EDGES
+        input_imports: ["*"]
+        stage: 0
+
+      - type: STAT_VAR_GROUPS
+        input_imports: ["*"]
+        stage: 0
+""")
+
+
+@patch('aggregation.orchestrator.BigQueryExecutor')
+class TestOrchestratorOrdering(unittest.TestCase):
+    """Tests for stage 0 prerequisite resolution and deterministic priority ordering."""
+
+    def setUp(self):
+        self.tmpdir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.tmpdir.cleanup)
+
+        self.config_path = os.path.join(self.tmpdir.name, "ordering_config.yaml")
+        with open(self.config_path, "w") as f:
+            f.write(ORDERING_CONFIG_YAML)
+
+        self.orchestrator = AggregationOrchestrator(
+            connection_id="conn",
+            project_id="proj",
+            instance_id="inst",
+            database_id="db",
+            config_file_path=self.config_path
+        )
+
+    def test_active_stages_includes_stage_0(self, mock_executor):
+        """Verifies stage 0 is included in active stages for matching imports."""
+        active_stages = self.orchestrator.get_active_stages(["TestImport"])
+        self.assertEqual(active_stages, [0, 1])
+
+    def test_calculation_priority_sorting(self, mock_executor):
+        """Verifies calculations are sorted deterministically by (stage, priority)."""
+        types_in_order = [calc["type"] for calc in self.orchestrator.calculations]
+        expected_order = [
+            "LINKED_EDGES",
+            "STAT_VAR_GROUPS",
+            "PROVENANCE_SUMMARY",
+            "PLACE_AGGREGATION"
+        ]
+        self.assertEqual(types_in_order, expected_order)
+
+
+class TestConfigSanity(unittest.TestCase):
+    """Sanity checks for calculation configurations and metadata."""
+
+    def test_all_calculation_types_have_priority(self):
+        """Ensures all CalculationType enum members have a priority defined in CALCULATION_TYPE_PRIORITY."""
+        calc_type_values = {t.value for t in CalculationType}
+        priority_keys = set(CALCULATION_TYPE_PRIORITY.keys())
+        self.assertSetEqual(
+            calc_type_values,
+            priority_keys,
+            f"Mismatch between CalculationType and CALCULATION_TYPE_PRIORITY. "
+            f"Missing: {calc_type_values ^ priority_keys}"
+        )
+
+
 if __name__ == '__main__':
     unittest.main()
+
+

pipeline/workflow/aggregation-helper/aggregation/schema.json

@@ -45,6 +45,16 @@
           "disabled": { "type": "boolean" }
         },
         "allOf": [
+          {
+            "if": {
+              "not": {
+                "properties": { "type": { "const": "EMBEDDING_GENERATION" } }
+              }
+            },
+            "then": {
+              "required": ["input_imports"]
+            }
+          },
           {
             "if": { "properties": { "type": { "const": "PLACE_AGGREGATION" } } },
             "then": {