fix: one finest-grain stored summary serves several groupings (quantiles at the estimate level, sums at the state level)

crates/integration-tests/tests/reuse_rules.rs

@@ -112,29 +112,43 @@ fn samples() -> Vec<(&'static str, &'static str, i64, f64)> {
     rows
 }
 
-/// Execute `root` at query time, every raw `TimeRange` input fed `samples`.
+/// Execute `root` at query time, every raw `TimeRange` input fed the
+/// `samples` in its window: `(end - range, end]`, where `end` is the
+/// evaluation time moved back by a `TimeShift` directly below (a tumbling
+/// pane's).
 fn execute(root: &Rc<OperatorNode>) -> Vec<String> {
     let timed = apply_materialization_timings(
         root,
         &MaterializationAssignment::default(),
         &mut TimingMemo::new(),
     )
     .unwrap();
+    // Tumbling panes run as their panes, each over its own window.
+    let timed = asap_types::ir::operator::pane_chain::expand_pane_chains(&timed).unwrap();
     let compiled = compile_physical_asap_dag_with_node_ids(&timed).unwrap();
     let dag = compiled.dag;
     let inputs: BTreeMap<u64, Batch> = dag
         .nodes
         .iter()
-        .filter(|node| {
-            matches!(
-                compiled.node_ids.operator_node(node.id).unwrap().operator,
-                Operator::NonASAP(NonASAPOp::TimeRange { .. })
-            )
+        .filter_map(|node| {
+            let Operator::NonASAP(NonASAPOp::TimeRange { range, child, .. }) =
+                &compiled.node_ids.operator_node(node.id).unwrap().operator
+            else {
+                return None;
+            };
+            let back_ms = match child.non_asap() {
+                Some(NonASAPOp::TimeShift { shift, .. }) => shift.offset_ms,
+                _ => 0,
+            };
+            let end = EVALUATION_MS - back_ms;
+            let start = end - range.as_millis() as i64;
+            Some((node, start, end))
         })
-        .map(|node| {
+        .map(|(node, start, end)| {
             let schema: Arc<Schema> = Arc::new(node.output_schema.clone());
             let rows = samples()
                 .into_iter()
+                .filter(|&(_, _, ts, _)| start < ts && ts <= end)
                 .map(|(job, series, ts, value)| {
                     let labels = BTreeMap::from([
                         ("__name__".to_string(), "m".to_string()),

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

@@ -29,11 +29,16 @@
 //!
 //! Composition sizes each consumer's summary build alone, for its own
 //! accuracy target. The *members* of a [`CapabilitySharingCandidate`] are summary
-//! builds of one candidate that chose the same family (same
-//! [`SketchAlgorithm`]) over the same input sub-DAG (window, `WHERE` and
-//! time range included), item and weight, grouping and filter, and whose
-//! every reader is an estimate the family answers
-//! ([`FamilyCapability::answers`]); they differ only in their sizes.
+//! builds of one candidate that chose the same family (the same
+//! [`SketchAlgorithm`](asap_types::ir::schema::SketchAlgorithm) at any
+//! size, or the same per-series exact accumulator) over the same input
+//! sub-DAG (window, `WHERE` and time range included), item and weight,
+//! grouping and filter, and whose every reader is an estimate the family
+//! answers ([`FamilyCapability::answers`]) or, for an exact accumulator,
+//! its finalization. They differ in their sizes, or, for per-series builds
+//! over one metric's scans without label matchers, in the labels those
+//! scans project: the shared build is over the scan with every member's
+//! labels, and each other member's readers keep only their own columns.
 //! [`summary_capability_pass2_candidates`] returns one Pass 2 candidate per
 //! choice, for each capability sharing candidate, of no sharing or one
 //! subset of two or more of its members
@@ -104,8 +109,7 @@ use asap_types::ir::cse::share_common_sub_dags;
 use asap_types::ir::operator::operator_properties::{Reduction, Source};
 use asap_types::ir::operator::AggIntent;
 use asap_types::ir::schema::{
-    ExactKind, FamilyCapability, FieldDataType, GroupingStrategy, MergeInputRule, SketchAlgorithm,
-    SketchKind,
+    ExactKind, FamilyCapability, FieldDataType, GroupingStrategy, MergeInputRule, SketchKind,
 };
 use asap_types::ir::SortKey;
 use asap_types::ir::{
@@ -354,13 +358,16 @@ impl Sharing<'_> {
 }
 
 /// Summary builds the summary-capability rule can make one build: builds
-/// of one family over the same input data that differ only in size, in the
-/// order a children-first walk of the roots meets them.
+/// of one family over the same input data that differ only in size, or in
+/// the labels their per-series scans project, in the order a children-first
+/// walk of the roots meets them.
 #[derive(Debug, Clone)]
 pub struct CapabilitySharingCandidate {
-    /// The sketch family every member builds.
-    pub family: SketchAlgorithm,
-    /// The summary builds (`SummaryAgg`), at least two, not all of one size.
+    /// The family the first member builds: every member builds this sketch
+    /// algorithm at some size, or this exact accumulator.
+    pub family: FieldDataType,
+    /// The summary builds (`SummaryAgg`), at least two, not all equal (in
+    /// size or in the labels their scans project).
     pub members: Vec<Rc<OperatorNode>>,
 }
 
@@ -400,7 +407,7 @@ pub fn capability_sharing_candidates(roots: &[QueryRoot]) -> Vec<CapabilityShari
         };
         match candidates
             .iter_mut()
-            .find(|c| c.family == family && same_input_data(&c.members[0], node))
+            .find(|c| same_family(&c.family, &family) && same_input_data(&c.members[0], node))
         {
             Some(candidate) => candidate.members.push(Rc::clone(node)),
             None => candidates.push(CapabilitySharingCandidate {
@@ -428,34 +435,59 @@ fn collect_children_first(
 }
 
 /// The family of a per-subpopulation sketch build whose every reader is an
-/// estimate the family answers; `None` for any other node.
+/// estimate the family answers, or of a per-series exact accumulator whose
+/// every reader finalizes it; `None` for any other node.
 fn shareable_family(
     node: &OperatorNode,
     readers: Option<&Vec<&Rc<OperatorNode>>>,
-) -> Option<SketchAlgorithm> {
-    let Some(ASAPOp::SummaryAgg {
-        family: family @ FieldDataType::Sketch(kind, GroupingStrategy::PerSubpopulationInstance),
-        ..
-    }) = node.asap()
-    else {
-        return None;
-    };
-    let capability = FamilyCapability::of(family);
-    let answered = |reader: &&Rc<OperatorNode>| {
-        matches!(
-            reader.asap(),
-            Some(ASAPOp::SummaryEstimate { query, .. }) if capability.answers(query)
-        )
-    };
+) -> Option<FieldDataType> {
     let readers = readers?;
-    readers
-        .iter()
-        .all(answered)
-        .then(|| kind.algorithm().clone())
+    match node.asap() {
+        Some(ASAPOp::SummaryAgg {
+            family: family @ FieldDataType::Sketch(_, GroupingStrategy::PerSubpopulationInstance),
+            ..
+        }) => {
+            let capability = FamilyCapability::of(family);
+            readers
+                .iter()
+                .all(|reader| {
+                    matches!(
+                        reader.asap(),
+                        Some(ASAPOp::SummaryEstimate { query, .. }) if capability.answers(query)
+                    )
+                })
+                .then(|| family.clone())
+        }
+        Some(ASAPOp::SummaryAgg {
+            family: family @ FieldDataType::ExactAggregate(..),
+            reduction: Reduction::PerEntity,
+            ..
+        }) => readers
+            .iter()
+            .all(|reader| matches!(reader.asap(), Some(ASAPOp::FinalizeExactAccumulator { .. })))
+            .then(|| family.clone()),
+        _ => None,
+    }
+}
+
+/// Whether builds of families `a` and `b` can be one build: one sketch
+/// algorithm and grouping at any sizes, or one exact accumulator.
+fn same_family(a: &FieldDataType, b: &FieldDataType) -> bool {
+    match (a, b) {
+        (FieldDataType::Sketch(a, grouping), FieldDataType::Sketch(b, grouping_b)) => {
+            a.algorithm() == b.algorithm() && grouping == grouping_b
+        }
+        _ => a == b,
+    }
 }
 
 /// Whether two builds of one family read the same data the same way: the
-/// same input sub-DAG, item and weight, grouping and filter.
+/// same input sub-DAG, item and weight, grouping and filter. Two
+/// per-series builds over one metric's samples also do when their scans
+/// project different labels ([`same_rows`]): a per-series state is keyed by
+/// the series identity, which determines every label, so a build over the
+/// scan with the most labels holds every member's states and its readers
+/// keep the labels each needs ([`share_summaries`]).
 fn same_input_data(a: &OperatorNode, b: &OperatorNode) -> bool {
     match (a.asap(), b.asap()) {
         (
@@ -476,7 +508,14 @@ fn same_input_data(a: &OperatorNode, b: &OperatorNode) -> bool {
                 ..
             }),
         ) => {
-            (Rc::ptr_eq(child, child_b) || child == child_b)
+            let same_child = Rc::ptr_eq(child, child_b)
+                || child == child_b
+                || (*reduction == Reduction::PerEntity
+                    && filter.is_none()
+                    && child.schema.has_promql_series_identity()
+                    && child_b.schema.has_promql_series_identity()
+                    && same_rows(child, child_b));
+            same_child
                 && input == input_b
                 && reduction == reduction_b
                 && grouping == grouping_b
@@ -486,6 +525,20 @@ fn same_input_data(a: &OperatorNode, b: &OperatorNode) -> bool {
     }
 }
 
+/// The names of the columns of a build's input: for a per-series build,
+/// the labels its scan projects, among others.
+fn input_column_names(build: &OperatorNode) -> Vec<&str> {
+    match build.asap() {
+        Some(ASAPOp::SummaryAgg { child, .. }) => child
+            .schema
+            .fields
+            .iter()
+            .map(|f| f.name.as_str())
+            .collect(),
+        _ => Vec::new(),
+    }
+}
+
 fn all_equal<'a>(mut nodes: impl Iterator<Item = &'a Rc<OperatorNode>>) -> bool {
     let Some(first) = nodes.next() else {
         return true;
@@ -594,13 +647,49 @@ fn share_summaries(
     shared_builds: &[Vec<Rc<OperatorNode>>],
 ) -> Result<Option<Vec<QueryRoot>>, PerComputationCandidateError> {
     let mut replaced: HashMap<*const OperatorNode, Rc<OperatorNode>> = HashMap::new();
+    // Every node that reads each build, by the build's address.
+    let mut nodes = Vec::new();
+    let mut seen = std::collections::HashSet::new();
+    for node in roots.iter().flat_map(root_operators) {
+        collect_children_first(node, &mut seen, &mut nodes);
+    }
+    let mut readers: HashMap<*const OperatorNode, Vec<&Rc<OperatorNode>>> = HashMap::new();
+    for node in &nodes {
+        for child in node.children() {
+            readers.entry(Rc::as_ptr(child)).or_default().push(node);
+        }
+    }
     for members in shared_builds {
-        let members: Vec<&Rc<OperatorNode>> = members.iter().collect();
+        let mut members: Vec<&Rc<OperatorNode>> = members.iter().collect();
+        // The build is over the input with every member's columns (the
+        // scan with the most labels); none such, no shared build.
+        let Some(host) = members.iter().position(|candidate| {
+            let columns = input_column_names(candidate);
+            members.iter().all(|member| {
+                input_column_names(member)
+                    .iter()
+                    .all(|c| columns.contains(c))
+            })
+        }) else {
+            return Ok(None);
+        };
+        members.swap(0, host);
         let Some(build) = sized_for_all(&members)? else {
             return Ok(None);
         };
         for member in members {
-            replaced.insert(Rc::as_ptr(member), Rc::clone(&build));
+            if member.schema == build.schema {
+                replaced.insert(Rc::as_ptr(member), Rc::clone(&build));
+                continue;
+            }
+            // A member over fewer labels: each of its estimates reads the
+            // shared build and keeps the member's columns, by name.
+            for reader in readers.get(&Rc::as_ptr(member)).into_iter().flatten() {
+                let Some(estimate) = relabeled_estimate(reader, &build)? else {
+                    return Ok(None);
+                };
+                replaced.insert(Rc::as_ptr(reader), estimate);
+            }
         }
     }
     let mut rebuilt: HashMap<*const OperatorNode, Rc<OperatorNode>> = HashMap::new();
@@ -623,11 +712,64 @@ fn share_summaries(
     Ok(Some(recompose_guarantees(&out)?))
 }
 
+/// `reader`, an estimate or finalization of a per-series build over fewer
+/// labels, as the same read of `build` with only `reader`'s columns, by
+/// name; `None` for another reader or a missing column.
+fn relabeled_estimate(
+    reader: &OperatorNode,
+    build: &Rc<OperatorNode>,
+) -> Result<Option<Rc<OperatorNode>>, PerComputationCandidateError> {
+    use asap_types::ir::ProjectItem;
+    let operator = match reader.asap() {
+        Some(ASAPOp::SummaryEstimate { query, .. }) => ASAPOp::SummaryEstimate {
+            summary_input: Rc::clone(build),
+            query: query.clone(),
+        },
+        Some(ASAPOp::FinalizeExactAccumulator { .. }) => ASAPOp::FinalizeExactAccumulator {
+            child: Rc::clone(build),
+        },
+        _ => return Ok(None),
+    };
+    let estimate = Rc::new(OperatorNode::new(Operator::ASAP(operator))?);
+    let mut cols = Vec::with_capacity(reader.schema.fields.len());
+    for field in &reader.schema.fields {
+        let Some(column) = estimate
+            .schema
+            .fields
+            .iter()
+            .position(|f| f.name == field.name && f.table == field.table)
+        else {
+            return Ok(None);
+        };
+        cols.push(ProjectItem {
+            alias: Some(field.name.clone()),
+            expr: ScalarExpr::Column(column),
+        });
+    }
+    Ok(Some(Rc::new(OperatorNode::new(Operator::NonASAP(
+        NonASAPOp::Project {
+            cols,
+            qualifier: None,
+            child: estimate,
+        },
+    ))?)))
+}
+
 /// One build serving every member: the first member with the members'
 /// shared size ([`size_for_all_readers`]).
 fn sized_for_all(
     members: &[&Rc<OperatorNode>],
 ) -> Result<Option<Rc<OperatorNode>>, PerComputationCandidateError> {
+    // An exact accumulator has one size: the first member serves them all.
+    if matches!(
+        members[0].asap(),
+        Some(ASAPOp::SummaryAgg {
+            family: FieldDataType::ExactAggregate(..),
+            ..
+        })
+    ) {
+        return Ok(Some(Rc::clone(members[0])));
+    }
     let kinds: Vec<&SketchKind> = members
         .iter()
         .filter_map(|member| match member.asap() {
@@ -1556,6 +1698,7 @@ mod tests {
     use crate::test_support::lower_promql;
     use asap_types::ir::properties::AccuracyTarget;
     use asap_types::ir::schema::schema_support::with_promql_series_identity;
+    use asap_types::ir::schema::SketchAlgorithm;
 
     fn roots(queries: &[&str]) -> Vec<QueryRoot> {
         queries
@@ -1896,7 +2039,9 @@ mod tests {
         assert!(own[0].1 < own[1].1);
         let sharing = capability_sharing_candidates(&roots);
         assert_eq!(sharing.len(), 1);
-        assert_eq!(sharing[0].family, SketchAlgorithm::Kll);
+        assert!(
+            matches!(&sharing[0].family, FieldDataType::Sketch(kind, _) if kind.algorithm() == &SketchAlgorithm::Kll)
+        );
         let pass2_candidates = summary_capability_pass2_candidates(&roots, &sharing).unwrap();
         assert_eq!(pass2_candidates.len(), 1);
         let [shared_builds] = pass2_candidates[0].shared_builds.as_slice() else {
@@ -2013,10 +2158,10 @@ mod tests {
             SketchAlgorithm::DDSketch,
         );
         assert_eq!(capability_sharing_candidates(&families).len(), 1);
-        assert_eq!(
-            capability_sharing_candidates(&families)[0].family,
-            SketchAlgorithm::DDSketch
-        );
+        assert!(matches!(
+            &capability_sharing_candidates(&families)[0].family,
+            FieldDataType::Sketch(kind, _) if kind.algorithm() == &SketchAlgorithm::DDSketch
+        ));
     }
 
     /// Quantiles over other groupings share no capability build.