fix(planner): Stage 3 rejects plans the executor cannot compile; no PromQL top-k heap over its instant vector

crates/devtools/tests/stage_pipeline.rs

@@ -163,40 +163,34 @@ fn display_cap_is_recorded_and_keeps_the_selection() {
     );
 }
 
-/// A choice composition cannot substitute is listed with its error, and the
-/// other candidates are still written.
+/// A top-k partitioned by every label of its input has no item, so a heap
+/// over it could not be composed; such a heap reads one row per item and
+/// is pruned in Pass 1 (`TopKHeapOverOneRowPerItem`), listed under
+/// `pass1_pruned`, and every candidate left composes.
 #[test]
-fn uncomposable_candidate_is_recorded_with_its_error() {
-    let out =
-        std::env::temp_dir().join(format!("stage_pipeline_{}_error.json", std::process::id()));
-    let status = Command::new(env!("CARGO_BIN_EXE_stage_pipeline"))
-        .args([
-            "--promql",
-            // Top-K partitioned by every label of its input has no item.
-            "topk by (job) (3, sum by (job) (rate(m[5m])))",
-            "--max-error",
-            "0.01",
-            "--failure-probability",
-            "0.001",
-            "--out",
-        ])
-        .arg(&out)
-        .status()
-        .unwrap();
-    assert!(status.success());
-    let document: Value = serde_json::from_str(&std::fs::read_to_string(&out).unwrap()).unwrap();
-    std::fs::remove_file(out).unwrap();
-    let candidates = document["stage1_logical_asap_optimization"]["candidates"]
-        .as_array()
-        .unwrap();
-    let (composed, failed): (Vec<_>, Vec<_>) = candidates
-        .iter()
-        .partition(|candidate| candidate.get("dag").is_some());
-    assert!(!composed.is_empty() && !failed.is_empty());
-    for candidate in failed {
-        assert!(candidate["error"].as_str().unwrap().contains("Top-K"));
-    }
-    for candidate in composed {
+fn a_topk_heap_over_one_row_per_item_is_pruned_with_its_reason() {
+    let document = document_of(&[
+        "--promql",
+        "topk by (job) (3, sum by (job) (rate(m[5m])))",
+        "--max-error",
+        "0.01",
+        "--failure-probability",
+        "0.001",
+    ]);
+    let stage1 = &document["stage1_logical_asap_optimization"];
+    assert_eq!(
+        stage1["pass1_pruned"],
+        serde_json::json!([{
+            "target": 0,
+            "query": "Q1",
+            "pruned": [
+                { "candidate": "CMS+heap", "reason": "TopKHeapOverOneRowPerItem" },
+                { "candidate": "CountSketch+heap", "reason": "TopKHeapOverOneRowPerItem" },
+            ],
+            "fused_pruned": [],
+        }])
+    );
+    for candidate in stage1["candidates"].as_array().unwrap() {
         validated_dag(&candidate["dag"], 1);
     }
 }
@@ -405,12 +399,13 @@ fn workload_above_the_candidate_limit_is_rejected() {
     );
 }
 
-/// The default limit is 10^4: nine approximate Top-K queries (3^9 = 19683
-/// Pass 1 combinations) are rejected once the count passes it.
+/// The default limit is 10^4: nine approximate quantile queries, three
+/// candidates each (3^9 = 19683 Pass 1 combinations), are rejected once the
+/// count passes it.
 #[test]
 fn default_candidate_limit_is_ten_thousand() {
     let queries: Vec<String> = (1..=9)
-        .map(|i| format!("topk(5, rate(m{i}[5m]))"))
+        .map(|i| format!("quantile_over_time(0.99, m{i}[5m])"))
         .collect();
     let mut args = Vec::new();
     for query in &queries {

crates/logical-asap-optimization/src/per_computation_candidates.rs

@@ -21,7 +21,10 @@
 //!   grouped `sum` over a `sum` (`ExactAccumulator([Sum])`).
 //! - Summary candidates are unsized; they are sized after Pass 2 sharing.
 //! - Candidates are pruned only when provably unable to meet the accuracy
-//!   requirement, and each pruned candidate keeps its [`PruneReason`].
+//!   requirement or to run, or when the exact candidate provably dominates
+//!   them (a PromQL top-k's heap over its instant vector, one row per item:
+//!   [`PruneReason::TopKHeapOverOneRowPerItem`]); each pruned candidate
+//!   keeps its [`PruneReason`].
 //! - A measure filter or `HAVING` prunes nothing. The target records the
 //!   measure filter (applied as `SummaryAgg.filter`) and one [`HavingEntry`]
 //!   per consumer `Filter` directly above the aggregate: a key-only predicate
@@ -178,6 +181,14 @@ pub enum PruneReason {
     /// +Inf: no sketch estimates that, and the executor's sketches refuse
     /// such a φ, so only the exact answer remains.
     QuantileOutsideUnitInterval,
+    /// A PromQL top-k's heap (Count-Min or CountSketch with a heap) over its
+    /// own input, an instant vector with one row per ranked series or
+    /// group: the heap reads every row and adds
+    /// nothing to it, so an exact top-k over the same rows, which keeps only
+    /// `k` of them, is exact and no costlier. Only the whole-expression heap
+    /// over the inner aggregate's input, where an item's rows repeat, is
+    /// offered ([`FusedComputation`]).
+    TopKHeapOverOneRowPerItem,
 }
 
 /// Unranked candidates of one measure and the candidates pruned for it.
@@ -671,6 +682,7 @@ pub fn generate_per_computation_candidates<Id: std::fmt::Debug>(
         }
     }
     add_whole_expression_targets(&mut targets, &consumers);
+    prune_topk_heaps_over_one_row_per_item(&mut targets);
     add_nested_sum_targets(&mut targets, &consumers);
     Ok(PerComputationCandidates { roots, targets })
 }
@@ -855,6 +867,46 @@ fn add_whole_expression_targets(
     }
 }
 
+/// Prunes the own heap candidates of each PromQL top-k
+/// ([`PruneReason::TopKHeapOverOneRowPerItem`]); a fused whole-expression
+/// heap ([`add_whole_expression_targets`]) has already moved out of them.
+fn prune_topk_heaps_over_one_row_per_item(targets: &mut [ComputationTarget]) {
+    for target in targets.iter_mut() {
+        let Some(NonASAPOp::Aggregate {
+            measures, child, ..
+        }) = target.aggregate.non_asap()
+        else {
+            continue;
+        };
+        if !matches!(measures.get(target.measure), Some(AggIntent::TopK { .. })) {
+            continue;
+        }
+        // A PromQL top-k ranks an instant vector: one row per series or
+        // group. (A SQL top-k keeps its heaps.)
+        if !reads_only_time_series(child) {
+            continue;
+        }
+        target
+            .choices
+            .prune(PruneReason::TopKHeapOverOneRowPerItem, |candidate| {
+                matches!(
+                    candidate,
+                    ComputationCandidate::Summary(
+                        SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap
+                    )
+                )
+            });
+    }
+}
+
+/// Whether every scan below `node` reads a time series (PromQL rows).
+fn reads_only_time_series(node: &OperatorNode) -> bool {
+    match node.non_asap() {
+        Some(NonASAPOp::Scan { source, .. }) => matches!(source, Source::TimeSeries { .. }),
+        _ => node.children().iter().all(|c| reads_only_time_series(c)),
+    }
+}
+
 /// The input, update and partitions of a whole-expression top-k over
 /// `target`'s inner aggregate, by the keyed-additive rule, or `None` when it
 /// does not apply. Rows that carry the full series identity rank it as a

crates/physical-asap-optimization/src/lib.rs

@@ -75,7 +75,8 @@
 //! columns, so an exact root has the same column names as a summary
 //! candidate's and as the frontend declares. Over PromQL rows it also keeps
 //! the series identity, after them: `topk` answers whole input series, so
-//! each answer carries the chosen series' full label set.
+//! each answer carries the chosen series' full label set. PromQL rows without
+//! one (e.g. of `sum by (job)`) keep their other string label columns instead.
 //! `topk_k` carries the ranked value's type and nullability (the frontend
 //! declares a non-null `Utf8` placeholder); casting the value to `Utf8` would
 //! not make it the same answer.
@@ -96,7 +97,7 @@ use asap_types::ir::physical_export::{
 use asap_types::ir::properties::{ExecutionDataStateError, ExecutionTiming};
 use asap_types::ir::scalar::column_resolution::resolve_column_ref;
 use asap_types::ir::scalar::expr_ir::ColumnRef;
-use asap_types::ir::schema::{ExactKind, FieldDataType};
+use asap_types::ir::schema::{DataType, ExactKind, FieldDataType};
 use asap_types::ir::{
     apply_materialization_timings, is_filtered_build, split_shared_by_timing, ASAPOp,
     MaterializationAssignment, NonASAPOp, Operator, OperatorNode, ProjectItem, QueryRoot,
@@ -779,6 +780,20 @@ fn implement_topk(
             alias: Some(asap_types::ir::schema::PROMQL_SERIES_IDENTITY.into()),
             expr: ScalarExpr::Column(identity),
         });
+    } else if reads_only_time_series(&child) {
+        // PromQL rows without it (e.g. those of `sum by (job)`) carry their
+        // labels as plain string columns: each answer keeps the ones the
+        // logical columns do not, so `topk(k, sum by (job) (...))` answers
+        // each job. A SQL result keeps only its logical columns.
+        for (index, field) in child.schema.fields.iter().enumerate() {
+            let is_label = field.dtype == FieldDataType::Plain(DataType::Utf8);
+            if is_label && index != value && !key_fields.iter().any(|k| k.name == field.name) {
+                cols.push(ProjectItem {
+                    alias: Some(field.name.clone()),
+                    expr: ScalarExpr::Column(index),
+                });
+            }
+        }
     }
     Ok(Rc::new(
         OperatorNode::new(Operator::NonASAP(NonASAPOp::Project {
@@ -888,7 +903,8 @@ mod tests {
         let root = lower_promql("topk(5, sum by (job) (m))", AccuracyTarget::exact());
         let candidate = first_candidate(&[root.into()]);
         let physical = operator_root(&candidate, 0);
-        assert_eq!(names(&physical), ["topk_5"]);
+        // The answers keep their `job` label.
+        assert_eq!(names(&physical), ["topk_5", "job"]);
         let limit = topk_limit(&physical);
         assert_eq!(names(&limit), ["job", "sum"]);
         assert_eq!(topk_sort_key(&limit), (ScalarExpr::Column(1), false));

crates/plan-selection/src/lib.rs

@@ -269,9 +269,9 @@ pub mod pane_sharing;
 /// The Stage 1 to Stage 3 driver and the optimization-pass interface.
 pub mod pass;
 pub use pass::{
-    optimize, run_stage1_to_stage3, run_stage1_to_stage3_with, OptimizationInput,
-    OptimizationInputError, OptimizationPass, OptimizeError, PassNameConflict, PassRegistry,
-    PlanOutput, PlanningModels, QueryPlan, Stage1ToStage3Pass, Stage1ToStage3Run,
+    optimize, run_stage1_to_stage3, run_stage1_to_stage3_with, ExecutorCompileCheck,
+    OptimizationInput, OptimizationInputError, OptimizationPass, OptimizeError, PassNameConflict,
+    PassRegistry, PlanOutput, PlanningModels, QueryPlan, Stage1ToStage3Pass, Stage1ToStage3Run,
     Stage1ToStage3Search,
 };
 
@@ -749,6 +749,11 @@ pub enum RejectionCause {
     /// executor cannot compile it
     /// ([`asap_types::ir::operator::agg_has_exact_per_entity_kernel`]).
     PerEntityAggregateNotExecutable,
+    /// The deployment's executor cannot compile its plan
+    /// ([`PlanningModels::executor_compile_check`]), e.g. a summary readout
+    /// whose output shape the executor does not build. The reason is the
+    /// executor's error.
+    ExecutorCannotCompile,
     /// A query's output has no accuracy guarantee (Stage 1 has no
     /// analytical rule for it and the deployment's evidence supplies none),
     /// and the deployment's accuracy model did not accept it directly
@@ -791,6 +796,7 @@ impl RejectionCause {
                 | Self::AccuracyTargetMissed
                 | Self::CountMinWeightsNotProvenNonNegative
                 | Self::PerEntityAggregateNotExecutable
+                | Self::ExecutorCannotCompile
                 | Self::NoAccuracyGuarantee
                 | Self::RawDataNotRetained
                 | Self::LatencyBoundExceeded
@@ -1208,6 +1214,17 @@ impl<'a> PlanSelector<'a> {
             if !selectable {
                 return Ok(());
             }
+        }
+        // Only a candidate that may be selected is compiled: the check runs
+        // the executor's own compiler, the costliest check.
+        if let Some(check) = self.models.executor_compile_check {
+            if let Err(reason) = check.compiles(&sized.roots) {
+                self.reject(&candidate.id, RejectionCause::ExecutorCannotCompile, reason);
+                return Ok(());
+            }
+        }
+        if !self.keep_every_candidate {
+            let kept_cost = |kept: &KeptCandidate| &self.costs[&kept.built.id];
             let dominated: Vec<bool> = self
                 .kept
                 .iter()
@@ -5381,6 +5398,56 @@ mod tests {
         }
     }
 
+    /// An executor that compiles every plan but the one it is first asked
+    /// about.
+    struct RefusesTheFirstPlan(std::cell::Cell<bool>);
+
+    impl crate::ExecutorCompileCheck for RefusesTheFirstPlan {
+        fn compiles(&self, _roots: &[QueryRoot]) -> Result<(), String> {
+            if self.0.replace(true) {
+                Ok(())
+            } else {
+                Err("invalid keyed evaluation shape or capacity".into())
+            }
+        }
+    }
+
+    /// A candidate the deployment's executor does not compile is rejected
+    /// as `ExecutorCannotCompile` with the executor's reason, and another
+    /// is selected; without the check the same candidate is valid.
+    #[test]
+    fn a_plan_the_executor_cannot_compile_is_rejected() {
+        let target = AccuracyTarget::max_error(0.01)
+            .failure_probability(0.001)
+            .accuracy_metric(AccuracyMetric::ABSOLUTE);
+        let entries = [entry_with(target, repeating(10_000))];
+        // More than one valid candidate: every output accepted.
+        let settings = PlanSelectionSettings {
+            ignore_accuracy_guarantees: true,
+            ..PlanSelectionSettings::default()
+        };
+        let models = PlanningModels::builtin().with_plan_selection(&settings);
+        let unchecked = select_plan(&candidates(), &entries, &data(), models).unwrap();
+        let check = RefusesTheFirstPlan(std::cell::Cell::new(false));
+        let checked = select_plan(
+            &candidates(),
+            &entries,
+            &data(),
+            models.with_executor_compile_check(&check),
+        )
+        .unwrap();
+        let refused = checked
+            .rejected
+            .iter()
+            .find(|r| r.cause == RejectionCause::ExecutorCannotCompile)
+            .expect("the first plan is refused");
+        assert_eq!(refused.reason, "invalid keyed evaluation shape or capacity");
+        assert!(refused.cause.failed_a_check());
+        assert!(unchecked.costs.contains_key(&refused.id), "{}", refused.id);
+        assert!(!checked.costs.contains_key(&refused.id));
+        assert_ne!(checked.selected, refused.id);
+    }
+
     /// `ignore_accuracy_guarantees` (issue #774) does not override the
     /// executor's requirement on Count-Min update weights either: over
     /// samples of a gauge, P2's CmsWithHeap is still rejected as

crates/plan-selection/src/pass/mod.rs

@@ -80,6 +80,20 @@ pub struct PlanningModels<'a> {
     /// whether it keeps query-time results across evaluations (Stage 2);
     /// default [`DeploymentCapabilities::default`].
     pub deployment_capabilities: &'a DeploymentCapabilities,
+    /// Whether the deployment's executor compiles a candidate's plan; Stage
+    /// 3 rejects a candidate it does not
+    /// ([`RejectionCause::ExecutorCannotCompile`](crate::RejectionCause::ExecutorCannotCompile)).
+    /// `None` (the default) checks nothing; `asap_planner::plan_installable`
+    /// sets the Planner's executor when the caller leaves it `None`.
+    pub executor_compile_check: Option<&'a dyn ExecutorCompileCheck>,
+}
+
+/// The deployment executor's compiler, asked by Stage 3 whether it compiles
+/// a candidate's plan, so selection and execution cannot drift.
+pub trait ExecutorCompileCheck {
+    /// `Err` with the executor's reason when it cannot compile the plan of
+    /// `roots`: one candidate's timed query roots, in workload order.
+    fn compiles(&self, roots: &[asap_types::ir::QueryRoot]) -> Result<(), String>;
 }
 
 impl<'a> PlanningModels<'a> {
@@ -94,6 +108,7 @@ impl<'a> PlanningModels<'a> {
             evidence,
             plan_selection: &DEFAULT_PLAN_SELECTION,
             deployment_capabilities: &DEFAULT_DEPLOYMENT,
+            executor_compile_check: None,
         }
     }
 
@@ -104,6 +119,7 @@ impl<'a> PlanningModels<'a> {
             evidence: &NO_ACCURACY_EVIDENCE,
             plan_selection: &DEFAULT_PLAN_SELECTION,
             deployment_capabilities: &DEFAULT_DEPLOYMENT,
+            executor_compile_check: None,
         }
     }
 
@@ -136,6 +152,12 @@ impl<'a> PlanningModels<'a> {
         self
     }
 
+    /// These models with the deployment executor's compile check.
+    pub fn with_executor_compile_check(mut self, check: &'a dyn ExecutorCompileCheck) -> Self {
+        self.executor_compile_check = Some(check);
+        self
+    }
+
     /// These models with the deployment's capabilities (supported
     /// summaries, raw-data and query-time retention).
     pub fn with_deployment_capabilities(