feat(planner): export a planned workload as one physical ASAP DAG

crates/asap-aware-mapping/src/pass/mod.rs

@@ -267,6 +267,31 @@ impl PlanOutput {
         nodes
     }
 
+    /// The workload as one physical ASAP DAG, with shared sub-DAGs exported
+    /// once: operator queries become its `roots` and standalone scalar
+    /// queries its `scalar_roots`, keyed by position in the batch. That
+    /// position is the entry index, because `check_contract` requires dense
+    /// entry indices.
+    ///
+    /// An empty workload is an error
+    /// ([`PhysicalASAPDAGValidationError::NoRoots`](asap_types::ir::physical_export::PhysicalASAPDAGValidationError::NoRoots)).
+    /// A workload of only literal scalars yields a DAG with no nodes: just
+    /// its `scalar_roots`.
+    pub fn execution_timed_dag(
+        &self,
+    ) -> Result<
+        asap_types::ir::physical_export::PhysicalASAPDAG,
+        crate::SummaryMaintenanceTimingError,
+    > {
+        let plans: Vec<_> = self
+            .plans
+            .iter()
+            .chain(&self.scalar_plans)
+            .map(|p| &p.plan)
+            .collect();
+        crate::summary_maintenance_lifecycle::execution_timed_workload_dag(&self.roots(), &plans)
+    }
+
     pub fn len(&self) -> usize {
         self.plans.len() + self.scalar_roots.len()
     }

crates/asap-aware-mapping/src/summary_maintenance_lifecycle.rs

@@ -22,11 +22,11 @@ use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
 use asap_types::ir::physical_export::{
-    compile_physical_asap_dag_with_node_ids, PhysicalASAPDAG, PhysicalASAPDAGValidationError,
-    PhysicalASAPNodeId,
+    compile_physical_asap_dag_with_node_ids, compile_physical_asap_workload_with_node_ids,
+    PhysicalASAPDAG, PhysicalASAPDAGValidationError, PhysicalASAPNodeId,
 };
 use asap_types::ir::timing::{apply_lifecycle_timings, LifecycleAssignment, TimingMemo};
-use asap_types::ir::{ASAPOp, Operator, OperatorNode};
+use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot};
 use asap_types::post_asap::{
     EvaluationSchedule, ExecutionDataStateError, ExecutionTiming, OutputRepresentation,
     ResultGuarantee, SummaryMaintenanceLifecycle, SummaryMaintenanceLifecycleGuarantee,
@@ -247,64 +247,94 @@ impl SummaryMaintenanceLifecyclePlan {
     /// a `SummaryAgg` is one of its inputs. Timings already on the root are
     /// ignored.
     pub fn execution_timed_dag(&self) -> Result<PhysicalASAPDAG, SummaryMaintenanceTimingError> {
-        let mut memo = TimingMemo::new();
-        let timed = apply_lifecycle_timings(
-            &self.root,
-            &LifecycleAssignment::default_maintained(),
-            &mut memo,
-        )?;
-        let compiled = compile_physical_asap_dag_with_node_ids(&timed)?;
-        let dag = compiled.dag;
-        for population in &standalone_populations(&self.root) {
-            let id = compiled
-                .node_ids
-                .node_id(memo.timed(population).expect("population was timed"))
-                .expect("collected population belongs to the compiled DAG");
-            if !self
-                .deployments
-                .iter()
-                .any(|deployment| deployment.post_asap_node_id == id)
-            {
+        execution_timed_workload_dag(&[QueryRoot::Operator(Rc::clone(&self.root))], &[self])
+    }
+}
+
+/// One physical ASAP DAG for a workload: `roots` are the queries in
+/// workload order, operator roots becoming `roots` and standalone scalars
+/// `scalar_roots` of the export, with sub-DAGs shared between queries
+/// exported once. `plans` are the lifecycle plans of every planned sub-DAG
+/// `roots` reach; timing follows the selected lifecycles of their
+/// deployments, as in [`SummaryMaintenanceLifecyclePlan::execution_timed_dag`].
+pub(crate) fn execution_timed_workload_dag(
+    roots: &[QueryRoot],
+    plans: &[&SummaryMaintenanceLifecyclePlan],
+) -> Result<PhysicalASAPDAG, SummaryMaintenanceTimingError> {
+    // One memo, so a node shared by several roots is timed and exported once.
+    let mut memo = TimingMemo::new();
+    let assignment = LifecycleAssignment::default_maintained();
+    let mut timed = Vec::with_capacity(roots.len());
+    for root in roots {
+        timed.push(match root {
+            QueryRoot::Operator(node) => {
+                QueryRoot::Operator(apply_lifecycle_timings(node, &assignment, &mut memo)?)
+            }
+            QueryRoot::Scalar(expr) => {
+                for node in expr.operator_refs() {
+                    apply_lifecycle_timings(node, &assignment, &mut memo)?;
+                }
+                QueryRoot::Scalar(expr.map_operator_refs(&mut |node| {
+                    Rc::clone(memo.timed(node).expect("scalar reference was timed"))
+                }))
+            }
+        });
+    }
+    let compiled = compile_physical_asap_workload_with_node_ids(&timed)?;
+    let id_of = |node: &Rc<OperatorNode>| {
+        compiled
+            .node_ids
+            .node_id(memo.timed(node).expect("plan node was timed"))
+            .expect("timed plan node belongs to the compiled DAG")
+    };
+    let deployments: Vec<_> = plans
+        .iter()
+        .flat_map(|plan| &plan.deployments)
+        .map(|deployment| (id_of(&deployment.summary), deployment))
+        .collect();
+    for plan in plans {
+        for population in &standalone_populations(&plan.root) {
+            let id = id_of(population);
+            if !deployments.iter().any(|(deployed, _)| *deployed == id) {
                 return Err(SummaryMaintenanceTimingError::UnplannedMaintainedState(id));
             }
         }
-        let mut pending = Vec::new();
-        for deployment in &self.deployments {
-            let guarantee = deployment
-                .summary_maintenance_lifecycle_guarantee
-                .as_ref()
-                .ok_or(SummaryMaintenanceTimingError::UnselectedLifecycle(
-                    deployment.post_asap_node_id,
-                ))?;
-            if guarantee.summary_maintenance_lifecycle != SummaryMaintenanceLifecycle::Ephemeral {
-                pending.push(deployment.post_asap_node_id);
-            }
+    }
+    let dag = compiled.dag;
+    let mut pending = Vec::new();
+    for (id, deployment) in &deployments {
+        let guarantee = deployment
+            .summary_maintenance_lifecycle_guarantee
+            .as_ref()
+            .ok_or(SummaryMaintenanceTimingError::UnselectedLifecycle(*id))?;
+        if guarantee.summary_maintenance_lifecycle != SummaryMaintenanceLifecycle::Ephemeral {
+            pending.push(*id);
         }
-        let mut ingestion = HashSet::new();
-        while let Some(id) = pending.pop() {
-            if ingestion.insert(id) {
-                pending.extend(
-                    dag.edges
-                        .iter()
-                        .filter(|edge| edge.consumer == id)
-                        .map(|edge| edge.producer),
-                );
-            }
+    }
+    let mut ingestion = HashSet::new();
+    while let Some(id) = pending.pop() {
+        if ingestion.insert(id) {
+            pending.extend(
+                dag.edges
+                    .iter()
+                    .filter(|edge| edge.consumer == id)
+                    .map(|edge| edge.producer),
+            );
         }
-        let phases = dag
-            .nodes
-            .iter()
-            .map(|node| {
-                let timing = if ingestion.contains(&node.id) {
-                    ExecutionTiming::IngestionTime
-                } else {
-                    ExecutionTiming::QueryTime
-                };
-                (node.id, timing)
-            })
-            .collect();
-        Ok(dag.with_execution_phases(&phases)?)
     }
+    let phases = dag
+        .nodes
+        .iter()
+        .map(|node| {
+            let timing = if ingestion.contains(&node.id) {
+                ExecutionTiming::IngestionTime
+            } else {
+                ExecutionTiming::QueryTime
+            };
+            (node.id, timing)
+        })
+        .collect();
+    Ok(dag.with_execution_phases(&phases)?)
 }
 
 /// Explicit association between a materialized target and the normalized

crates/integration-tests/tests/operator_design_examples.rs

@@ -398,4 +398,22 @@ async fn batch_planning_replaces_and_shares_summary_operators() {
             rows
         );
     }
+    // The batch exports as one physical DAG: a root per query and the shared
+    // SUM state once.
+    let workload_dag = output.execution_timed_dag().unwrap();
+    assert_eq!(workload_dag.roots.len(), 2);
+    assert_ne!(workload_dag.roots[0], workload_dag.roots[1]);
+    assert_eq!(
+        workload_dag
+            .nodes
+            .iter()
+            .filter(|n| matches!(
+                n.payload,
+                asap_types::ir::physical_export::PhysicalASAPOperatorPayload::ASAP(
+                    asap_types::ir::ASAPOp::SummaryAgg { .. }
+                )
+            ))
+            .count(),
+        1
+    );
 }