Skip to content

refactor(ir): single canonical Aggregate output-schema derivation (#41) #92

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

Merged
zzylol merged 1 commit into from
Jul 4, 2026
Merged

refactor(ir): single canonical Aggregate output-schema derivation (#41) #92

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 3 additions & 3 deletions crates/ir/src/intent_algebra/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ pub use agg_intent::{
pub use expr_ir::{ArithOp, ColumnRef, CompareOp, Expr, L2Expr, L3Expr, L3Scalar};
pub use names::{BindingName, QueryId};
pub use query_expr::{
BinaryOpKind, BindingScope, DataModel, GroupKeys, GroupSide, JoinKind, Predicate, ProjectItem,
QueryExpr, QueryExprError, SetOpKind, SortKey, Source, VectorGrouping, VectorMatch,
VectorMatchKind, WindowFuncKind, WindowKind,
aggregate_output_schema, BinaryOpKind, BindingScope, DataModel, GroupKeys, GroupSide, JoinKind,
Predicate, ProjectItem, QueryExpr, QueryExprError, SetOpKind, SortKey, Source, VectorGrouping,
VectorMatch, VectorMatchKind, WindowFuncKind, WindowKind,
};
pub use schema::{cse_reuse_is_legal, Column, ColumnId, CseError, DataType, Schema};
167 changes: 92 additions & 75 deletions crates/ir/src/intent_algebra/query_expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,81 +481,9 @@ impl QueryExpr {
child.as_ref(),
QueryExpr::TimeRange { .. } | QueryExpr::Subquery { .. }
);
if by.is_empty() && aggs.len() == 1 && (aggs[0].is_per_series() || is_range_child) {
return Ok(per_series_reduction_schema(&in_schema, &aggs[0]));
}

let mut out_cols: Vec<Column> = Vec::with_capacity(by.len() + aggs.len());
for &id in by {
let c =
in_schema
.columns
.get(id)
.ok_or(QueryExprError::InvalidGroupByColumn(
id,
in_schema.columns.len(),
))?;
out_cols.push(c.clone());
}
let value_col_idx = in_schema
.column_id("value")
.or_else(|| (0..in_schema.columns.len()).find(|i| !by.contains(i)));
let probe = value_col_idx
.and_then(|i| in_schema.columns.get(i))
.cloned()
.unwrap_or_else(|| Column::new("value", DataType::Float64, false));
// Each reducer types off its own input column (`SUM(bytes)` vs
// `AVG(latency)` in one node); `None` falls back to the sample-
// value probe (PromQL's single-column convention). A non-empty
// `output_names[i]` overrides the synthetic output column name.
for (i, intent) in aggs.iter().enumerate() {
// `count_values("l", v)` emits TWO columns: the synthesized
// `Utf8` label `l` (the stringified sample value it groups
// by) and the per-value count. `output_names[i]` still
// overrides the count column's name if set. If `l` collides
// with a group-by key of the same name, PromQL's synthesized
// label takes precedence — emit a single column, never a
// duplicate.
if let AggIntent::CountValues { label } = intent {
if !out_cols.iter().any(|c| c.name == *label) {
out_cols.push(Column::new(label.clone(), DataType::Utf8, false));
}
let mut cnt = intent.output_column(&probe);
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
cnt.name = name.clone();
}
out_cols.push(cnt);
continue;
}
let in_col = intent
.input_col()
.and_then(|id| in_schema.columns.get(id))
.unwrap_or(&probe);
let mut out = intent.output_column(in_col);
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
out.name = name.clone();
}
out_cols.push(out);
}
// `count_values` groups by (by-keys ∪ the synthesized value
// label), so the by-keys alone are not a unique key — be
// conservative and claim none.
let has_count_values =
aggs.iter().any(|a| matches!(a, AggIntent::CountValues { .. }));
let unique_keys = if by.is_empty() || has_count_values {
Vec::new()
} else {
vec![(0..by.len()).collect()]
};
Ok(Schema {
columns: out_cols,
time_index: None,
unique_keys,
// A cross-series aggregate enumerates exactly `by ++ aggs`,
// so its output is closed even over an open input — this is
// where an open schema freezes to closed.
closed: true,
})
let per_series =
by.is_empty() && aggs.len() == 1 && (aggs[0].is_per_series() || is_range_child);
aggregate_output_schema(&in_schema, by, aggs, output_names, per_series)
}

QueryExpr::LetBinding { name, expr, child } => {
Expand Down Expand Up @@ -797,6 +725,95 @@ fn per_series_reduction_schema(input: &Schema, agg: &AggIntent) -> Schema {
}
}

/// The output schema of an `Aggregate { by, aggs }` over `in_schema` — the
/// **single** canonical derivation shared by [`QueryExpr::output_schema_in`]'s
/// `Aggregate` arm and the converter's HAVING-resolution path
/// (`column_resolution::output_schema_for_aggregate`), so the two can never
/// drift (issue #41).
///
/// `per_series` selects the label-preserving [`per_series_reduction_schema`]
/// (`rate`/`increase`/`*_over_time`) instead of the cross-series `by ++ aggs`
/// shape. The caller supplies it because the decision depends on the *child
/// node* (a `TimeRange`/`Subquery` marker), which this function does not see;
/// the child-independent part is `by.is_empty() && aggs.len() == 1 &&
/// aggs[0].is_per_series()`.
pub fn aggregate_output_schema(
in_schema: &Schema,
by: &[ColumnId],
aggs: &[AggIntent],
output_names: &[String],
per_series: bool,
) -> Result<Schema, QueryExprError> {
if per_series {
debug_assert_eq!(aggs.len(), 1, "a per-series reduction is single-aggregate");
return Ok(per_series_reduction_schema(in_schema, &aggs[0]));
}

let mut out_cols: Vec<Column> = Vec::with_capacity(by.len() + aggs.len());
for &id in by {
let c = in_schema
.columns
.get(id)
.ok_or(QueryExprError::InvalidGroupByColumn(id, in_schema.columns.len()))?;
out_cols.push(c.clone());
}
let value_col_idx = in_schema
.column_id("value")
.or_else(|| (0..in_schema.columns.len()).find(|i| !by.contains(i)));
let probe = value_col_idx
.and_then(|i| in_schema.columns.get(i))
.cloned()
.unwrap_or_else(|| Column::new("value", DataType::Float64, false));
// Each reducer types off its own input column (`SUM(bytes)` vs `AVG(latency)`
// in one node); `None` falls back to the sample-value probe (PromQL's
// single-column convention). A non-empty `output_names[i]` overrides the
// synthetic output column name.
for (i, intent) in aggs.iter().enumerate() {
// `count_values("l", v)` emits TWO columns: the synthesized `Utf8` label
// `l` (the stringified sample value it groups by) and the per-value
// count. If `l` collides with a group-by key of the same name, PromQL's
// synthesized label takes precedence — emit a single column, never a
// duplicate.
if let AggIntent::CountValues { label } = intent {
if !out_cols.iter().any(|c| c.name == *label) {
out_cols.push(Column::new(label.clone(), DataType::Utf8, false));
}
let mut cnt = intent.output_column(&probe);
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
cnt.name = name.clone();
}
out_cols.push(cnt);
continue;
}
let in_col = intent
.input_col()
.and_then(|id| in_schema.columns.get(id))
.unwrap_or(&probe);
let mut out = intent.output_column(in_col);
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
out.name = name.clone();
}
out_cols.push(out);
}
// `count_values` groups by (by-keys ∪ the synthesized value label), so the
// by-keys alone are not a unique key — be conservative and claim none.
let has_count_values = aggs.iter().any(|a| matches!(a, AggIntent::CountValues { .. }));
let unique_keys = if by.is_empty() || has_count_values {
Vec::new()
} else {
vec![(0..by.len()).collect()]
};
Ok(Schema {
columns: out_cols,
time_index: None,
unique_keys,
// A cross-series aggregate enumerates exactly `by ++ aggs`, so its output
// is closed even over an open input — this is where an open schema
// freezes to closed.
closed: true,
})
}

/// Infer the `(DataType, nullable)` a scalar [`L3Expr`] produces against an
/// input [`Schema`]. Used by `Project` schema derivation. Approximate at L3:
/// unknown columns and bare `FunctionCall`s fall back to a permissive default
Expand Down
102 changes: 63 additions & 39 deletions crates/l2/src/column_resolution.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use thiserror::Error;
use asap_ir::intent_algebra::agg_intent::AggIntent;
use asap_ir::intent_algebra::expr_ir::ColumnRef;
use asap_ir::intent_algebra::expr_ir::{L2Expr, L3Expr};
use asap_ir::intent_algebra::query_expr::{aggregate_output_schema, QueryExprError};
use crate::relational::QueryExpr;
use asap_ir::intent_algebra::schema::{Column, ColumnId, DataType, Schema};

Expand Down Expand Up @@ -169,44 +170,15 @@ pub fn output_schema_for_aggregate(
by: &[ColumnId],
aggs: &[AggIntent],
output_names: &[String],
) -> Schema {
let mut out_cols: Vec<Column> = Vec::with_capacity(by.len() + aggs.len());
for &id in by {
if let Some(c) = input.columns.get(id) {
out_cols.push(c.clone());
}
}
let value_col_idx = input
.column_id("value")
.or_else(|| (0..input.columns.len()).find(|i| !by.contains(i)));
let probe = value_col_idx
.and_then(|i| input.columns.get(i))
.cloned()
.unwrap_or_else(|| Column::new("value", DataType::Float64, false));
for (i, intent) in aggs.iter().enumerate() {
let in_col = intent
.input_col()
.and_then(|id| input.columns.get(id))
.unwrap_or(&probe);
let mut out = intent.output_column(in_col);
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
out.name = name.clone();
}
out_cols.push(out);
}
let unique_keys = if by.is_empty() {
Vec::new()
} else {
vec![(0..by.len()).collect()]
};
Schema {
columns: out_cols,
time_index: None,
unique_keys,
// A cross-series aggregate fully determines its output columns, so the
// result is closed even over an open input (mirrors `output_schema_in`).
closed: true,
}
) -> Result<Schema, QueryExprError> {
// Delegate to the single canonical derivation so HAVING resolution can never
// drift from `QueryExpr::output_schema_in` (issue #41). HAVING is SQL-only
// and cross-series, but detect the child-independent per-series case anyway
// (a lone `rate`/`increase`/`*_over_time` intent) so the two agree on every
// shared input — the `TimeRange`/`Subquery` marker the canonical arm also
// keys off is not visible here, and never co-occurs with HAVING.
let per_series = by.is_empty() && aggs.len() == 1 && aggs[0].is_per_series();
aggregate_output_schema(input, by, aggs, output_names, per_series)
}

#[cfg(test)]
Expand Down Expand Up @@ -285,11 +257,63 @@ mod tests {
.columns
.push(Column::new("host", DataType::Utf8, false));
let out =
output_schema_for_aggregate(&input, &[2usize], &[AggIntent::Sum { col: None }], &[]);
output_schema_for_aggregate(&input, &[2usize], &[AggIntent::Sum { col: None }], &[])
.expect("valid group-by column");
assert_eq!(out.columns.len(), 2); // host, sum
assert_eq!(out.columns[0].name, "host");
assert_eq!(out.columns[1].name, "sum");
assert!(out.time_index.is_none());
assert_eq!(out.unique_keys, vec![vec![0]]);
}

#[test]
fn having_schema_agrees_with_canonical_for_a_per_series_reduction() {
// Issue #41: `output_schema_for_aggregate` (HAVING resolution) and the
// canonical `QueryExpr::output_schema_in` must produce identical schemas
// for the same aggregate. Before the dedup this diverged on a per-series
// reduction — the HAVING mirror lacked the per-series branch and would
// collapse `[ts, value]` to a single `rate` column.
use asap_ir::intent_algebra::query_expr::{QueryExpr as L3, Source};
use std::time::Duration;

let leaf_schema = Schema::with_time_index(
vec![
Column::new("ts", DataType::Timestamp, false),
Column::new("value", DataType::Float64, false),
],
0,
vec![],
);
let scan = L3::Scan {
source: Source::TimeSeries { metric: "m".into() },
predicates: vec![],
schema: leaf_schema.clone(),
};
// Aggregate{ by: [], [Rate], child: TimeRange{ Scan } } — a per-series
// reduction (label-preserving).
let agg = L3::Aggregate {
by: Default::default(),
aggs: vec![AggIntent::Rate],
output_names: vec![],
having: None,
child: Box::new(L3::TimeRange {
range: Duration::from_secs(300),
child: Box::new(scan),
}),
};
let canonical = agg.output_schema().expect("canonical schema");

// The HAVING-resolution derivation gets only the input schema (the
// TimeRange passes the leaf schema through).
let having_side =
output_schema_for_aggregate(&leaf_schema, &[], &[AggIntent::Rate], &[]).unwrap();

assert_eq!(
canonical, having_side,
"the two aggregate-schema derivations must agree (issue #41)"
);
// Sanity: it really is the label-preserving per-series shape, not `[rate]`.
assert!(having_side.columns.iter().any(|c| c.name == "value"));
assert!(having_side.time_index.is_some());
}
}
2 changes: 1 addition & 1 deletion crates/l2/src/lower.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,7 @@ pub fn convert(
.as_ref()
.map(|h| -> Result<Predicate, ConvertError> {
let out_schema =
output_schema_for_aggregate(&child_schema, &by, &intents, &output_names);
output_schema_for_aggregate(&child_schema, &by, &intents, &output_names)?;
Ok(Predicate(resolve_expr(h, &out_schema)?))
})
.transpose()?;
Expand Down
Loading