From 0a5d6b1a09adbb8eea6b84235d014463dc8fb2b0 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 4 Nov 2026 10:21:30 -0500 Subject: [PATCH 1/2] refactor(stage1): remove Pass 2's production imports of Pass 1 (#1191) - Pass 2's rewrite helpers return a private CseError instead of Pass 1's PerComputationCandidateError, which loses TooManyCandidates (only Pass 2 produced it). Recompose wraps SchemaDerivationError directly. - column_ref, exact_family and partitions_over move to a crate-private ir_construction module both passes use. - rolled_up builds its Sum/Min/Max exact accumulator inline instead of calling Pass 1's exact_accumulator. - pass_layering walks each pass recursively and forbids pass2 -> pass1 too. Co-Authored-By: Claude Opus 5.5 --- .../src/ir_construction.rs | 64 +++++++++++ crates/logical-asap-optimization/src/lib.rs | 2 + .../src/pass1/per_computation_candidates.rs | 80 ++----------- .../src/pass1/realization.rs | 14 +-- .../src/pass2/asap_aware_cse.rs | 105 +++++++++++------- .../src/pass2/memo/enumeration.rs | 7 +- .../tests/pass_layering.rs | 53 ++++++--- 7 files changed, 187 insertions(+), 138 deletions(-) create mode 100644 crates/logical-asap-optimization/src/ir_construction.rs diff --git a/crates/logical-asap-optimization/src/ir_construction.rs b/crates/logical-asap-optimization/src/ir_construction.rs new file mode 100644 index 000000000..c5129537c --- /dev/null +++ b/crates/logical-asap-optimization/src/ir_construction.rs @@ -0,0 +1,64 @@ +//! IR-building helpers both Stage 1 passes use. Neither pass imports the +//! other (`tests/pass_layering.rs`), so what they share lives here. + +use asap_types::ir::operator::Reduction; +use asap_types::ir::scalar::expr_ir::ColumnRef; +use asap_types::ir::schema::{ExactKind, ExactParams, Field, FieldDataType, Schema}; + +/// A reference to `column`: qualified by its table when it has one. +pub(crate) fn column_ref(column: &Field) -> ColumnRef { + match &column.table { + Some(t) => ColumnRef::Qualified { + table: t.clone(), + name: column.name.clone(), + }, + None => ColumnRef::Named(column.name.clone()), + } +} + +/// `FieldDataType::ExactAggregate` of accumulator `kind`, with its parameters +/// (none has any, so they mirror `kind`). +pub(crate) fn exact_family(kind: ExactKind) -> FieldDataType { + let params = match kind { + ExactKind::Sum => ExactParams::Sum, + ExactKind::Count => ExactParams::Count, + ExactKind::Min => ExactParams::Min, + ExactKind::Max => ExactParams::Max, + ExactKind::Rate => ExactParams::Rate, + ExactKind::IRate => ExactParams::IRate, + ExactKind::Increase => ExactParams::Increase, + }; + FieldDataType::ExactAggregate(kind, params) +} + +/// `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. +pub(crate) 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 column = |schema: &Schema, index: usize| column_ref(&schema.fields[index]); + let keys = keys + .keys() + .iter() + .map(|&index| { + inner.fields.get(index)?; + let key = column(inner, index); + let mut matches = (0..input.fields.len()).filter(|&i| column(input, i) == key); + match (matches.next(), matches.next()) { + (Some(i), None) => Some(i), + _ => None, + } + }) + .collect::>()?; + Some(Reduction::by(keys)) +} diff --git a/crates/logical-asap-optimization/src/lib.rs b/crates/logical-asap-optimization/src/lib.rs index 213790d6d..b4c709862 100644 --- a/crates/logical-asap-optimization/src/lib.rs +++ b/crates/logical-asap-optimization/src/lib.rs @@ -62,6 +62,8 @@ /// The deployment's capabilities: supported summaries and what it retains. mod deployment_capabilities; +/// IR-building helpers both Stage 1 passes use. +mod ir_construction; mod pass1; mod pass2; /// Stage 0: the memo groups of the workload's query roots. diff --git a/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs b/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs index e18050e22..9e6867d16 100644 --- a/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs @@ -113,9 +113,8 @@ use asap_types::ir::scalar::expr_ir::ArithmeticOpKind; use asap_types::ir::scalar::{resolve_column_ref, ColumnRef, CompareOpKind, ScalarValue}; use asap_types::ir::schema::{DataType, Schema}; use asap_types::ir::schema::{ - EntityIdentity, ExactKind, ExactParams, FieldDataType, GroupingStrategy, - NonNegativeWeightProof, SketchAlgorithm, SketchKind, StatisticKind, SummaryInputExpr, - SummaryUpdate, WeightDomain, + EntityIdentity, ExactKind, FieldDataType, GroupingStrategy, NonNegativeWeightProof, + SketchAlgorithm, SketchKind, StatisticKind, SummaryInputExpr, SummaryUpdate, WeightDomain, }; use asap_types::ir::{ ASAPOp, BinaryOperator, ExprSemantics, NonASAPOp, Operator, OperatorNode, OperatorResultKind, @@ -137,27 +136,13 @@ use asap_types::workload::bounds::{declared_non_negative, series_read, values_pe use asap_types::workload::DataWorkload; use crate::deployment_capabilities::{summary_supported, SupportedSummary}; +use crate::ir_construction::{exact_family, partitions_over}; use crate::pass1::realization::summary_candidates; use asap_types::stages::{ ComputationTarget, DagOfLogicalMemoGroups, FusedComputation, HavingEntry, HavingPlacement, HowToComputeAggregateMeasure, MemoGroupId, PruneReason, }; -/// `FieldDataType::ExactAggregate` of accumulator `kind`, with its parameters -/// (none has any, so they mirror `kind`). -fn exact_family(kind: ExactKind) -> FieldDataType { - let params = match kind { - ExactKind::Sum => ExactParams::Sum, - ExactKind::Count => ExactParams::Count, - ExactKind::Min => ExactParams::Min, - ExactKind::Max => ExactParams::Max, - ExactKind::Rate => ExactParams::Rate, - ExactKind::IRate => ExactParams::IRate, - ExactKind::Increase => ExactParams::Increase, - }; - FieldDataType::ExactAggregate(kind, params) -} - /// Unranked candidates of one measure and the candidates pruned for it. #[derive(Debug, Clone, PartialEq)] pub(super) struct ComputationChoices { @@ -702,18 +687,6 @@ pub enum PerComputationCandidateError { InvalidChoice, #[error("unsupported substitution in logical ASAP-aware optimization: {0}")] Unsupported(&'static str), - /// A Pass 2 rule found more than `max` candidates of one kind, whose - /// 2^n subsets it would enumerate: the workload is too large to plan - /// this way, not bad input. - #[error("found {found} {counted}, more than {max}")] - TooManyCandidates { - /// What was counted, e.g. "rollup candidates". - counted: &'static str, - /// How many were found, more than `max`. - found: usize, - /// The cap that was exceeded. - max: usize, - }, } /// Enumerate exact, rewrite and summary choices of one measure in stable catalog @@ -1344,40 +1317,6 @@ fn whole_expression_input( 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. -pub(crate) 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 column = |schema: &Schema, index: usize| { - crate::pass1::realization::column_ref(&schema.fields[index]) - }; - let keys = keys - .keys() - .iter() - .map(|&index| { - inner.fields.get(index)?; - let key = column(inner, index); - let mut matches = (0..input.fields.len()).filter(|&i| column(input, i) == key); - match (matches.next(), matches.next()) { - (Some(i), None) => Some(i), - _ => None, - } - }) - .collect::>()?; - Some(Reduction::by(keys)) -} - /// `update` with the open requirement that its weights are samples of /// `metric` ([`WeightDomain::RawDataSamplesOfMetric`]), when it is a keyed heap /// update whose weights are not otherwise proven. @@ -2348,7 +2287,7 @@ fn counted_column( schema .fields .get(id) - .map(crate::pass1::realization::column_ref) + .map(crate::ir_construction::column_ref) }; filter .map_or_else(Vec::new, |filter| not_null(&filter.0)) @@ -2411,7 +2350,7 @@ fn summary_update( .enumerate() .filter(|(index, _)| *index != value && !groups.contains(index)) .map(|(_, field)| { - SummaryInputExpr::Column(crate::pass1::realization::column_ref(field)) + SummaryInputExpr::Column(crate::ir_construction::column_ref(field)) }) .collect(); if items.is_empty() { @@ -2424,7 +2363,7 @@ fn summary_update( let excluding = groups .iter() .filter_map(|&index| child.fields.get(index)) - .map(crate::pass1::realization::column_ref) + .map(crate::ir_construction::column_ref) .collect(); SummaryInputExpr::EntityIdentity(EntityIdentity::PromqlLabelSet { excluding }) }; @@ -3640,7 +3579,7 @@ mod tests { }; assert_eq!( *family, - FieldDataType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) + FieldDataType::ExactAggregate(ExactKind::Sum, asap_types::ir::schema::ExactParams::Sum) ); assert!(matches!( child.non_asap(), @@ -4705,7 +4644,10 @@ mod tests { }; assert_eq!( sum_family, - &FieldDataType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) + &FieldDataType::ExactAggregate( + ExactKind::Sum, + asap_types::ir::schema::ExactParams::Sum + ) ); assert!(!leaf.contains_asap()); assert!(child diff --git a/crates/logical-asap-optimization/src/pass1/realization.rs b/crates/logical-asap-optimization/src/pass1/realization.rs index 1965c20f3..7691730dd 100644 --- a/crates/logical-asap-optimization/src/pass1/realization.rs +++ b/crates/logical-asap-optimization/src/pass1/realization.rs @@ -10,7 +10,7 @@ use asap_types::ir::scalar::expr_ir::ColumnRef; use asap_types::ir::schema::topk::ranked_column; use asap_types::ir::schema::ColumnId; use asap_types::ir::schema::{ - EntityIdentity, FamilyCapability, Field, FieldDataType, NonNegativeWeightProof, Schema, + EntityIdentity, FamilyCapability, FieldDataType, NonNegativeWeightProof, Schema, SketchAlgorithm, SummaryInputExpr, SummaryUpdate, WeightDomain, }; use asap_types::ir::{NonASAPOp, OperatorNode}; @@ -18,6 +18,8 @@ use asap_types::ir::{NonASAPOp, OperatorNode}; use asap_types::ir::{Operator, ScalarExpr}; use std::rc::Rc; +use crate::ir_construction::column_ref; + use asap_accuracy::sizing::sketch_statistic; /// The sketch families the planner sizes and offers, in the order Pass 1 @@ -92,16 +94,6 @@ pub(crate) fn ranked_value(child: &OperatorNode) -> Result<(ColumnId, ColumnRef) Ok((ranked, column)) } -pub(crate) fn column_ref(column: &Field) -> ColumnRef { - match &column.table { - Some(t) => ColumnRef::Qualified { - table: t.clone(), - name: column.name.clone(), - }, - None => ColumnRef::Named(column.name.clone()), - } -} - /// What the summary consumes per input row. An intent that reads one column (or /// none) feeds that column; `COUNT(DISTINCT a, b)` feeds the whole tuple as one /// item, so the distinct-count sketch hashes `(a, b)` rather than `a` — the diff --git a/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs b/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs index 55bf79e40..dc3ed2a56 100644 --- a/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs +++ b/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs @@ -110,9 +110,10 @@ use std::rc::Rc; use asap_types::ir::operator::operator_properties::{Reduction, Source}; use asap_types::ir::operator::AggIntent; +use asap_types::ir::scalar::resolve_column_ref; use asap_types::ir::schema::{ ExactKind, FamilyCapability, FieldDataType, GroupingStrategy, MergeInputRule, MergedError, - Schema, SketchKind, PROMQL_SERIES_IDENTITY, + Schema, SketchKind, SummaryInputExpr, SummaryUpdate, WeightDomain, PROMQL_SERIES_IDENTITY, }; use asap_types::ir::ScalarExpr; use asap_types::ir::SortKey; @@ -121,10 +122,35 @@ use asap_types::ir::{ TimeRangeKind, }; -use crate::pass1::per_computation_candidates::PerComputationCandidateError; +use crate::ir_construction::{column_ref, exact_family, partitions_over}; use asap_accuracy::estimators::size_for_all_readers; use asap_accuracy::propagation::recompose_guarantees; +use asap_types::ir::SchemaDerivationError; + +/// Why a sharing rewrite builds no candidate. The rules report it as text +/// ([`Pass2RuleError`](crate::pass2::memo::state_groups::Pass2RuleError)) or +/// drop the candidate. +#[derive(Debug, thiserror::Error)] +pub enum CseError { + #[error(transparent)] + Structure(#[from] SchemaDerivationError), + #[error("unsupported substitution in logical ASAP-aware optimization: {0}")] + Unsupported(&'static str), + /// A rule found more than `max` candidates of one kind, whose 2^n subsets + /// it would enumerate: the workload is too large to plan this way, not + /// bad input. + #[error("found {found} {counted}, more than {max}")] + TooManyCandidates { + /// What was counted, e.g. "rollup choices". + counted: &'static str, + /// How many were found, more than `max`. + found: usize, + /// The most the rule enumerates. + max: usize, + }, +} + /// `roots` with the time-series scans of one metric made identical when /// they project different labels, so that interning finds them identical: the PromQL frontend projects only the labels each query uses /// (`ts, value, job, …` for `quantile by (job) (0.5, data)`), and @@ -139,9 +165,7 @@ use asap_accuracy::propagation::recompose_guarantees; /// and the other not, when a node would change arity, or when a consumer /// would have to narrow a summary state or a grouped or filtered tumbling-window chain /// would read a `Project`. -pub fn with_union_scan_labels( - roots: &[QueryRoot], -) -> Result>, PerComputationCandidateError> { +pub fn with_union_scan_labels(roots: &[QueryRoot]) -> Result>, CseError> { let mut nodes = Vec::new(); let mut seen = HashSet::new(); for node in roots.iter().flat_map(QueryRoot::operator_refs) { @@ -310,7 +334,7 @@ fn passes_labels_on(node: &OperatorNode) -> bool { fn narrowed( wide: &Rc, schema: &Schema, -) -> Result>, PerComputationCandidateError> { +) -> Result>, CseError> { use asap_types::ir::ProjectItem; if wide.schema.fields == schema.fields { return Ok(Some(Rc::clone(wide))); @@ -370,9 +394,7 @@ fn same_build_family(a: &FieldDataType, b: &FieldDataType) -> bool { /// One build serving every member: the first member with the members' /// shared size ([`size_for_all_readers`]). -fn sized_for_all( - members: &[&Rc], -) -> Result>, PerComputationCandidateError> { +fn sized_for_all(members: &[&Rc]) -> Result>, CseError> { // An exact accumulator has one size: the first member serves them all. if matches!( members[0].asap(), @@ -468,7 +490,7 @@ pub struct WindowCompositionPass2Candidate { /// build has its own chain. pub shared_chains: Vec, /// The candidate's roots with those builds in tumbling windows. - pub roots: Result, PerComputationCandidateError>, + pub roots: Result, CseError>, } /// One tumbling-window chain several tumbling-window candidates read (window-composition rule): @@ -596,9 +618,9 @@ fn splittable_window_ms(build: &OperatorNode) -> Option { pub fn window_composition_pass2_candidates( roots: &[QueryRoot], candidates: &[TumblingWindowCandidate], -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { if candidates.len() > MAX_TUMBLING_WINDOW_CANDIDATES { - return Err(PerComputationCandidateError::TooManyCandidates { + return Err(CseError::TooManyCandidates { counted: "summary builds that could be built in tumbling windows", found: candidates.len(), max: MAX_TUMBLING_WINDOW_CANDIDATES, @@ -1027,7 +1049,7 @@ fn longest(candidates: &[TumblingWindowCandidate], group: &[usize]) -> usize { fn shared_chain_build( candidates: &[TumblingWindowCandidate], group: &[usize], -) -> Result>, PerComputationCandidateError> { +) -> Result>, CseError> { // Of the longest windows, the finest grouping: a rollup group's host. let host = *group .iter() @@ -1053,7 +1075,7 @@ fn in_tumbling_windows( roots: &[QueryRoot], candidates: &[TumblingWindowCandidate], shared_chains: &[(SharedTumblingWindowChain, Rc)], -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { let mut replaced: HashMap<*const OperatorNode, Rc> = HashMap::new(); for (shared, build) in shared_chains { let chain = tumbling_window_chain(&TumblingWindowCandidate { @@ -1074,7 +1096,7 @@ fn in_tumbling_windows( let read = if group_key_count(member) == group_key_count(build) { merge } else { - regrouped(member, &merge).ok_or(PerComputationCandidateError::Unsupported( + regrouped(member, &merge).ok_or(CseError::Unsupported( "a coarser tumbling-window build whose keys the shared chain does not hold", ))? }; @@ -1128,7 +1150,7 @@ pub(crate) fn with_tumbling_window_chain( family: FieldDataType, tumbling_window_width_ms: u64, tumbling_window_count: u32, -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { let Some(ASAPOp::SummaryMerge { children, newest_tumbling_windows: own_newest, @@ -1203,7 +1225,7 @@ pub(crate) fn with_resized_build( roots: &[QueryRoot], build: &Rc, family: FieldDataType, -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { let Some(ASAPOp::SummaryAgg { child, input, @@ -1253,7 +1275,7 @@ pub(crate) fn with_resized_build( /// output names. fn tumbling_window_chain( candidate: &TumblingWindowCandidate, -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { let Some(ASAPOp::SummaryAgg { child: range, family, @@ -1293,7 +1315,7 @@ fn rebuild( node: &Rc, replaced: &HashMap<*const OperatorNode, Rc>, rebuilt: &mut HashMap<*const OperatorNode, Rc>, -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { if let Some(done) = rebuilt.get(&Rc::as_ptr(node)) { return Ok(Rc::clone(done)); } @@ -1350,7 +1372,7 @@ pub struct RollupPass2Candidate { /// The candidate's roots with those builds rolled up and every guarantee /// composed again; `Err` when the rebuilt DAG does not derive or /// compose, which drops only this candidate. - pub roots: Result, PerComputationCandidateError>, + pub roots: Result, CseError>, } /// The most rollup or top-k reuse candidates one candidate may have: every @@ -1836,7 +1858,6 @@ fn rolled_up( finer: &Rc, coarse_value: &OperatorNode, ) -> Option> { - use crate::pass1::per_computation_candidates::{exact_accumulator, partitions_over}; let Some(ASAPOp::SummaryAgg { child, family, @@ -1900,12 +1921,23 @@ fn rolled_up( return None; } let kind = rollup_kind(family)?; - let intent = match kind { - ExactKind::Min => AggIntent::Min { col: Some(value) }, - ExactKind::Max => AggIntent::Max { col: Some(value) }, - _ => AggIntent::Sum { col: Some(value) }, - }; - let state = exact_accumulator(kind, &intent, finer, regrouped).ok()?; + // An exact accumulator over the finer value column: what Pass 1 builds + // for a Sum, Min or Max of that column (no item, signed weights). + let weight = column_ref(&finer.schema.fields[value]); + resolve_column_ref(&weight, &finer.schema).ok()?; + let state = OperatorNode::new_shared(Operator::ASAP(ASAPOp::SummaryAgg { + child: Rc::clone(finer), + input: SummaryUpdate { + item: None, + weight: SummaryInputExpr::Column(weight), + weight_domain: WeightDomain::UnknownOrSigned, + }, + family: exact_family(kind), + reduction: regrouped, + grouping: GroupingStrategy::default(), + filter: None, + })) + .ok()?; let mut finalized = OperatorNode::new_shared(Operator::ASAP(ASAPOp::FinalizeExactAccumulator { child: state, @@ -1981,7 +2013,7 @@ fn is_sketch_build(node: &OperatorNode) -> bool { pub fn rollup_pass2_candidates( roots: &[QueryRoot], candidates: &[RollupCandidate], -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { Ok(subsets(&rollup_choices(candidates), "rollup")? .into_iter() .map(|choices| { @@ -2021,11 +2053,9 @@ pub fn rollup_pass2_candidates( /// coarse builds have no shared size, or a node does not derive. fn sketch_rollups( choices: &[Vec], -) -> Result)>, PerComputationCandidateError> { +) -> Result)>, CseError> { let unsupported = || { - PerComputationCandidateError::Unsupported( - "a sketch rollup whose finest and coarse builds have no shared size", - ) + CseError::Unsupported("a sketch rollup whose finest and coarse builds have no shared size") }; // A finest build can also be the coarse build of another choice, read // through a compensating filter. Each build is replaced once: as a @@ -2075,13 +2105,10 @@ fn sketch_rollups( /// Every non-empty subset of `candidates`, in binary-counter order (the /// first candidate the lowest bit); an error past [`MAX_REUSE_CANDIDATES`]. -fn subsets( - candidates: &[T], - rule: &'static str, -) -> Result>, PerComputationCandidateError> { +fn subsets(candidates: &[T], rule: &'static str) -> Result>, CseError> { let n = candidates.len(); if n > MAX_REUSE_CANDIDATES { - return Err(PerComputationCandidateError::TooManyCandidates { + return Err(CseError::TooManyCandidates { counted: match rule { "rollup" => "rollup choices", _ => "top-k reuse candidates", @@ -2106,12 +2133,12 @@ fn subsets( fn replace_all( roots: &[QueryRoot], replaced: &HashMap<*const OperatorNode, Rc>, -) -> Result, PerComputationCandidateError> { +) -> Result, CseError> { fn walk( node: &Rc, replaced: &HashMap<*const OperatorNode, Rc>, rebuilt: &mut HashMap<*const OperatorNode, Rc>, - ) -> Result, PerComputationCandidateError> { + ) -> Result, CseError> { if let Some(done) = rebuilt.get(&Rc::as_ptr(node)) { return Ok(Rc::clone(done)); } diff --git a/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs b/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs index 21d4d1e36..d849a5e75 100644 --- a/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs +++ b/crates/logical-asap-optimization/src/pass2/memo/enumeration.rs @@ -19,10 +19,9 @@ use std::collections::{BTreeMap, HashMap, HashSet}; use std::rc::Rc; -use asap_types::ir::{OperatorNode, QueryRoot}; +use asap_types::ir::{OperatorNode, QueryRoot, SchemaDerivationError}; use super::rules::reuse::replace; -use crate::pass1::per_computation_candidates::PerComputationCandidateError; use asap_accuracy::propagation::recompose_guarantees; use asap_types::stages::Memo; use asap_types::stages::MemoGroupId; @@ -61,7 +60,7 @@ pub enum ErrorBuildingLogicalASAPDagsFromChosenCandidates { Cycle { group: MemoGroupId }, /// Re-deriving the composed nodes' schemas and guarantees failed. #[error("recomposing the composed DAG: {0}")] - Recompose(#[from] PerComputationCandidateError), + Recompose(#[from] SchemaDerivationError), } /// Every choice of one candidate per needed group, starting from the query @@ -198,7 +197,7 @@ pub fn build_logical_asap_dags_from_chosen_candidates( .collect::, ErrorBuildingLogicalASAPDagsFromChosenCandidates>>()? .into_iter() .unzip(); - let roots = recompose_guarantees(&roots).map_err(PerComputationCandidateError::from)?; + let roots = recompose_guarantees(&roots)?; Ok(groups .into_iter() .zip(roots) diff --git a/crates/logical-asap-optimization/tests/pass_layering.rs b/crates/logical-asap-optimization/tests/pass_layering.rs index 76f9fd92e..2d365db77 100644 --- a/crates/logical-asap-optimization/tests/pass_layering.rs +++ b/crates/logical-asap-optimization/tests/pass_layering.rs @@ -1,21 +1,26 @@ -//! Pass 2 builds on Pass 1, never the reverse (#1096): Pass 1's production -//! code does not import `crate::pass2`. Test modules may use Pass 2 to build -//! fixtures. Pass 2 → Pass 1 is still allowed (#1191). +//! Stage 1's passes are independent modules (#1096, #1191): neither pass's +//! production code imports the other. Test modules may use the other pass to +//! build fixtures. What both passes need lives in `crate::ir_construction`. -use std::path::Path; +use std::path::{Path, PathBuf}; -/// Pass 1's production source does not reference `pass2`. -#[test] -fn pass1_does_not_import_pass2() { - let src = Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); - let mut files = vec![src.join("pass1.rs")]; - for entry in std::fs::read_dir(src.join("pass1")).unwrap() { - let path = entry.unwrap().path(); - // `*_tests.rs` files are `#[cfg(test)]` modules. - if !path.to_string_lossy().ends_with("_tests.rs") { - files.push(path); +/// Production lines under `src/.rs` and `src//` that name a path +/// through `needle`'s module. +fn production_lines_mentioning(pass: &str, needle: &str) -> Vec { + fn walk(dir: &Path, files: &mut Vec) { + for entry in std::fs::read_dir(dir).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + walk(&path, files); + // `*_tests.rs` files are `#[cfg(test)]` modules. + } else if !path.to_string_lossy().ends_with("_tests.rs") { + files.push(path); + } } } + let src = Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); + let mut files = vec![src.join(format!("{pass}.rs"))]; + walk(&src.join(pass), &mut files); let mut offending = Vec::new(); for file in files { let text = std::fs::read_to_string(&file).unwrap(); @@ -23,14 +28,32 @@ fn pass1_does_not_import_pass2() { // parse with `syn` if inline `#[cfg(test)]` items ever need excluding. let production = text.split("#[cfg(test)]\nmod tests").next().unwrap(); for (n, line) in production.lines().enumerate() { - if line.contains("pass2") { + if line.contains(needle) { offending.push(format!("{}:{}: {}", file.display(), n + 1, line.trim())); } } } + offending +} + +/// Pass 1's production source does not reference `pass2`. +#[test] +fn pass1_does_not_import_pass2() { + let offending = production_lines_mentioning("pass1", "pass2::"); assert!( offending.is_empty(), "Pass 1 imports Pass 2:\n{}", offending.join("\n") ); } + +/// Pass 2's production source does not reference `pass1`. +#[test] +fn pass2_does_not_import_pass1() { + let offending = production_lines_mentioning("pass2", "pass1::"); + assert!( + offending.is_empty(), + "Pass 2 imports Pass 1:\n{}", + offending.join("\n") + ); +} From 63c24929a651707d819b6362e692079d88e9faad Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 4 Nov 2026 10:57:48 -0500 Subject: [PATCH 2/2] refactor(types): move Stage 1's shared IR helpers into asap-types (#1191) Review on #1212: the helpers both passes share are IR operations, and column_ref was duplicated. - Field::column_ref replaces ir_construction::column_ref and the copies in schema/topk.rs and pass1's schema_column_ref. - FieldDataType::exact replaces exact_family. - Reduction::partitions_over replaces partitions_over, with a unit test. - ir_construction is deleted. Co-Authored-By: Claude Opus 5.5 --- .../src/ir_construction.rs | 64 ----------------- crates/logical-asap-optimization/src/lib.rs | 2 - .../src/pass1/per_computation_candidates.rs | 21 +++--- .../src/pass1/realization.rs | 27 ++----- .../src/pass2/asap_aware_cse.rs | 7 +- .../tests/pass_layering.rs | 2 +- .../src/ir/operator/operator_properties.rs | 70 +++++++++++++++++++ crates/types/src/ir/schema/mod.rs | 27 +++++++ crates/types/src/ir/schema/topk.rs | 10 +-- 9 files changed, 118 insertions(+), 112 deletions(-) delete mode 100644 crates/logical-asap-optimization/src/ir_construction.rs diff --git a/crates/logical-asap-optimization/src/ir_construction.rs b/crates/logical-asap-optimization/src/ir_construction.rs deleted file mode 100644 index c5129537c..000000000 --- a/crates/logical-asap-optimization/src/ir_construction.rs +++ /dev/null @@ -1,64 +0,0 @@ -//! IR-building helpers both Stage 1 passes use. Neither pass imports the -//! other (`tests/pass_layering.rs`), so what they share lives here. - -use asap_types::ir::operator::Reduction; -use asap_types::ir::scalar::expr_ir::ColumnRef; -use asap_types::ir::schema::{ExactKind, ExactParams, Field, FieldDataType, Schema}; - -/// A reference to `column`: qualified by its table when it has one. -pub(crate) fn column_ref(column: &Field) -> ColumnRef { - match &column.table { - Some(t) => ColumnRef::Qualified { - table: t.clone(), - name: column.name.clone(), - }, - None => ColumnRef::Named(column.name.clone()), - } -} - -/// `FieldDataType::ExactAggregate` of accumulator `kind`, with its parameters -/// (none has any, so they mirror `kind`). -pub(crate) fn exact_family(kind: ExactKind) -> FieldDataType { - let params = match kind { - ExactKind::Sum => ExactParams::Sum, - ExactKind::Count => ExactParams::Count, - ExactKind::Min => ExactParams::Min, - ExactKind::Max => ExactParams::Max, - ExactKind::Rate => ExactParams::Rate, - ExactKind::IRate => ExactParams::IRate, - ExactKind::Increase => ExactParams::Increase, - }; - FieldDataType::ExactAggregate(kind, params) -} - -/// `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. -pub(crate) 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 column = |schema: &Schema, index: usize| column_ref(&schema.fields[index]); - let keys = keys - .keys() - .iter() - .map(|&index| { - inner.fields.get(index)?; - let key = column(inner, index); - let mut matches = (0..input.fields.len()).filter(|&i| column(input, i) == key); - match (matches.next(), matches.next()) { - (Some(i), None) => Some(i), - _ => None, - } - }) - .collect::>()?; - Some(Reduction::by(keys)) -} diff --git a/crates/logical-asap-optimization/src/lib.rs b/crates/logical-asap-optimization/src/lib.rs index b4c709862..213790d6d 100644 --- a/crates/logical-asap-optimization/src/lib.rs +++ b/crates/logical-asap-optimization/src/lib.rs @@ -62,8 +62,6 @@ /// The deployment's capabilities: supported summaries and what it retains. mod deployment_capabilities; -/// IR-building helpers both Stage 1 passes use. -mod ir_construction; mod pass1; mod pass2; /// Stage 0: the memo groups of the workload's query roots. diff --git a/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs b/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs index 9e6867d16..ea2b891ca 100644 --- a/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs +++ b/crates/logical-asap-optimization/src/pass1/per_computation_candidates.rs @@ -136,7 +136,6 @@ use asap_types::workload::bounds::{declared_non_negative, series_read, values_pe use asap_types::workload::DataWorkload; use crate::deployment_capabilities::{summary_supported, SupportedSummary}; -use crate::ir_construction::{exact_family, partitions_over}; use crate::pass1::realization::summary_candidates; use asap_types::stages::{ ComputationTarget, DagOfLogicalMemoGroups, FusedComputation, HavingEntry, HavingPlacement, @@ -1051,7 +1050,7 @@ fn add_nested_sum_targets(targets: &mut [ComputationTarget], logical_nodes: &[Rc } /// The input rows, the inner sum and the outer grouping (moved onto the -/// input rows by column, [`partitions_over`]) of a nested-sum merge of +/// input rows by column, [`Reduction::partitions_over`]) of a nested-sum merge of /// `target`, or `None` when it does not apply: `target` is a grouped `sum` /// (not by `without`) of the value column of a per-series or grouped `sum` /// directly beneath it, neither filtered nor with `HAVING`, and not grouped @@ -1116,7 +1115,7 @@ fn nested_sum_input(target: &OperatorNode) -> Option<(Rc, AggInten { return None; } - let partitions = partitions_over(reduction, &inner.schema, &rows.schema)?; + let partitions = reduction.partitions_over(&inner.schema, &rows.schema)?; Some((Rc::clone(rows), inner_sum.clone(), partitions)) } @@ -1253,7 +1252,7 @@ fn reads_only_time_series(node: &OperatorNode) -> bool { /// `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`]). +/// onto the input the heap reads ([`Reduction::partitions_over`]). fn whole_expression_input( target: &OperatorNode, ) -> Option<(Rc, SummaryUpdate, Reduction)> { @@ -1299,7 +1298,7 @@ fn whole_expression_input( .. }) ); - let partitions = partitions_over(reduction, &child.schema, &realized.child.schema)?; + let partitions = reduction.partitions_over(&child.schema, &realized.child.schema)?; if per_series && realized.child.schema.has_promql_series_identity() { input.item = Some(asap_types::ir::schema::topk::series_identity_item( &realized.child.schema, @@ -1859,7 +1858,7 @@ pub(crate) fn exact_accumulator( input: &Rc, reduction: Reduction, ) -> Result, PerComputationCandidateError> { - let family = exact_family(kind); + let family = FieldDataType::exact(kind); let read = measure_input(intent, input, None)?; Ok(OperatorNode::new_shared(Operator::ASAP( ASAPOp::SummaryAgg { @@ -1907,7 +1906,7 @@ fn substitute_measure( ))?) }; let finalized = |kind: &ExactKind, intent: &AggIntent, read| { - let state = state(exact_family(*kind), intent, read)?; + let state = state(FieldDataType::exact(*kind), intent, read)?; Ok::<_, PerComputationCandidateError>(OperatorNode::new_shared(Operator::ASAP( ASAPOp::FinalizeExactAccumulator { child: state }, ))?) @@ -2287,7 +2286,7 @@ fn counted_column( schema .fields .get(id) - .map(crate::ir_construction::column_ref) + .map(asap_types::ir::schema::Field::column_ref) }; filter .map_or_else(Vec::new, |filter| not_null(&filter.0)) @@ -2349,9 +2348,7 @@ fn summary_update( .iter() .enumerate() .filter(|(index, _)| *index != value && !groups.contains(index)) - .map(|(_, field)| { - SummaryInputExpr::Column(crate::ir_construction::column_ref(field)) - }) + .map(|(_, field)| SummaryInputExpr::Column(field.column_ref())) .collect(); if items.is_empty() { return Err(PerComputationCandidateError::Unsupported( @@ -2363,7 +2360,7 @@ fn summary_update( let excluding = groups .iter() .filter_map(|&index| child.fields.get(index)) - .map(crate::ir_construction::column_ref) + .map(asap_types::ir::schema::Field::column_ref) .collect(); SummaryInputExpr::EntityIdentity(EntityIdentity::PromqlLabelSet { excluding }) }; diff --git a/crates/logical-asap-optimization/src/pass1/realization.rs b/crates/logical-asap-optimization/src/pass1/realization.rs index 7691730dd..d5dc1f09c 100644 --- a/crates/logical-asap-optimization/src/pass1/realization.rs +++ b/crates/logical-asap-optimization/src/pass1/realization.rs @@ -10,7 +10,7 @@ use asap_types::ir::scalar::expr_ir::ColumnRef; use asap_types::ir::schema::topk::ranked_column; use asap_types::ir::schema::ColumnId; use asap_types::ir::schema::{ - EntityIdentity, FamilyCapability, FieldDataType, NonNegativeWeightProof, Schema, + EntityIdentity, FamilyCapability, Field, FieldDataType, NonNegativeWeightProof, Schema, SketchAlgorithm, SummaryInputExpr, SummaryUpdate, WeightDomain, }; use asap_types::ir::{NonASAPOp, OperatorNode}; @@ -18,8 +18,6 @@ use asap_types::ir::{NonASAPOp, OperatorNode}; use asap_types::ir::{Operator, ScalarExpr}; use std::rc::Rc; -use crate::ir_construction::column_ref; - use asap_accuracy::sizing::sketch_statistic; /// The sketch families the planner sizes and offers, in the order Pass 1 @@ -78,7 +76,7 @@ pub(crate) fn summarised_column(intent: &AggIntent, child_schema: &Schema) -> Co .first() .and_then(|id| child_schema.fields.get(*id)) { - Some(c) => column_ref(c), + Some(c) => c.column_ref(), None => ColumnRef::SampleValue, } } @@ -89,7 +87,7 @@ pub(crate) fn ranked_value(child: &OperatorNode) -> Result<(ColumnId, ColumnRef) let ranked = ranked_column(child)?; let column = match resolve_column_ref(&ColumnRef::SampleValue, &child.schema) { Ok(value) if value == ranked => ColumnRef::SampleValue, - _ => column_ref(&child.schema.fields[ranked]), + _ => child.schema.fields[ranked].column_ref(), }; Ok((ranked, column)) } @@ -115,7 +113,7 @@ pub(crate) fn summarised_input( } let legs = cols .iter() - .map(|id| child_schema.fields.get(*id).map(column_ref)) + .map(|id| child_schema.fields.get(*id).map(Field::column_ref)) .collect::>>()?; Some(SummaryInputExpr::Tuple( legs.into_iter().map(SummaryInputExpr::Column).collect(), @@ -209,7 +207,7 @@ pub(crate) fn realize_keyed_additive_summary_input( } [AggIntent::Sum { col }] => SummaryInputExpr::Column(match col { None => ColumnRef::SampleValue, - Some(index) => match schema_column_ref(raw_child, *index) { + Some(index) => match raw_child.schema.fields.get(*index).map(Field::column_ref) { Some(column) => column, None => { // Unsupported: sum-ranked Top-K value column is outside the raw input schema. @@ -239,7 +237,7 @@ pub(crate) fn realize_keyed_additive_summary_input( Reduction::PerEntity => vec![], Reduction::Reduce(keys) => keys .iter() - .filter_map(|index| schema_column_ref(child, *index)) + .filter_map(|index| child.schema.fields.get(*index).map(Field::column_ref)) .collect(), }; let item = match reduction { @@ -249,7 +247,7 @@ pub(crate) fn realize_keyed_additive_summary_input( Reduction::Reduce(keys) if !keys.is_without() && !keys.is_empty() => { let Some(columns) = keys .iter() - .map(|index| schema_column_ref(raw_child, *index)) + .map(|index| raw_child.schema.fields.get(*index).map(Field::column_ref)) .collect::>>() else { // Unsupported: ranked item column is outside the raw input schema. @@ -349,17 +347,6 @@ fn summed_per_series_samples(grouped: &OperatorNode) -> Option<(Rc Some((Rc::clone(rows), inner.clone())) } -pub(crate) fn schema_column_ref(child: &OperatorNode, index: usize) -> Option { - let column = child.schema.fields.get(index)?; - Some(match &column.table { - Some(table) => ColumnRef::Qualified { - table: table.clone(), - name: column.name.clone(), - }, - None => ColumnRef::Named(column.name.clone()), - }) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs b/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs index dc3ed2a56..991e96f40 100644 --- a/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs +++ b/crates/logical-asap-optimization/src/pass2/asap_aware_cse.rs @@ -122,7 +122,6 @@ use asap_types::ir::{ TimeRangeKind, }; -use crate::ir_construction::{column_ref, exact_family, partitions_over}; use asap_accuracy::estimators::size_for_all_readers; use asap_accuracy::propagation::recompose_guarantees; @@ -1913,7 +1912,7 @@ fn rolled_up( { return None; } - let regrouped = partitions_over(reduction, &child.schema, &finer.schema)?; + let regrouped = reduction.partitions_over(&child.schema, &finer.schema)?; if regrouped .group_keys() .is_some_and(|keys| keys.iter().any(|&k| k >= value)) @@ -1923,7 +1922,7 @@ fn rolled_up( let kind = rollup_kind(family)?; // An exact accumulator over the finer value column: what Pass 1 builds // for a Sum, Min or Max of that column (no item, signed weights). - let weight = column_ref(&finer.schema.fields[value]); + let weight = finer.schema.fields[value].column_ref(); resolve_column_ref(&weight, &finer.schema).ok()?; let state = OperatorNode::new_shared(Operator::ASAP(ASAPOp::SummaryAgg { child: Rc::clone(finer), @@ -1932,7 +1931,7 @@ fn rolled_up( weight: SummaryInputExpr::Column(weight), weight_domain: WeightDomain::UnknownOrSigned, }, - family: exact_family(kind), + family: FieldDataType::exact(kind), reduction: regrouped, grouping: GroupingStrategy::default(), filter: None, diff --git a/crates/logical-asap-optimization/tests/pass_layering.rs b/crates/logical-asap-optimization/tests/pass_layering.rs index 2d365db77..daf0595a7 100644 --- a/crates/logical-asap-optimization/tests/pass_layering.rs +++ b/crates/logical-asap-optimization/tests/pass_layering.rs @@ -1,6 +1,6 @@ //! Stage 1's passes are independent modules (#1096, #1191): neither pass's //! production code imports the other. Test modules may use the other pass to -//! build fixtures. What both passes need lives in `crate::ir_construction`. +//! build fixtures. IR helpers both passes need live in `asap-types`. use std::path::{Path, PathBuf}; diff --git a/crates/types/src/ir/operator/operator_properties.rs b/crates/types/src/ir/operator/operator_properties.rs index 0edf68a59..e2a6ef1da 100644 --- a/crates/types/src/ir/operator/operator_properties.rs +++ b/crates/types/src/ir/operator/operator_properties.rs @@ -447,6 +447,38 @@ pub enum Reduction { PerEntity, } +impl Reduction { + /// These keys, which index `inner`, as indices into `input`, 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. + pub fn partitions_over( + &self, + inner: &crate::ir::schema::Schema, + input: &crate::ir::schema::Schema, + ) -> Option { + let Reduction::Reduce(keys) = self else { + return Some(Reduction::PerEntity); + }; + if keys.is_without() { + return None; + } + let keys = keys + .keys() + .iter() + .map(|&index| { + let key = inner.fields.get(index)?.column_ref(); + let mut matches = + (0..input.fields.len()).filter(|&i| input.fields[i].column_ref() == key); + match (matches.next(), matches.next()) { + (Some(i), None) => Some(i), + _ => None, + } + }) + .collect::>()?; + Some(Reduction::by(keys)) + } +} + impl Reduction { /// Shorthand for the common case — group by these (possibly empty) /// keys, kept rather than excluded. @@ -475,3 +507,41 @@ impl Reduction { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::ir::schema::{DataType, Field, Schema}; + + /// A key moves to the input's column of the same name and qualifier; a + /// key that is missing or ambiguous there, or a `without` grouping, has no + /// partitions over the input. + #[test] + fn partitions_over_matches_keys_by_column() { + let field = |name: &str| Field::plain(name, DataType::Utf8, false); + let inner = Schema::new(vec![field("job"), field("value")]); + let input = Schema::new(vec![field("ts"), field("job"), field("value")]); + assert_eq!( + Reduction::by(vec![0]).partitions_over(&inner, &input), + Some(Reduction::by(vec![1])) + ); + assert_eq!( + Reduction::PerEntity.partitions_over(&inner, &input), + Some(Reduction::PerEntity) + ); + let ambiguous = Schema::new(vec![field("job"), field("job")]); + assert_eq!( + Reduction::by(vec![0]).partitions_over(&inner, &ambiguous), + None + ); + let qualified = Schema::new(vec![field("job").with_table("t")]); + assert_eq!( + Reduction::by(vec![0]).partitions_over(&inner, &qualified), + None + ); + assert_eq!( + Reduction::Reduce(GroupKeys::without(vec![0])).partitions_over(&inner, &input), + None + ); + } +} diff --git a/crates/types/src/ir/schema/mod.rs b/crates/types/src/ir/schema/mod.rs index 6b4196ceb..90796a76b 100644 --- a/crates/types/src/ir/schema/mod.rs +++ b/crates/types/src/ir/schema/mod.rs @@ -83,6 +83,18 @@ impl Field { self.table = Some(table.into()); self } + + /// A reference to this field: qualified by its table when it has one. + pub fn column_ref(&self) -> crate::ir::scalar::ColumnRef { + use crate::ir::scalar::ColumnRef; + match &self.table { + Some(table) => ColumnRef::Qualified { + table: table.clone(), + name: self.name.clone(), + }, + None => ColumnRef::Named(self.name.clone()), + } + } } impl Field { @@ -168,6 +180,21 @@ pub enum FieldDataType { } impl FieldDataType { + /// The state of exact accumulator `kind`, with its parameters (none has + /// any, so they mirror `kind`). + pub fn exact(kind: ExactKind) -> Self { + let params = match kind { + ExactKind::Sum => ExactParams::Sum, + ExactKind::Count => ExactParams::Count, + ExactKind::Min => ExactParams::Min, + ExactKind::Max => ExactParams::Max, + ExactKind::Rate => ExactParams::Rate, + ExactKind::IRate => ExactParams::IRate, + ExactKind::Increase => ExactParams::Increase, + }; + FieldDataType::ExactAggregate(kind, params) + } + pub(crate) fn is_plain(&self) -> bool { matches!(self, FieldDataType::Plain(_)) } diff --git a/crates/types/src/ir/schema/topk.rs b/crates/types/src/ir/schema/topk.rs index 11ab729f3..1ad5c2186 100644 --- a/crates/types/src/ir/schema/topk.rs +++ b/crates/types/src/ir/schema/topk.rs @@ -130,15 +130,7 @@ pub fn series_identity_item( && field.dtype == FieldDataType::Plain(DataType::Utf8) && field.name != PROMQL_SERIES_IDENTITY }) - .map(|(_, field)| { - SummaryInputExpr::Column(match &field.table { - Some(table) => ColumnRef::Qualified { - table: table.clone(), - name: field.name.clone(), - }, - None => ColumnRef::Named(field.name.clone()), - }) - }) + .map(|(_, field)| SummaryInputExpr::Column(field.column_ref())) .collect(); items.push(SummaryInputExpr::Column(ColumnRef::Named( PROMQL_SERIES_IDENTITY.into(),