From d889c14c6baae020ec843ac5c11409cbdbb11ce4 Mon Sep 17 00:00:00 2001 From: zzylol <50204836+zzylol@users.noreply.github.com> Date: Mon, 5 Oct 2026 13:10:46 +0000 Subject: [PATCH] fix(logical-optimizer): move whole-expression top-k partitions onto the heap input The whole-expression heap reads the absorbed aggregate's input, but kept the outer top-k's partition keys, which index that aggregate's output: `topk by(job)(k, sum by(service, job)(rate(x[5m])))` partitioned by `ts` and lost `job`. Resolve the keys by column onto the heap's input, and do not offer the heap when they cannot be resolved. Fixes #639 Co-Authored-By: Claude Opus 5.5 --- .../executor/tests/weighted_topk_binding.rs | 41 ++++--- .../src/pass1/logical_candidates.rs | 106 ++++++++++++++++-- 2 files changed, 120 insertions(+), 27 deletions(-) diff --git a/crates/executor/tests/weighted_topk_binding.rs b/crates/executor/tests/weighted_topk_binding.rs index b2268c2d6..637843623 100644 --- a/crates/executor/tests/weighted_topk_binding.rs +++ b/crates/executor/tests/weighted_topk_binding.rs @@ -14,8 +14,6 @@ use planner_types::ir::schema::{DataType, *}; use planner_types::types::AccuracyTarget; use std::{collections::BTreeMap, rc::Rc, sync::Arc}; #[test] -#[ignore = "Stage 1's whole-expression heap keeps the outer top-k's partition keys, which \ - index the inner aggregate's output, over the raw rates it reads"] fn planner_weighted_topk_binds_at_either_deployment_phase() { assert_weighted_binding(SketchAlgorithm::CmsWithHeap); assert_weighted_binding(SketchAlgorithm::CountSketchWithHeap); @@ -124,23 +122,34 @@ fn assert_weighted_binding(algorithm: SketchAlgorithm) { } output }); - assert_eq!(output.len(), 4); - let mut scores = output + // Two items per `job`: the outer `by(job)` partitions the ranking. + let schema = &dag + .nodes .iter() - .map(|row| { - row.iter() - .find_map(|v| { - if let Value::Float64(v) = v { - Some(*v) - } else { - None - } - }) - .unwrap() + .find(|node| node.id == dag.roots[0]) + .unwrap() + .output_schema; + let column = |name: &str| schema.fields.iter().position(|f| f.name == name).unwrap(); + let (job, value) = (column("job"), column("value")); + let mut ranked = output + .iter() + .map(|row| match (&row[job], &row[value]) { + (Value::Utf8(job), Value::Float64(score)) => (job.to_string(), *score), + other => panic!("unexpected (job, value) {other:?}"), }) .collect::>(); - scores.sort_by(f64::total_cmp); - assert_eq!(scores, vec![0.3125, 0.375, 80., 100.]); + ranked.sort_by(|a, b| a.0.cmp(&b.0).then(a.1.total_cmp(&b.1))); + let expected = [ + ("api", 0.3125), + ("api", 0.375), + ("batch", 80.), + ("batch", 100.), + ]; + assert_eq!( + ranked, + expected.map(|(job, score)| (job.to_string(), score)), + "{phase:?}" + ); } } diff --git a/crates/logical-optimizer/src/pass1/logical_candidates.rs b/crates/logical-optimizer/src/pass1/logical_candidates.rs index 0db14e932..a9a3b38d6 100644 --- a/crates/logical-optimizer/src/pass1/logical_candidates.rs +++ b/crates/logical-optimizer/src/pass1/logical_candidates.rs @@ -311,11 +311,14 @@ fn count_item<'a>( }) } -/// The input and update 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 column, as -/// [`summary_update`] does. -fn whole_expression_input(target: &OperatorNode) -> Option<(Rc, SummaryUpdate)> { +/// 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 +/// column, as [`summary_update`] does. The partitions are `target`'s, moved +/// onto the input the heap reads ([`partitions_over`]). +fn whole_expression_input( + target: &OperatorNode, +) -> Option<(Rc, SummaryUpdate, Reduction)> { let Some(NonASAPOp::Aggregate { child, reduction, @@ -363,7 +366,34 @@ fn whole_expression_input(target: &OperatorNode) -> Option<(Rc, Su // An encoded label set needs an open PromQL schema. return None; } - Some((realized.child, input)) + let partitions = partitions_over(reduction, &child.schema, &realized.child.schema)?; + Some((realized.child, input, partitions)) +} + +/// `reduction`'s keys, which index `inner` (the absorbed aggregate's +/// output), as indices into `input` (the rows the heap reads), matched by +/// column. `None` when a key is missing or ambiguous in `input`, or the keys +/// are given by `without`: `input` carries columns `inner` dropped. +fn partitions_over(reduction: &Reduction, inner: &Schema, input: &Schema) -> Option { + let Reduction::Reduce(keys) = reduction else { + return Some(Reduction::PerEntity); + }; + if keys.is_without() { + return None; + } + let keys = keys + .iter() + .map(|&index| { + let key = column_ref(inner.fields.get(index)?); + let mut matches = + (0..input.fields.len()).filter(|&i| column_ref(&input.fields[i]) == key); + match (matches.next(), matches.next()) { + (Some(i), None) => Some(i), + _ => None, + } + }) + .collect::>()?; + Some(Reduction::by(keys)) } /// Number of whole-workload candidates: one per choice of an alternative for @@ -646,9 +676,9 @@ fn realize( )?), false => None, }; - let child = match &whole { - Some((input, _)) => memo[&Rc::as_ptr(input)].clone(), - None => memo[&Rc::as_ptr(child)].clone(), + let (child, reduction) = match &whole { + Some((input, _, partitions)) => (memo[&Rc::as_ptr(input)].clone(), partitions), + None => (memo[&Rc::as_ptr(child)].clone(), reduction), }; let (family, query) = match realization { Realization::ExactAggregate { kind, params } => ( @@ -661,8 +691,8 @@ fn realize( ), _ => return Err(LogicalCandidateError::Unsupported("summary family")), }; - let mut input = match whole { - Some((_, update)) => update, + let mut input = match &whole { + Some((_, update, _)) => update.clone(), None if *grouping != GroupingStrategy::PerSubpopulationInstance => { hydra_update(reduction, &child.schema) .ok_or(LogicalCandidateError::Unsupported("Hydra item column"))? @@ -1008,6 +1038,60 @@ mod tests { } } + /// The whole-expression heap reads the inner sum's input, so the outer + /// top-k's partition keys (indices into the sum's output) are moved onto + /// that input by column. + #[test] + fn whole_expression_topk_partitions_index_the_heap_input() { + let root = lower_promql( + "topk by (job) (2, sum by (service, job) (rate(m[1m])))", + AccuracyTarget::Epsilon(0.1), + ); + let inventory = enumerate_local_logical_candidates( + vec![(0, QueryRoot::Operator(root))], + &BTreeMap::new(), + ) + .unwrap(); + let mut heaps = 0; + for choice in enumerate_choices(&inventory, usize::MAX) { + if !choice + .iter() + .enumerate() + .any(|(t, &c)| inventory.targets[t].absorbs[c].is_some()) + { + continue; + } + let roots = compose_logical_candidate(&inventory, &choice).unwrap(); + let QueryRoot::Operator(root) = &roots[0].1 else { + panic!("operator root") + }; + for node in OperatorNode::reachable(root) { + if let Some(ASAPOp::SummaryAgg { + child, + reduction, + family: FieldDataType::Sketch(kind, _), + .. + }) = node.asap() + { + if !matches!( + kind.algorithm(), + SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap + ) { + continue; + } + let names: Vec<_> = reduction + .expect_reduce() + .iter() + .map(|&index| child.schema.fields[index].name.as_str()) + .collect(); + assert_eq!(names, ["job"], "{choice:?}"); + heaps += 1; + } + } + } + assert!(heaps > 0, "a whole-expression heap"); + } + /// The weight domain of every heap-sketch update (Count-Min or /// CountSketch + heap) in every candidate of `query`. fn heap_weight_domains(query: &str, metric_types: &[(&str, MetricType)]) -> Vec {