Skip to content

Commit d96e032

Browse files
zzylolclaude
zzylol
and
claude
committed
feat(core+sql): thread aggregate output names so Project-over-Aggregate resolves
Closes the output_names gap (PR #5 item 5). DataFusion names aggregate outputs in its plan schema (e.g. "sum(metrics.bytes)") and the enclosing Projection references them by those names; the L3 Aggregate named its outputs synthetically ("sum"), so the Project's column refs didn't resolve and fell back to Utf8. - query_expr: Aggregate gains output_names: Vec<String> (parallel to aggs); a non-empty entry overrides the synthetic AggIntent::output_column name. Schema derivation (+ output_schema_for_aggregate helper) honors it. - lower: convert threads AggItem.alias → output_names in all three Aggregate constructions (fused / multi-agg / topk). - sql: lower_aggregate sets each AggItem.alias to DataFusion's aggregate output field name (schema.fields() past the group cols) — the names the Projection uses. - promql: aggregate alias is now empty (was "value") → keeps intent-keyed output names ("sum", "quantile_0_99", …); PromQL output naming unchanged. - cse: carry output_names through the shared-producer rewrite. Test: SELECT SUM(bytes), AVG(latency) → the root Projection's output schema now resolves to Int64 / Float64 (was the Utf8 fallback). 110 tests green, clippy -D warnings clean, fmt clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 6ae087b commit d96e032

7 files changed

Lines changed: 97 additions & 15 deletions

File tree

‎crates/core/src/intent_algebra/column_resolution.rs‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,12 @@ pub fn resolve_named_keys(keys: &[String], schema: &Schema) -> Result<Vec<Column
101101
/// Mirrors `QueryExpr::output_schema_in`'s `Aggregate` arm; out-of-range `by`
102102
/// ids are silently dropped (callers needing the strict check resolve `by`
103103
/// via [`resolve_named_keys`], which surfaces `NotFound`).
104-
pub fn output_schema_for_aggregate(input: &Schema, by: &[ColumnId], aggs: &[AggIntent]) -> Schema {
104+
pub fn output_schema_for_aggregate(
105+
input: &Schema,
106+
by: &[ColumnId],
107+
aggs: &[AggIntent],
108+
output_names: &[String],
109+
) -> Schema {
105110
let mut out_cols: Vec<Column> = Vec::with_capacity(by.len() + aggs.len());
106111
for &id in by {
107112
if let Some(c) = input.columns.get(id) {
@@ -119,12 +124,16 @@ pub fn output_schema_for_aggregate(input: &Schema, by: &[ColumnId], aggs: &[AggI
119124
dtype: DataType::Float64,
120125
nullable: false,
121126
});
122-
for intent in aggs {
127+
for (i, intent) in aggs.iter().enumerate() {
123128
let in_col = intent
124129
.input_col()
125130
.and_then(|id| input.columns.get(id))
126131
.unwrap_or(&probe);
127-
out_cols.push(intent.output_column(in_col));
132+
let mut out = intent.output_column(in_col);
133+
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
134+
out.name = name.clone();
135+
}
136+
out_cols.push(out);
128137
}
129138
let unique_keys = if by.is_empty() {
130139
Vec::new()
@@ -171,7 +180,8 @@ mod tests {
171180
dtype: DataType::Utf8,
172181
nullable: false,
173182
});
174-
let out = output_schema_for_aggregate(&input, &[2usize], &[AggIntent::Sum { col: None }]);
183+
let out =
184+
output_schema_for_aggregate(&input, &[2usize], &[AggIntent::Sum { col: None }], &[]);
175185
assert_eq!(out.columns.len(), 2); // host, sum
176186
assert_eq!(out.columns[0].name, "host");
177187
assert_eq!(out.columns[1].name, "sum");

‎crates/core/src/intent_algebra/cse.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,11 +89,13 @@ pub fn dedupe_subtrees(roots: Vec<(QueryId, QueryExpr)>) -> CseWorkloadPlan {
8989
QueryExpr::Aggregate {
9090
by,
9191
aggs,
92+
output_names,
9293
having,
9394
child,
9495
} if *child == shared_expr => QueryExpr::Aggregate {
9596
by,
9697
aggs,
98+
output_names,
9799
having,
98100
child: Box::new(QueryExpr::Ref {
99101
name: binding_name.clone(),
@@ -162,6 +164,7 @@ mod tests {
162164
q: 0.99,
163165
accuracy: AccuracyTarget::Epsilon(0.01),
164166
}],
167+
output_names: vec![],
165168
having: None,
166169
child: Box::new(windowed_scan()),
167170
};
@@ -178,6 +181,7 @@ mod tests {
178181
q,
179182
accuracy: AccuracyTarget::Epsilon(0.01),
180183
}],
184+
output_names: vec![],
181185
having: None,
182186
child: Box::new(windowed_scan()),
183187
};
@@ -219,6 +223,7 @@ mod tests {
219223
let mk = || QueryExpr::Aggregate {
220224
by: vec![],
221225
aggs: vec![AggIntent::Sum { col: None }],
226+
output_names: vec![],
222227
having: None,
223228
child: Box::new(scan_no_uk.clone()),
224229
};

‎crates/core/src/intent_algebra/lower.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,7 @@ pub fn convert(
110110
let aggregate = CQueryExpr::Aggregate {
111111
by: Vec::new(),
112112
aggs: vec![intent],
113+
output_names: vec![aggs[0].alias.clone()],
113114
having: None,
114115
child: Box::new(agg_child),
115116
};
@@ -150,6 +151,7 @@ pub fn convert(
150151
CQueryExpr::Aggregate {
151152
by,
152153
aggs: intents,
154+
output_names: aggs.iter().map(|item| item.alias.clone()).collect(),
153155
having: having.clone().map(Predicate),
154156
child: Box::new(child),
155157
}
@@ -198,6 +200,7 @@ pub fn convert(
198200
k: *k as usize,
199201
accuracy: acc.clone(),
200202
}],
203+
output_names: vec![],
201204
having: None,
202205
child: Box::new(child),
203206
}
@@ -425,10 +428,11 @@ mod tests {
425428
]
426429
);
427430

428-
// Output schema types each reducer off its own input column.
431+
// Output schema types each reducer off its own input column, and names
432+
// it from the AggItem alias (threaded via Aggregate.output_names).
429433
let out = l3.output_schema().unwrap();
430-
assert_eq!(out.columns[0], col("sum", DataType::Int64)); // SUM(bytes:Int64)
431-
assert_eq!(out.columns[1], col("avg", DataType::Float64)); // AVG(latency)→Float64
434+
assert_eq!(out.columns[0], col("total_bytes", DataType::Int64)); // SUM(bytes:Int64)
435+
assert_eq!(out.columns[1], col("avg_latency", DataType::Float64)); // AVG(latency)→Float64
432436
}
433437

434438
/// PromQL's single sample-value reducer stays `col: None`.

‎crates/core/src/intent_algebra/query_expr.rs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,14 @@ pub enum QueryExpr {
237237
Aggregate {
238238
by: Vec<ColumnId>,
239239
aggs: Vec<AggIntent>,
240+
/// Output column names parallel to `aggs`. A non-empty entry overrides
241+
/// the synthetic intent-keyed name — SQL threads DataFusion's generated
242+
/// name (e.g. `"sum(metrics.bytes)"`) here so an enclosing `Project`
243+
/// resolves the aggregate output by the name it references. An empty
244+
/// entry (or empty vec) falls back to `AggIntent::output_column`'s name
245+
/// (PromQL's convention).
246+
#[serde(default)]
247+
output_names: Vec<String>,
240248
#[serde(default)]
241249
having: Option<Predicate>,
242250
child: Box<QueryExpr>,
@@ -333,7 +341,11 @@ impl QueryExpr {
333341
QueryExpr::Window { child, .. } => child.output_schema_in(scope),
334342

335343
QueryExpr::Aggregate {
336-
by, aggs, child, ..
344+
by,
345+
aggs,
346+
output_names,
347+
child,
348+
..
337349
} => {
338350
let in_schema = child.output_schema_in(scope)?;
339351
let mut out_cols: Vec<Column> = Vec::with_capacity(by.len() + aggs.len());
@@ -361,13 +373,18 @@ impl QueryExpr {
361373
});
362374
// Each reducer types off its own input column (`SUM(bytes)` vs
363375
// `AVG(latency)` in one node); `None` falls back to the sample-
364-
// value probe (PromQL's single-column convention).
365-
for intent in aggs {
376+
// value probe (PromQL's single-column convention). A non-empty
377+
// `output_names[i]` overrides the synthetic output column name.
378+
for (i, intent) in aggs.iter().enumerate() {
366379
let in_col = intent
367380
.input_col()
368381
.and_then(|id| in_schema.columns.get(id))
369382
.unwrap_or(&probe);
370-
out_cols.push(intent.output_column(in_col));
383+
let mut out = intent.output_column(in_col);
384+
if let Some(name) = output_names.get(i).filter(|s| !s.is_empty()) {
385+
out.name = name.clone();
386+
}
387+
out_cols.push(out);
371388
}
372389
let unique_keys = if by.is_empty() {
373390
Vec::new()

‎crates/lower/src/promql.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -398,7 +398,9 @@ fn windowed_aggregate(inner: Inner, keys: Vec<String>, func: AggFunc) -> L2 {
398398
L2::Aggregate {
399399
keys,
400400
aggs: vec![AggItem {
401-
alias: "value".into(),
401+
// Empty alias → the converter keeps PromQL's intent-keyed output
402+
// names ("sum", "quantile_0_99", …) instead of overriding them.
403+
alias: String::new(),
402404
func,
403405
col: ColumnRef::SampleValue,
404406
distinct: false,
@@ -415,7 +417,9 @@ fn outer_aggregate(keys: Vec<String>, func: AggFunc, input: L2) -> L2 {
415417
L2::Aggregate {
416418
keys,
417419
aggs: vec![AggItem {
418-
alias: "value".into(),
420+
// Empty alias → the converter keeps PromQL's intent-keyed output
421+
// names ("sum", "quantile_0_99", …) instead of overriding them.
422+
alias: String::new(),
419423
func,
420424
col: ColumnRef::SampleValue,
421425
distinct: false,

‎crates/lower/src/sql/mod.rs‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -209,11 +209,30 @@ impl<'a> SqlLowerer<'a> {
209209
.iter()
210210
.map(expr_to_group_name)
211211
.collect::<Result<Vec<_>, _>>()?;
212+
// DataFusion names the aggregate outputs in its own schema (e.g.
213+
// "sum(metrics.bytes)") — the same names the enclosing Projection
214+
// references. The schema is [group fields …, aggregate fields …], so
215+
// skip the group fields and thread the rest as L2 aliases → L3
216+
// `Aggregate.output_names`, letting that Projection resolve them.
217+
let out_names: Vec<String> = agg
218+
.schema
219+
.fields()
220+
.iter()
221+
.skip(agg.group_expr.len())
222+
.map(|f| f.name().to_string())
223+
.collect();
212224
let aggs = agg
213225
.aggr_expr
214226
.iter()
215-
.map(lower_agg_item)
216-
.collect::<Result<Vec<_>, _>>()?;
227+
.enumerate()
228+
.map(|(i, e)| {
229+
let mut item = lower_agg_item(e)?;
230+
if let Some(name) = out_names.get(i) {
231+
item.alias = name.clone();
232+
}
233+
Ok(item)
234+
})
235+
.collect::<Result<Vec<_>, LoweringError>>()?;
217236
Ok(L2::Aggregate {
218237
keys,
219238
aggs,

‎crates/lower/tests/sql_lowering.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,29 @@ async fn multi_aggregate_group_by_binds_columns_positionally() {
111111
);
112112
}
113113

114+
#[tokio::test]
115+
async fn projection_over_aggregate_resolves_output_types_via_output_names() {
116+
// The enclosing Projection references the aggregates by DataFusion's
117+
// generated names (e.g. "sum(metrics.bytes)"); output_names threads those
118+
// onto the L3 Aggregate so the Project resolves real types — not the Utf8
119+
// fallback that an unresolved column would get.
120+
let qe = lower("SELECT SUM(bytes), AVG(latency) FROM metrics").await;
121+
let schema = qe
122+
.output_schema()
123+
.expect("root projection schema derivation");
124+
assert_eq!(schema.columns.len(), 2);
125+
assert_eq!(
126+
schema.columns[0].dtype,
127+
DataType::Int64,
128+
"SUM(bytes:Int64) resolves to Int64, not the Utf8 fallback"
129+
);
130+
assert_eq!(
131+
schema.columns[1].dtype,
132+
DataType::Float64,
133+
"AVG(latency) resolves to Float64"
134+
);
135+
}
136+
114137
#[tokio::test]
115138
async fn count_star_is_count_intent() {
116139
let qe = lower("SELECT COUNT(*) FROM metrics").await;

0 commit comments

Comments
 (0)