Skip to content

Commit df545ab

Browse files
zzylolclaude
andcommitted
feat(planner): compose a whole-workload candidate from Pass 1 choices
Replays #575's composition on the #561 Pass 1 API: choice[i] indexes targets[i].choices.candidates, several targets (measures) of one aggregate are grouped, an exact accumulator becomes SummaryAgg → FinalizeExactAccumulator and a summary SummaryAgg → SummaryEstimate. Pass 1 summaries are unsized, so composition sizes each for its measure's accuracy requirement. Summary coverage is derived from the subtree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01W7qG9aFyPij5uWsyAJCxDW
1 parent 15c190f commit df545ab

3 files changed

Lines changed: 293 additions & 6 deletions

File tree

‎crates/asap-aware-mapping/src/per_computation_candidates.rs‎

Lines changed: 255 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,21 @@ use std::collections::{HashMap, HashSet};
2828
use std::rc::Rc;
2929

3030
use asap_types::ir::operator_properties::Reduction;
31-
use asap_types::ir::{NonASAPOp, OperatorNode, Predicate, QueryRoot, SchemaDerivationError};
32-
use asap_types::post_asap::{ExactKind, ExactParams, Materialization, SketchAlgorithm};
33-
use asap_types::pre_asap::{resolve_column_ref, AggIntent, ColumnRef};
31+
use asap_types::ir::{
32+
ASAPOp, NonASAPOp, Operator, OperatorNode, Predicate, QueryRoot, SchemaDerivationError,
33+
};
34+
use asap_types::post_asap::{
35+
EntityIdentity, ExactKind, ExactParams, FieldDataType, GroupingStrategy, Materialization,
36+
NonNegativeWeightProof, SketchAlgorithm, SketchKind, SketchStatistic, SummaryInputExpr,
37+
SummaryUpdate, WeightDomain,
38+
};
39+
use asap_types::pre_asap::{resolve_column_ref, AggIntent, ColumnRef, Schema};
3440
use asap_types::types::AccuracyTarget;
3541
use thiserror::Error;
3642

37-
use crate::replacement::{accuracy_budget, accuracy_target, summary_candidates};
43+
use crate::replacement::{
44+
accuracy_budget, accuracy_target, default_size_params, summary_candidates,
45+
};
3846

3947
/// One Pass 1 candidate for a computation (one measure of an aggregate).
4048
///
@@ -170,6 +178,10 @@ pub enum PerComputationCandidateError {
170178
root: Option<String>,
171179
measure: Option<usize>,
172180
},
181+
#[error("choice must name one listed candidate per target")]
182+
InvalidChoice,
183+
#[error("unsupported substitution in logical ASAP-aware optimization: {0}")]
184+
Unsupported(&'static str),
173185
}
174186

175187
/// Enumerate exact, rewrite and summary choices of one measure in stable catalog
@@ -386,6 +398,245 @@ pub fn generate_per_computation_candidates<Id: std::fmt::Debug>(
386398
Ok(PerComputationCandidates { roots, targets })
387399
}
388400

401+
/// Build one whole-workload candidate (end of Pass 1): `choice[i]` indexes
402+
/// `inventory.targets[i].choices.candidates`. A target whose chosen
403+
/// candidate is not `Exact` is replaced by `SummaryAgg` followed by
404+
/// `SummaryEstimate` (summary) or `FinalizeExactAccumulator` (exact
405+
/// accumulator). A summary is sized here for its measure's accuracy
406+
/// requirement; Pass 1 leaves it unsized. Untouched sub-DAGs keep their
407+
/// identity, so sharing between roots is preserved.
408+
pub fn compose_logical_candidate<Id: Clone>(
409+
inventory: &PerComputationCandidates<Id>,
410+
choice: &[usize],
411+
) -> Result<Vec<(Id, QueryRoot)>, PerComputationCandidateError> {
412+
if choice.len() != inventory.targets.len() {
413+
return Err(PerComputationCandidateError::InvalidChoice);
414+
}
415+
let mut chosen: Chosen = HashMap::new();
416+
for (target, &index) in inventory.targets.iter().zip(choice) {
417+
let candidate = target
418+
.choices
419+
.candidates
420+
.get(index)
421+
.ok_or(PerComputationCandidateError::InvalidChoice)?;
422+
chosen
423+
.entry(Rc::as_ptr(&target.target))
424+
.or_default()
425+
.push((target, candidate));
426+
}
427+
let mut memo = HashMap::new();
428+
inventory
429+
.roots
430+
.iter()
431+
.map(|(id, root)| {
432+
let root = match root {
433+
QueryRoot::Operator(node) => {
434+
QueryRoot::Operator(rewrite(node, &chosen, &mut memo)?)
435+
}
436+
QueryRoot::Scalar(expr) => {
437+
for node in expr.operator_refs() {
438+
rewrite(node, &chosen, &mut memo)?;
439+
}
440+
QueryRoot::Scalar(
441+
expr.map_operator_refs(&mut |node| memo[&Rc::as_ptr(node)].clone()),
442+
)
443+
}
444+
};
445+
Ok((id.clone(), root))
446+
})
447+
.collect()
448+
}
449+
450+
/// Each aggregate's targets with their chosen candidates.
451+
type Chosen<'a> =
452+
HashMap<*const OperatorNode, Vec<(&'a ComputationTarget, &'a ComputationCandidate)>>;
453+
type Memo = HashMap<*const OperatorNode, Rc<OperatorNode>>;
454+
455+
fn rewrite(
456+
node: &Rc<OperatorNode>,
457+
chosen: &Chosen<'_>,
458+
memo: &mut Memo,
459+
) -> Result<Rc<OperatorNode>, PerComputationCandidateError> {
460+
if let Some(done) = memo.get(&Rc::as_ptr(node)) {
461+
return Ok(done.clone());
462+
}
463+
for child in node.children() {
464+
rewrite(child, chosen, memo)?;
465+
}
466+
let changed = node
467+
.children()
468+
.iter()
469+
.any(|child| !Rc::ptr_eq(child, &memo[&Rc::as_ptr(child)]));
470+
let rebuilt = match chosen.get(&Rc::as_ptr(node)) {
471+
Some(targets)
472+
if targets
473+
.iter()
474+
.any(|(_, candidate)| **candidate != ComputationCandidate::Exact) =>
475+
{
476+
realize(node, targets, memo)?
477+
}
478+
_ if changed => Rc::new(node.with_new_children(|child| memo[&Rc::as_ptr(child)].clone())?),
479+
_ => node.clone(),
480+
};
481+
memo.insert(Rc::as_ptr(node), rebuilt.clone());
482+
Ok(rebuilt)
483+
}
484+
485+
fn realize(
486+
target: &OperatorNode,
487+
targets: &[(&ComputationTarget, &ComputationCandidate)],
488+
memo: &Memo,
489+
) -> Result<Rc<OperatorNode>, PerComputationCandidateError> {
490+
let Some(NonASAPOp::Aggregate {
491+
child,
492+
reduction,
493+
measures,
494+
filters,
495+
having,
496+
..
497+
}) = target.non_asap()
498+
else {
499+
return Err(PerComputationCandidateError::Unsupported(
500+
"target is not an aggregate",
501+
));
502+
};
503+
let ([intent], [(_, candidate)]) = (measures.as_slice(), targets) else {
504+
return Err(PerComputationCandidateError::Unsupported(
505+
"multi-measure aggregate",
506+
));
507+
};
508+
if filters.iter().any(Option::is_some) || having.is_some() {
509+
return Err(PerComputationCandidateError::Unsupported(
510+
"filtered or HAVING aggregate",
511+
));
512+
}
513+
let child = memo[&Rc::as_ptr(child)].clone();
514+
let (family, query) = match candidate {
515+
ComputationCandidate::ExactAccumulator { kind, params } => (
516+
FieldDataType::ExactAggregate(kind.clone(), params.clone()),
517+
None,
518+
),
519+
ComputationCandidate::Summary(algorithm) => (
520+
FieldDataType::Sketch(sized(algorithm, intent), GroupingStrategy::default()),
521+
Some(statistic(intent)?),
522+
),
523+
_ => {
524+
return Err(PerComputationCandidateError::Unsupported(
525+
"Pass 1 candidate",
526+
))
527+
}
528+
};
529+
let input = summary_update(intent, &family, reduction, &child.schema)?;
530+
let state = OperatorNode::new_shared(Operator::ASAP(ASAPOp::SummaryAgg {
531+
child: child.clone(),
532+
family,
533+
input,
534+
reduction: reduction.clone(),
535+
grouping: GroupingStrategy::default(),
536+
filter: None,
537+
}))?;
538+
let evaluation = match query {
539+
Some(query) => ASAPOp::SummaryEstimate {
540+
summary_input: state,
541+
query,
542+
},
543+
None => ASAPOp::FinalizeExactAccumulator { child: state },
544+
};
545+
Ok(OperatorNode::new_shared(Operator::ASAP(evaluation))?)
546+
}
547+
548+
/// `algorithm` sized for `intent`'s accuracy requirement.
549+
fn sized(algorithm: &SketchAlgorithm, intent: &AggIntent) -> SketchKind {
550+
let (epsilon, delta) =
551+
accuracy_budget(accuracy_target(intent).unwrap_or(&AccuracyTarget::Exact));
552+
SketchKind::new(
553+
algorithm.clone(),
554+
default_size_params(algorithm.clone(), intent, epsilon, delta),
555+
)
556+
}
557+
558+
/// What each input row contributes, following the legacy realization rules:
559+
/// heap sketches rank series identities, frequency sketches count values, and
560+
/// every other family reads the measure's input column.
561+
fn summary_update(
562+
intent: &AggIntent,
563+
family: &FieldDataType,
564+
reduction: &Reduction,
565+
child: &Schema,
566+
) -> Result<SummaryUpdate, PerComputationCandidateError> {
567+
let algorithm = match family {
568+
FieldDataType::Sketch(kind, _) => Some(kind.algorithm()),
569+
_ => None,
570+
};
571+
let weight = crate::replacement::summarised_input(intent, child).map_err(|_| {
572+
PerComputationCandidateError::Unsupported("input column outside child schema")
573+
})?;
574+
Ok(match (intent, algorithm) {
575+
(AggIntent::TopK { .. }, Some(_)) => {
576+
// SQL rows carry no implicit series identity to rank.
577+
if child.closed {
578+
return Err(PerComputationCandidateError::Unsupported(
579+
"Top-K item identity for closed schemas",
580+
));
581+
}
582+
if reduction.group_keys().is_some_and(|keys| keys.is_without()) {
583+
return Err(PerComputationCandidateError::Unsupported(
584+
"Top-K partitions given by `without`",
585+
));
586+
}
587+
let excluding = reduction
588+
.group_keys()
589+
.into_iter()
590+
.flat_map(|keys| keys.iter())
591+
.filter_map(|&index| child.fields.get(index))
592+
.map(crate::replacement::column_ref)
593+
.collect();
594+
SummaryUpdate {
595+
item: Some(SummaryInputExpr::EntityIdentity(
596+
EntityIdentity::PromqlLabelSet { excluding },
597+
)),
598+
weight: SummaryInputExpr::Column(ColumnRef::SampleValue),
599+
// Not proven non-negative; selection decides whether CMS is admissible.
600+
weight_domain: WeightDomain::UnknownOrSigned,
601+
}
602+
}
603+
(_, Some(SketchAlgorithm::UnivMon))
604+
| (AggIntent::Count { .. }, Some(SketchAlgorithm::Cms | SketchAlgorithm::CountSketch)) => {
605+
SummaryUpdate {
606+
item: Some(weight),
607+
weight: SummaryInputExpr::Constant(1.0),
608+
weight_domain: WeightDomain::NonNegative {
609+
proof: NonNegativeWeightProof::UnitCount,
610+
},
611+
}
612+
}
613+
_ => SummaryUpdate {
614+
item: None,
615+
weight,
616+
weight_domain: WeightDomain::UnknownOrSigned,
617+
},
618+
})
619+
}
620+
621+
fn statistic(intent: &AggIntent) -> Result<SketchStatistic, PerComputationCandidateError> {
622+
Ok(match intent {
623+
AggIntent::Quantile { q, .. } => SketchStatistic::Quantile { q: *q },
624+
AggIntent::Cardinality { .. } => SketchStatistic::Cardinality,
625+
AggIntent::FrequencyL2 { .. } => SketchStatistic::FrequencyL2,
626+
AggIntent::FrequencyEntropy { .. } => SketchStatistic::FrequencyEntropy,
627+
AggIntent::TopK { k, .. } => SketchStatistic::TopK { k: *k },
628+
AggIntent::Count { .. } => SketchStatistic::PointCount {
629+
key: ColumnRef::SampleValue,
630+
value: None,
631+
},
632+
_ => {
633+
return Err(PerComputationCandidateError::Unsupported(
634+
"sketch evaluation",
635+
))
636+
}
637+
})
638+
}
639+
389640
#[cfg(test)]
390641
mod tests {
391642
use super::*;

‎crates/asap-aware-mapping/src/replacement.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3828,7 +3828,7 @@ fn summarised_column(intent: &AggIntent, child_schema: &Schema) -> ColumnRef {
38283828
}
38293829
}
38303830

3831-
fn column_ref(column: &Field) -> ColumnRef {
3831+
pub(crate) fn column_ref(column: &Field) -> ColumnRef {
38323832
match &column.table {
38333833
Some(t) => ColumnRef::Qualified {
38343834
table: t.clone(),
@@ -3846,7 +3846,7 @@ fn column_ref(column: &Field) -> ColumnRef {
38463846
/// A tuple leg outside the child schema is an error rather than
38473847
/// [`summarised_column`]'s sample-value fallback: a leg has no sample-value
38483848
/// reading, and silently dropping one would under-count.
3849-
fn summarised_input(
3849+
pub(crate) fn summarised_input(
38503850
intent: &AggIntent,
38513851
child_schema: &Schema,
38523852
) -> Result<SummaryInputExpr, RealizationError> {

‎crates/asap-aware-mapping/tests/per_computation_candidates.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -850,3 +850,39 @@ fn without_and_per_entity_having_stays_above() {
850850
);
851851
}
852852
}
853+
854+
/// A workload candidate replaces only chosen targets, derives the summary's
855+
/// coverage from its subtree, and keeps unchosen plans identical.
856+
#[test]
857+
fn composed_candidate_replaces_chosen_target_with_summary_evaluation() {
858+
use asap_aware_mapping::per_computation_candidates::compose_logical_candidate;
859+
let producer = aggregate(AggIntent::Cardinality {
860+
cols: vec![0],
861+
accuracy: approximate(),
862+
});
863+
let inventory =
864+
generate_per_computation_candidates(vec![(0, QueryRoot::Operator(producer.clone()))])
865+
.unwrap();
866+
let exact = compose_logical_candidate(&inventory, &[0]).unwrap();
867+
assert!(matches!(&exact[0].1, QueryRoot::Operator(node) if Rc::ptr_eq(node, &producer)));
868+
869+
let hll = compose_logical_candidate(&inventory, &[1]).unwrap();
870+
let QueryRoot::Operator(estimate) = &hll[0].1 else {
871+
panic!("operator root expected")
872+
};
873+
let Some(ASAPOp::SummaryEstimate { summary_input, .. }) = estimate.asap() else {
874+
panic!("summary evaluation expected")
875+
};
876+
// The state takes every row of its source: one unrestricted selection box.
877+
let coverage = summary_input.coverage().unwrap();
878+
assert!(OperatorNode::reachable(&coverage.definition).iter().any(|node| matches!(
879+
node.non_asap(),
880+
Some(NonASAPOp::Scan { source: Source::Table { table_ref }, .. }) if table_ref == "flows"
881+
)));
882+
assert_eq!(coverage.selection.len(), 1);
883+
assert!(
884+
coverage.selection[0].columns.is_empty() && coverage.selection[0].relative_time.is_none()
885+
);
886+
hll[0].1.validate_structure().unwrap();
887+
assert!(compose_logical_candidate(&inventory, &[99]).is_err());
888+
}

0 commit comments

Comments
 (0)