Start maintenance DAG execution at supplied materializations

data_plane/src/precompute_engine/maintenance_runtime.rs

@@ -28,19 +28,20 @@ struct OperatorAdapter<'a> {
 impl PrecomputeOperatorRegistry<SummaryState> for OperatorAdapter<'_> {
     type Error = String;
 
+    fn materialized_input(&self, node: &ExecutableDagNode) -> Result<Option<SummaryState>, String> {
+        Ok(matches!(
+            self.binding.node(node.id),
+            Some(BackendNodeBinding::Materialization { summary_definition })
+                if *summary_definition == self.source_definition
+        )
+        .then(|| Arc::clone(&self.source)))
+    }
+
     fn execute(
         &self,
         node: &ExecutableDagNode,
         inputs: &[Arc<SummaryState>],
     ) -> Result<SummaryState, Self::Error> {
-        let is_source = matches!(
-            self.binding.node(node.id),
-            Some(BackendNodeBinding::Materialization { summary_definition })
-                if *summary_definition == self.source_definition
-        );
-        if is_source && inputs.is_empty() {
-            return Ok(Arc::clone(&self.source));
-        }
         match &node.payload {
             ExecutableOperatorPayload::SummaryMerge => merge_inputs(inputs),
             ExecutableOperatorPayload::SummaryAgg { .. } => Err(
@@ -326,7 +327,7 @@ impl MaintenanceDagSink {
                 if !depends_on_any(&dag, *sink_node, &source_nodes) {
                     continue;
                 }
-                let reachable = dependencies(&dag, *sink_node);
+                let reachable = dependencies_until(&dag, *sink_node, &source_nodes);
                 let foreign_source = reachable.iter().any(|node| {
                     let has_input = dag.edges.iter().any(|edge| edge.consumer == *node);
                     !has_input
@@ -344,6 +345,7 @@ impl MaintenanceDagSink {
                 }
                 if dag.edges.iter().any(|edge| {
                     reachable.contains(&edge.consumer)
+                        && !source_nodes.contains(&edge.consumer)
                         && !matches!(
                             edge.grouping,
                             planner_types::post_asap::GroupingEdgeCompatibility::Identical
@@ -430,11 +432,19 @@ fn depends_on_any(
 fn dependencies(
     dag: &planner_types::post_asap::ExecutableDag,
     sink: PostAsapNodeId,
+) -> BTreeSet<PostAsapNodeId> {
+    dependencies_until(dag, sink, &BTreeSet::new())
+}
+
+fn dependencies_until(
+    dag: &planner_types::post_asap::ExecutableDag,
+    sink: PostAsapNodeId,
+    frontier: &BTreeSet<PostAsapNodeId>,
 ) -> BTreeSet<PostAsapNodeId> {
     let mut pending = vec![sink];
     let mut seen = BTreeSet::new();
     while let Some(node) = pending.pop() {
-        if seen.insert(node) {
+        if seen.insert(node) && !frontier.contains(&node) {
             pending.extend(
                 dag.edges
                     .iter()
@@ -696,7 +706,7 @@ mod tests {
         use crate::storage_engines::types::{
             ActivePhysicalPlan, HotReloadActivePhysicalPlan, StreamingConfig,
         };
-        use control_plane::physical::executable_binding::{InstalledPostAsapDag, OwnedPostAsapDag};
+        use asap_types::executable_plan::{InstalledPostAsapDag, OwnedPostAsapDag};
         use std::sync::atomic::{AtomicUsize, Ordering};
 
         #[derive(Default)]

data_plane/src/precompute_engine/subdag_scheduler.rs

@@ -21,6 +21,11 @@ pub struct MaterializationCommitKey {
 
 pub trait PrecomputeOperatorRegistry<V> {
     type Error;
+    /// A materialized input is an execution frontier: its absorbed semantic
+    /// dependencies have already run and must not be evaluated again.
+    fn materialized_input(&self, _node: &ExecutableDagNode) -> Result<Option<V>, Self::Error> {
+        Ok(None)
+    }
     fn execute(&self, node: &ExecutableDagNode, inputs: &[Arc<V>]) -> Result<V, Self::Error>;
 }
 
@@ -120,6 +125,14 @@ where
                 "query-time node {id} in precompute dependency path"
             )));
         }
+        if let Some(value) = registry
+            .materialized_input(node)
+            .map_err(ScheduleError::Operator)?
+        {
+            values.insert(id, Arc::new(value));
+            active.remove(&id);
+            return Ok(());
+        }
         let child_ids = inputs.get(&id).cloned().unwrap_or_default();
         for child in &child_ids {
             visit(*child, nodes, inputs, active, values, registry)?;
@@ -301,6 +314,52 @@ mod tests {
         assert_eq!(registry.0.lock().unwrap().values().sum::<usize>(), 4);
     }
 
+    #[test]
+    fn supplied_materialization_cuts_absorbed_dependencies_and_is_shared() {
+        struct FrontierRegistry(Registry);
+        impl PrecomputeOperatorRegistry<u32> for FrontierRegistry {
+            type Error = String;
+            fn materialized_input(&self, node: &ExecutableDagNode) -> Result<Option<u32>, String> {
+                Ok((node.id == PostAsapNodeId(1)).then_some(10))
+            }
+            fn execute(
+                &self,
+                node: &ExecutableDagNode,
+                inputs: &[Arc<u32>],
+            ) -> Result<u32, String> {
+                assert_ne!(
+                    node.id,
+                    PostAsapNodeId(0),
+                    "absorbed source subtree must not execute"
+                );
+                self.0.execute(node, inputs)
+            }
+        }
+        let mut raw = node(0);
+        raw.output_state = ExecutionDataState::READ_ROWS;
+        let mut query = node(4);
+        query.output_state = ExecutionDataState::READ_ROWS;
+        let dag = ExecutableDag {
+            nodes: vec![raw, node(1), node(2), node(3), query],
+            edges: vec![edge(0, 1), edge(1, 2), edge(1, 3), edge(2, 3), edge(3, 4)],
+            root: PostAsapNodeId(4),
+        };
+        let mut bindings = binding();
+        bindings
+            .nodes
+            .insert(PostAsapNodeId(0), BackendNodeBinding::QueryInput);
+        let registry = FrontierRegistry(Registry::default());
+        let sink = Sink::default();
+        let result =
+            execute_precompute_sink(&dag, &bindings, PostAsapNodeId(3), key(3), &registry, &sink)
+                .unwrap();
+        assert_eq!(*result, 25);
+        assert_eq!(
+            *registry.0 .0.lock().unwrap(),
+            BTreeMap::from([(2, 1), (3, 1)])
+        );
+    }
+
     #[test]
     fn rejects_query_node_in_precompute_path_and_mismatched_lineage_key() {
         let mut query_child = node(0);