Skip to content

refactor(stage1): rollups as one operator over groups (#1203 2a-4) #1221

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zzylol wants to merge 2 commits into from
Closed

refactor(stage1): rollups as one operator over groups (#1203 2a-4) #1221

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
71d7557
refactor(stage1): rollups as one operator over groups
zzylol Nov 5, 2026
947d5b3
test(integration): rollup plans follow each query's path of groups do…
zzylol Nov 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 47 additions & 17 deletions crates/integration-tests/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -539,28 +539,19 @@ pub mod planning {
if !is_build(&group.logical_node) || group.reached_by_queries().len() < queries {
continue;
}
// Each query's consumer of the state, and its own summary candidate
// above, which reads the consumer's group (Pass 1 adds the consumer
// as a group beneath the query's aggregate).
// Each query's own summary candidate, with the candidates on a
// path of groups from it down to the state (Pass 1 and the
// sharing rules add the consumers between them as groups).
let reading: Option<std::collections::BTreeMap<usize, usize>> = (0..queries)
.map(|q| {
let reader = own_choices(memo, q, |c| {
c.input_groups.iter().any(|child| child.0 as usize == state)
own_builds[q].iter().find_map(|&(group, index)| {
let mut path = path_to_state(memo, group, index, state)?;
path.push((group, index));
Some(path)
})
.into_iter()
.next()?;
let top = own_builds[q].iter().copied().find(|&(group, index)| {
asap_logical_asap_optimization::test_support::groups_read_through_added_groups(
memo,
&memo.candidates[group][index].input_groups,
)
.added_groups
.contains(&asap_types::stages::MemoGroupId(reader.0 as u32))
})?;
Some([top, reader])
})
.collect::<Option<Vec<_>>>()
.map(|pairs| pairs.into_iter().flatten().collect());
.map(|paths| paths.into_iter().flatten().collect());
if let Some(choice) = reading {
plans.push((
format!("shared{state}"),
Expand Down Expand Up @@ -604,6 +595,45 @@ pub mod planning {
(roots, separate, cheapest(&plans[separate_plans..]))
}

/// The candidates, one per group, on a path from candidate `index` of
/// `group` down to group `state` (excluded, and the start excluded):
/// each reads the next one's group. `None` when no path reaches `state`.
fn path_to_state(
memo: &asap_types::stages::Memo,
group: usize,
index: usize,
state: usize,
) -> Option<Vec<(usize, usize)>> {
fn below(
memo: &asap_types::stages::Memo,
group: usize,
state: usize,
seen: &mut std::collections::HashSet<usize>,
) -> Option<Vec<(usize, usize)>> {
if group == state {
return Some(Vec::new());
}
if !seen.insert(group) {
return None;
}
memo.candidates[group]
.iter()
.enumerate()
.find_map(|(index, candidate)| {
candidate.input_groups.iter().find_map(|input| {
let mut path = below(memo, input.0 as usize, state, seen)?;
path.push((group, index));
Some(path)
})
})
}
let mut seen = std::collections::HashSet::from([group]);
memo.candidates[group][index]
.input_groups
.iter()
.find_map(|input| below(memo, input.0 as usize, state, &mut seen))
}

/// The summary builds and rollups (regrouping merges) `roots` read, each
/// once.
pub fn builds_and_rollups(roots: &[Rc<OperatorNode>]) -> (usize, usize) {
Expand Down
8 changes: 4 additions & 4 deletions crates/logical-asap-optimization/src/added_groups.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,17 @@ use asap_types::stages::{Memo, MemoGroupId};
/// The groups some inputs read through the groups added beneath Stage 0's
/// groups ([`groups_read_through_added_groups`]).
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct GroupsReadThroughAddedGroups {
pub(crate) struct GroupsReadThroughAddedGroups {
/// The added groups passed through, in the order first reached.
pub added_groups: Vec<MemoGroupId>,
pub(crate) added_groups: Vec<MemoGroupId>,
/// The Stage 0 groups reached, in the order first reached.
pub stage0_groups: Vec<MemoGroupId>,
pub(crate) stage0_groups: Vec<MemoGroupId>,
}

/// The groups `inputs` read, walking through every added group (each read
/// through all its candidates' input groups) and stopping at Stage 0's
/// groups.
pub fn groups_read_through_added_groups(
pub(crate) fn groups_read_through_added_groups(
memo: &Memo,
inputs: &[MemoGroupId],
) -> GroupsReadThroughAddedGroups {
Expand Down
169 changes: 6 additions & 163 deletions crates/logical-asap-optimization/src/pass2/memo/rules/reuse.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,138 +1,19 @@
//! Shared boundary handling for producer-level reuse rules. Rule outputs retain
//! canonical input identities after guarantee recomposition rebuilds their DAGs.
//! What the producer-level reuse rules share: which candidates they read
//! ([`eligible`]), and how they add what they derive, as one operator over
//! groups ([`add_candidate_over_groups`], [`find_or_add_producer`]).
use crate::find_or_add_group::find_or_add_group;
use crate::operator_over_groups::{operator_over_groups, without_exact_guarantees_below};
use crate::pass2::memo::state_groups::{boundary_groups, Pass2RuleError};
use crate::pass2::memo::state_groups::Pass2RuleError;
use asap_types::ir::OperatorNode;
use asap_types::stages::MemoGroupId;
use asap_types::stages::{CandidateForMemoGroup, Memo, OriginOfCandidateForMemoGroup};
use asap_types::stages::{MemoGroup, MemoGroupId};
use std::collections::{HashMap, HashSet};
use std::collections::HashMap;
use std::rc::Rc;

pub(in crate::pass2::memo) fn eligible(candidate: &CandidateForMemoGroup) -> bool {
candidate.accuracy_budget_parts > 0
}

pub(in crate::pass2::memo) fn restore_inputs(
node: &Rc<OperatorNode>,
memo: &Memo,
inputs: &[MemoGroupId],
) -> Rc<OperatorNode> {
// These are the same declared logical input groups, rebound after a rule
// cloned the DAG. Value guarantees are recomposed from selected producers
// during reconstruction; state inputs additionally require the same error
// contract because the intermediate state's guarantee is part of its API.
let canonical: Vec<_> = inputs
.iter()
.flat_map(|id| OperatorNode::reachable(&memo.groups()[id.0 as usize].logical_node))
.collect();
fn visit(
node: &Rc<OperatorNode>,
canonical: &[Rc<OperatorNode>],
done: &mut HashMap<*const OperatorNode, Rc<OperatorNode>>,
) -> Rc<OperatorNode> {
if let Some(found) = done.get(&Rc::as_ptr(node)) {
return Rc::clone(found);
}
let operator = node
.operator
.map_children(|child| visit(child, canonical, done));
let rebuilt = Rc::new(OperatorNode {
operator,
..(**node).clone()
});
let result = canonical
.iter()
.find(|known| {
same_node(
known,
&rebuilt,
known.result_kind == asap_types::ir::OperatorResultKind::State,
)
})
.cloned()
.unwrap_or(rebuilt);
done.insert(Rc::as_ptr(node), Rc::clone(&result));
result
}
visit(node, &canonical, &mut HashMap::new())
}

pub(in crate::pass2::memo) fn same_node(
left: &OperatorNode,
right: &OperatorNode,
guarantees: bool,
) -> bool {
left.operator.map_children(|_| ()) == right.operator.map_children(|_| ())
&& left.schema == right.schema
&& left.result_kind == right.result_kind
&& left.timing == right.timing
&& left.materialization == right.materialization
&& left.children().len() == right.children().len()
&& left
.children()
.iter()
.zip(right.children())
.all(|(a, b)| Rc::ptr_eq(a, b))
&& serde_json::to_string(&(
left.operator.map_children(|_| ()),
&left.schema,
left.result_kind,
left.timing,
left.materialization,
))
.ok()
== serde_json::to_string(&(
right.operator.map_children(|_| ()),
&right.schema,
right.result_kind,
right.timing,
right.materialization,
))
.ok()
&& (!guarantees || left.guarantee == right.guarantee)
}

pub(in crate::pass2::memo) fn append_producer(
memo: &mut Memo,
node: Rc<OperatorNode>,
queries: &[usize],
origin: OriginOfCandidateForMemoGroup,
accuracy_budget_parts: usize,
) -> Result<MemoGroupId, Pass2RuleError> {
// Reuse only dedicated producer groups. A logical result group can contain
// other implementations and cannot enforce this producer's exact contract.
let root_groups = asap_types::stages::root_groups(&memo.query_roots);
if let Some((index, _)) = memo.groups().iter().enumerate().find(|(index, group)| {
!root_groups.contains(&MemoGroupId(*index as u32))
&& memo.candidates[*index].len() == 1
&& memo.candidates[*index][0].origin == origin
&& memo.candidates[*index][0].accuracy_budget_parts == accuracy_budget_parts
&& same_node(&group.logical_node, &node, true)
}) {
return Ok(MemoGroupId(index as u32));
}
// The new group's id must fit a `MemoGroupId`.
u32::try_from(memo.groups().len()).map_err(|_| Pass2RuleError::TooManyGroups)?;
let node = Rc::new((*node).clone());
let children = boundary_groups(memo, &node, |_| true);
Ok(memo.push_group(
MemoGroup {
logical_node: Rc::clone(&node),
child_groups: children.clone(),
parents: Vec::new(),
queries: queries.to_vec(),
},
vec![CandidateForMemoGroup {
fused_groups: Vec::new(),
root_node: node,
input_groups: children,
origin,
accuracy_budget_parts,
}],
))
}

/// Whether `actual` has `expected`'s output: the same result kind, fields,
/// time index and unique keys, and a closed label set where `expected`'s is
/// closed (a projection may enumerate every output of an open schema).
Expand Down Expand Up @@ -219,31 +100,6 @@ pub(in crate::pass2::memo) fn find_or_add_producer(
find_or_add_group(memo, &node, queries, origin, 1)
}

pub(in crate::pass2::memo) fn append_reader(
memo: &mut Memo,
owner: usize,
original: &CandidateForMemoGroup,
node: Rc<OperatorNode>,
origin: OriginOfCandidateForMemoGroup,
) -> Result<bool, Pass2RuleError> {
check_output_contract(original, &node, &origin)?;
let children = boundary_groups(memo, &node, |group| group.0 as usize != owner);
if memo.candidates[owner].iter().any(|known| {
known.input_groups == children
&& known.accuracy_budget_parts == original.accuracy_budget_parts
&& *known.root_node == *node
&& serde_json::to_string(&known.root_node).ok() == serde_json::to_string(&node).ok()
}) {
return Ok(false);
}
let mut candidate = original.clone();
candidate.root_node = node;
candidate.input_groups = children;
candidate.origin = origin;
memo.candidates[owner].push(candidate);
Ok(true)
}

pub(in crate::pass2::memo) fn replace(
node: &Rc<OperatorNode>,
replacements: &HashMap<*const OperatorNode, Rc<OperatorNode>>,
Expand Down Expand Up @@ -278,16 +134,3 @@ pub(in crate::pass2::memo) fn replace(
}
visit(node, replacements, &mut HashMap::new())
}

pub(in crate::pass2::memo) fn combined_inputs(
left: &CandidateForMemoGroup,
right: &CandidateForMemoGroup,
) -> Vec<MemoGroupId> {
let mut seen = HashSet::new();
left.input_groups
.iter()
.chain(&right.input_groups)
.copied()
.filter(|id| seen.insert(*id))
.collect()
}
Loading
Loading