Unify workload selection for startup and distributed publication

control_plane/src/main.rs

@@ -759,6 +759,7 @@ fn compile_physical_plan_request(
         .unwrap_or_default()
         .as_millis() as u64;
     let mut queries = Vec::with_capacity(request.queries.len());
+    let mut canonical_roots = Vec::with_capacity(request.queries.len());
     for query in request.queries {
         if query.query_id.trim().is_empty()
             || query.metric.trim().is_empty()
@@ -773,15 +774,11 @@ fn compile_physical_plan_request(
             Ok(expr) => expr,
             Err(error) => return Err((StatusCode::UNPROCESSABLE_ENTITY, error.to_string())),
         };
-        let post_asap = match physical::compiler::select_post_asap(
-            &expr,
-            query.accuracy.clone(),
-            &query.lifecycle,
-            request.evidence.get(&query.query_id),
-        ) {
+        let post_asap = match control_plane::planner_selection::keep_pre_asap(&expr) {
             Ok(plan) => plan,
             Err(error) => return Err((StatusCode::UNPROCESSABLE_ENTITY, error.to_string())),
         };
+        canonical_roots.push(std::rc::Rc::new(expr));
         queries.push(physical::compiler::PlanningQuery {
             query_id: query.query_id,
             query_string: query.query_string,
@@ -798,6 +795,12 @@ fn compile_physical_plan_request(
         });
     }
 
+    if let Err(error) =
+        physical::compiler::select_workload_roots(&mut queries, canonical_roots, &request.evidence)
+    {
+        return Err((StatusCode::UNPROCESSABLE_ENTITY, error.to_string()));
+    }
+
     let bundle = match physical::compiler::PhysicalCompiler.compile(
         physical::compiler::PlanningRequest {
             queries,

control_plane/src/physical/compiler.rs

@@ -1619,25 +1619,7 @@ impl BackendLocalPlanningSnapshot {
                 runtime_policy: RuntimeRulePolicy::default(),
             });
         }
-        // A cohort shares the same end-to-end requirement, not an inferred
-        // weakest common accuracy. Different targets are searched separately.
-        let mut cohorts: Vec<(AccuracyTarget, Vec<(usize, Rc<QueryExpr>)>)> = Vec::new();
-        for (index, root) in canonical_roots.into_iter().enumerate() {
-            let accuracy = &queries[index].accuracy;
-            if let Some((_, roots)) = cohorts.iter_mut().find(|(target, _)| target == accuracy) {
-                roots.push((index, root));
-            } else {
-                cohorts.push((accuracy.clone(), vec![(index, root)]));
-            }
-        }
-        for (accuracy, roots) in cohorts {
-            let model = ControlPlaneCostModel::new(accuracy.clone());
-            let selected = crate::planner_selection::select_workload(roots, accuracy, &model)
-                .map_err(|error| CompileError::Snapshot(error.to_string()))?;
-            for (index, node) in selected {
-                queries[index].post_asap = node;
-            }
-        }
+        select_workload_roots(&mut queries, canonical_roots, &HashMap::new())?;
         PhysicalCompiler.compile(
             PlanningRequest {
                 queries,
@@ -2072,6 +2054,52 @@ fn summary_agg_metric(node: &SummaryNode) -> Option<String> {
         .flatten()
 }
 
+/// Shared selection boundary for canonical startup and compile-and-publish.
+/// Certificate-bearing roots stay isolated: equal certificate values do not
+/// establish that the certificate's source scope covers another query.
+pub fn select_workload_roots(
+    queries: &mut [PlanningQuery],
+    roots: Vec<Rc<QueryExpr>>,
+    evidence: &HashMap<String, TopKMembershipEvidence>,
+) -> Result<(), CompileError> {
+    if roots.len() != queries.len() {
+        return Err(CompileError::Snapshot(
+            "canonical root/query mapping is incomplete".into(),
+        ));
+    }
+    let mut cohorts: Vec<(AccuracyTarget, Option<String>, Vec<(usize, Rc<QueryExpr>)>)> =
+        Vec::new();
+    for (index, root) in roots.into_iter().enumerate() {
+        let accuracy = &queries[index].accuracy;
+        let certificate_scope = evidence
+            .contains_key(&queries[index].query_id)
+            .then(|| queries[index].query_id.clone());
+        if let Some((_, _, roots)) = cohorts
+            .iter_mut()
+            .find(|(target, scope, _)| target == accuracy && scope == &certificate_scope)
+        {
+            roots.push((index, root));
+        } else {
+            cohorts.push((accuracy.clone(), certificate_scope, vec![(index, root)]));
+        }
+    }
+    for (accuracy, scope, roots) in cohorts {
+        let model = ControlPlaneCostModel::new(accuracy.clone());
+        let certificate = scope.as_ref().and_then(|id| evidence.get(id));
+        let selected = crate::planner_selection::select_workload_with_evidence(
+            roots,
+            accuracy,
+            &model,
+            &QueryEvidence(certificate),
+        )
+        .map_err(|error| CompileError::Snapshot(error.to_string()))?;
+        for (index, node) in selected {
+            queries[index].post_asap = node;
+        }
+    }
+    Ok(())
+}
+
 /// Planner-adapter selection step used before physical compilation. Keeping
 /// this separate makes the ownership boundary explicit: callers supply the
 /// selected post-ASAP DAG to [`PhysicalCompiler::compile`].
@@ -2611,6 +2639,51 @@ mod tests {
         request_with_evidence(query_id, promql, None).expect("post-ASAP selection")
     }
 
+    // Both production adapters preserve canonical root identity and select the
+    // whole evidence-free cohort, rather than independently binding roots.
+    #[test]
+    fn shared_selection_adapter_preserves_query_mapping() {
+        let mut workload = request("q90", "quantile_over_time(0.9, m[1m])");
+        workload
+            .queries
+            .extend(request("q99", "quantile_over_time(0.99, m[1m])").queries);
+        let roots = workload
+            .queries
+            .iter()
+            .map(|query| {
+                Rc::new(
+                    crate::query_parser::parse_query_expr_canonical(
+                        &query.query_string,
+                        query.accuracy.clone(),
+                    )
+                    .unwrap(),
+                )
+            })
+            .collect();
+        select_workload_roots(&mut workload.queries, roots, &workload.evidence).unwrap();
+        let bundle = PhysicalCompiler
+            .compile(workload, environment(10000))
+            .unwrap();
+        assert_eq!(bundle.query_plan.entries.len(), 2);
+        assert_eq!(bundle.collector_plans[0].materializations.len(), 1);
+        assert_eq!(
+            bundle
+                .query_plan
+                .entries
+                .values()
+                .map(|entry| entry.query_id.as_str())
+                .collect::<BTreeSet<_>>(),
+            BTreeSet::from(["q90", "q99"])
+        );
+    }
+
+    // A broken input mapping must be rejected, never silently drop a root.
+    #[test]
+    fn shared_selection_rejects_incomplete_root_mapping() {
+        let mut workload = request("q", "quantile_over_time(0.9, m[1m])");
+        assert!(select_workload_roots(&mut workload.queries, vec![], &workload.evidence).is_err());
+    }
+
     #[test]
     fn shared_materialization_is_emitted_once_for_every_runtime() {
         for target in [

control_plane/src/planner_selection.rs

@@ -158,12 +158,33 @@ pub fn select_workload(
     roots: Vec<(usize, Rc<QueryExpr>)>,
     accuracy: AccuracyTarget,
     cost_model: &dyn CostModel,
+) -> Result<Vec<(usize, Rc<SummaryNode>)>, SelectionError> {
+    select_workload_with_evidence(
+        roots,
+        accuracy,
+        cost_model,
+        &asap_aware_mapping::NoAccuracyEvidence,
+    )
+}
+
+/// The entire cohort uses the same scoped accuracy certificate; callers must
+/// not spread one query's evidence to unrelated workload roots.
+pub fn select_workload_with_evidence(
+    roots: Vec<(usize, Rc<QueryExpr>)>,
+    accuracy: AccuracyTarget,
+    cost_model: &dyn CostModel,
+    evidence: &dyn AccuracyEvidenceProvider,
 ) -> Result<Vec<(usize, Rc<SummaryNode>)>, SelectionError> {
     // Canonical CSE still runs inside search_workload_with_targets. Do not
     // offer CSE's per-invocation recompute alternative: this runtime currently
     // provisions continuously maintained, content-addressed state only.
     let strategies: Vec<Box<dyn ReplacementStrategy + '_>> = vec![
-        Box::new(SketchAlgorithmStrategy::new(cost_model)),
+        Box::new(SketchAlgorithmStrategy::with_models_and_evidence(
+            cost_model,
+            &asap_aware_mapping::DefaultAccuracyModel,
+            &asap_aware_mapping::EqualSplitAllocator,
+            evidence,
+        )),
         Box::new(asap_aware_mapping::SemanticEquivalentRewriteStrategy),
     ];
     let space = asap_aware_mapping::search_workload_with_targets(

data_plane/src/precompute_engine/operators/sum_accumulator.rs

@@ -49,12 +49,15 @@ impl SumAccumulator {
     }
 
     pub fn deserialize_from_bytes(buffer: &[u8]) -> Result<Self, Box<dyn std::error::Error>> {
-        if buffer.len() < 4 {
-            return Err("Buffer too short for f32".into());
+        match buffer.len() {
+            // Legacy Python scalar sums carry no sample-count evidence.
+            4 => Ok(Self::with_sum(f32::from_le_bytes(buffer.try_into()?) as f64)),
+            // Counted sums use the same fixed layout as the Collector Sum payload.
+            16 => Self::from_sum_bytes(buffer),
+            len => {
+                Err(format!("Invalid persisted Sum payload length: {len} (want 4 or 16)").into())
+            }
         }
-        // Python uses struct.pack("<f", self.sum) which is 4-byte little-endian float
-        let sum = f32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as f64;
-        Ok(Self::with_sum(sum))
     }
 
     /// Decode the fixed Sum payload produced by the first-class Sum
@@ -96,8 +99,15 @@ impl SerializableToSink for SumAccumulator {
     }
 
     fn serialize_to_bytes(&self) -> Vec<u8> {
-        // Match Python's struct.pack("<f", self.sum) - 4-byte little-endian float
-        (self.sum as f32).to_le_bytes().to_vec()
+        match self.observation_count {
+            Some(count) => {
+                let mut bytes = Vec::with_capacity(16);
+                bytes.extend_from_slice(&self.sum.to_le_bytes());
+                bytes.extend_from_slice(&count.to_le_bytes());
+                bytes
+            }
+            None => (self.sum as f32).to_le_bytes().to_vec(),
+        }
     }
 }
 
@@ -272,6 +282,40 @@ mod tests {
         assert_eq!(merged.aux_stats().count, None);
     }
 
+    // Persistence retains known counts, including zero and the full u64 range.
+    #[test]
+    fn counted_sum_binary_round_trip() {
+        for count in [0, 3, u64::MAX] {
+            let acc = SumAccumulator {
+                sum: 1.0000000000001,
+                observation_count: Some(count),
+            };
+            let bytes = acc.serialize_to_bytes();
+            assert_eq!(bytes.len(), 16);
+            let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap();
+            assert_eq!(restored.sum, acc.sum);
+            assert_eq!(restored.observation_count, Some(count));
+        }
+    }
+
+    // Existing scalar-only files remain readable without inventing counts.
+    #[test]
+    fn legacy_binary_sum_has_unknown_count() {
+        let bytes = 42.5f32.to_le_bytes();
+        let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap();
+        assert_eq!(restored.sum, 42.5);
+        assert_eq!(restored.observation_count, None);
+        assert_eq!(restored.serialize_to_bytes(), bytes);
+    }
+
+    // Truncated counted payloads must not silently decode as scalar sums.
+    #[test]
+    fn persisted_sum_rejects_invalid_lengths() {
+        for len in [0, 3, 5, 8, 15, 17] {
+            assert!(SumAccumulator::deserialize_from_bytes(&vec![0; len]).is_err());
+        }
+    }
+
     #[test]
     fn test_sum_accumulator_creation() {
         let acc = SumAccumulator::new();

data_plane/src/storage_engines/sketch_db/index/mod.rs

@@ -3924,6 +3924,70 @@ mod tests {
         drop(p);
     }
 
+    // Raw sample counts must survive production durable storage.
+    #[test]
+    fn raw_count_survives_disk_eviction() {
+        use crate::storage_engines::types::AggregationType;
+        let tmp = tempfile::TempDir::new().unwrap();
+        let idx = Arc::new(SketchStore::new());
+        // Register an ExactAgg(Sum) sid keyed by `zone`.
+        let mut m = meta(8001);
+        m.metric_name = "http_requests_total".into();
+        m.group_by_keys = ["zone".to_string()].into_iter().collect();
+        m.agg_kind = AggKind::ExactAgg {
+            agg_type: AggregationType::Sum,
+            parameters_canonical: String::new(),
+            spatial_filter_canonical: String::new(),
+        };
+        idx.register(m);
+        let p = idx
+            .start_persistence(durable_cfg(tmp.path().to_path_buf()))
+            .unwrap();
+
+        let lv_zone = |v: &str| {
+            let mut x = BTreeMap::new();
+            x.insert("zone".to_string(), v.to_string());
+            x
+        };
+        for i in 0..10u64 {
+            let s = i * 30_000;
+            idx.append_precompute(
+                8001,
+                lv_zone("z0"),
+                (s, s + 30_000),
+                Box::new({
+                    let mut acc = crate::precompute_engine::operators::SumAccumulator::new();
+                    acc.update((i + 1) as f64);
+                    acc.update(10.0);
+                    acc
+                }),
+            );
+        }
+        assert!(
+            wait_until(
+                || idx.approx_memory_bytes() == 0 && idx.list_sealed_epochs_len() == 0,
+                std::time::Duration::from_secs(5),
+            ),
+            "exact-agg windows never fully evicted"
+        );
+        // Query the EVICTED portion [0, 150_000) — must come back from disk.
+        let series = idx.query_exact_agg_range(8001, 0, 150_000);
+        assert!(
+            !series.is_empty(),
+            "exact-agg query returned no result after flush and eviction"
+        );
+        let (_label, samples) = &series[0];
+        assert!(
+            samples.contains_key(&30_000),
+            "evicted exact-agg window missing from disk"
+        );
+        let stats = samples[&30_000].aux_stats();
+        assert_eq!(stats.count, Some(2));
+        assert_eq!(stats.sum, Some(11.0));
+        assert_eq!(stats.sum.unwrap() / stats.count.unwrap() as f64, 5.5);
+        drop(p);
+    }
+
     /// BUG #3: the memory diagnostic + the flusher's memory-pressure
     /// trigger must account for `current_epoch`, not just sealed epochs.
     /// On origin/main `approx_memory_bytes()` sums ONLY sealed epochs, so