refactor(stage1): top-k reuse and finer-aggregate reads as one operator over groups (#1203 2a-3)

crates/logical-asap-optimization/src/pass2/memo/rules/finer_aggregate.rs

@@ -1,12 +1,12 @@
 //! Read a SQL aggregate from a finer aggregate's result
 //! ([`match_finer`](crate::pass2::finer_aggregate::match_finer)): filter the finer
 //! rows, then merge them into the query's groups with each measure's rollup.
-use super::reuse::append_reader;
+use super::reuse::add_candidate_over_groups;
 use crate::pass2::finer_aggregate::match_finer;
 use crate::pass2::memo::state_groups::Pass2RuleError;
 use asap_types::ir::operator::{AggIntent, Reduction};
 use asap_types::ir::{NonASAPOp, Operator, OperatorNode, OperatorResultKind};
-use asap_types::stages::{Memo, OriginOfCandidateForMemoGroup};
+use asap_types::stages::{Memo, MemoGroupId, OriginOfCandidateForMemoGroup};
 use std::collections::HashMap;
 use std::rc::Rc;
 
@@ -167,9 +167,9 @@ fn add_read(memo: &mut Memo, owner: usize, finer: usize) -> Result<usize, Pass2R
         return Ok(0);
     }
     let reader = Rc::new(OperatorNode::with_schema(reader.operator.clone(), schema));
-    Ok(usize::from(append_reader(
+    Ok(usize::from(add_candidate_over_groups(
         memo,
-        owner,
+        MemoGroupId(owner as u32),
         &original,
         reader,
         OriginOfCandidateForMemoGroup::Rollup,

crates/logical-asap-optimization/src/pass2/memo/rules/topk.rs

@@ -1,7 +1,8 @@
 //! Exact ranking reuse with explicit exact-input groups. No global selection
 //! flags are required: only exact implementations can supply these inputs.
-use super::reuse::{append_producer, append_reader, eligible, replace};
+use super::reuse::{add_candidate_over_groups, eligible, find_or_add_producer, replace};
 use super::shared_producer::{add_shared_producer_candidates, Reader, SharedProducer};
+use crate::find_or_add_group::find_or_add_group;
 use crate::pass2::asap_aware_cse::topk_reuse_candidates;
 use crate::pass2::memo::state_groups::{boundary_ptrs, nodes_within, Pass2RuleError};
 use asap_types::ir::operator::{AggIntent, Reduction};
@@ -123,25 +124,23 @@ impl SharedProducer for TopKRankings {
         let sort = covering.ranking.children()[0]
             .with_new_children(|_| Rc::clone(&input))
             .map_err(|error| Pass2RuleError::ReuseRule(error.to_string()))?;
-        let sort_group = append_producer(
+        let sort_group = find_or_add_producer(
             memo,
-            Rc::new(sort),
+            &Rc::new(sort),
             &queries,
-            OriginOfCandidateForMemoGroup::TopKReuse,
-            1,
-        )?;
+            &OriginOfCandidateForMemoGroup::TopKReuse,
+        );
         let sort = Rc::clone(&memo.groups()[sort_group.0 as usize].logical_node);
         let ranking = covering
             .ranking
             .with_new_children(|_| Rc::clone(&sort))
             .map_err(|error| Pass2RuleError::ReuseRule(error.to_string()))?;
-        let ranking_group = append_producer(
+        let ranking_group = find_or_add_producer(
             memo,
-            Rc::new(ranking),
+            &Rc::new(ranking),
             &queries,
-            OriginOfCandidateForMemoGroup::TopKReuse,
-            1,
-        )?;
+            &OriginOfCandidateForMemoGroup::TopKReuse,
+        );
         let ranking = Rc::clone(&memo.groups()[ranking_group.0 as usize].logical_node);
         let owner = reader.group;
         let logical = Rc::clone(&memo.groups()[owner].logical_node);
@@ -176,9 +175,9 @@ impl SharedProducer for TopKRankings {
             child: limited,
         }))
         .map_err(|error| Pass2RuleError::ReuseRule(error.to_string()))?;
-        Ok(usize::from(append_reader(
+        Ok(usize::from(add_candidate_over_groups(
             memo,
-            owner,
+            MemoGroupId(owner as u32),
             &original,
             reader,
             OriginOfCandidateForMemoGroup::TopKReuse,
@@ -251,17 +250,29 @@ fn exact_input_group(
         AccuracyGuarantee::exact("only exact implementations and exact inputs"),
     )));
     let queries = memo.groups()[group.0 as usize].queries.clone();
-    let id = append_producer(
+    // The exact guarantee is the group's identity: it keeps the group apart
+    // from `group`, whose candidates may be approximate.
+    let groups_before = memo.groups().len();
+    let id = find_or_add_group(
         memo,
-        node,
+        &node,
         &queries,
-        OriginOfCandidateForMemoGroup::TopKReuse,
+        &OriginOfCandidateForMemoGroup::TopKReuse,
         1,
-    )?;
+    );
     for candidate in &mut candidates {
         candidate.fused_groups = Vec::new();
     }
-    memo.candidates[id.0 as usize] = candidates;
+    if id.0 as usize >= groups_before {
+        // Added: the exact candidates replace the group's own candidate.
+        memo.candidates[id.0 as usize] = candidates;
+    } else {
+        // Found (another group's exact input equals it): keep its
+        // candidates and add the new ones.
+        for candidate in candidates {
+            memo.add_candidate(id, candidate);
+        }
+    }
     known.insert(group, Some(id));
     Ok(Some(id))
 }