Skip to content

Commit 947e6a9

Browse files
committed
Add lambda support for APPLY transformer on * and COLUMNS()
Adds support for lambda expressions in APPLY transformers: - `* APPLY x -> toString(x)` - `COLUMNS(...) APPLY x -> expr` Previously only simple function names were supported (e.g., `* APPLY toString`). This adds: - ApplyLambda field to ast.ColumnTransformer struct - Lambda detection and parsing in parseAsteriskApply and parseColumnsApply Fixes 14 statements in 02378_analyzer_projection_names (16→2 pending) and improves several other tests.
1 parent d1edb0d commit 947e6a9

9 files changed

Lines changed: 35 additions & 61 deletions

File tree

ast/ast.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1215,11 +1215,12 @@ func (r *ReplaceExpr) End() token.Position { return r.Position }
12151215

12161216
// ColumnTransformer represents a single transformer (APPLY, EXCEPT, or REPLACE) in order.
12171217
type ColumnTransformer struct {
1218-
Position token.Position `json:"-"`
1219-
Type string `json:"type"` // "apply", "except", "replace"
1220-
Apply string `json:"apply,omitempty"` // function name for APPLY
1221-
Except []string `json:"except,omitempty"` // column names for EXCEPT
1222-
Replaces []*ReplaceExpr `json:"replaces,omitempty"` // replacement expressions for REPLACE
1218+
Position token.Position `json:"-"`
1219+
Type string `json:"type"` // "apply", "except", "replace"
1220+
Apply string `json:"apply,omitempty"` // function name for APPLY
1221+
ApplyLambda Expression `json:"apply_lambda,omitempty"` // lambda expression for APPLY x -> expr
1222+
Except []string `json:"except,omitempty"` // column names for EXCEPT
1223+
Replaces []*ReplaceExpr `json:"replaces,omitempty"` // replacement expressions for REPLACE
12231224
}
12241225

12251226
// ColumnsMatcher represents COLUMNS('pattern') or COLUMNS(col1, col2) expression.

parser/expression.go

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2557,8 +2557,17 @@ func (p *Parser) parseAsteriskApply(asterisk *ast.Asterisk) ast.Expression {
25572557
p.nextToken() // skip (
25582558
}
25592559

2560-
// Parse function name (can be IDENT or keyword like sum, avg, etc.)
2561-
if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2560+
// Check for lambda expression: x -> expr
2561+
if p.currentIs(token.IDENT) && p.peekIs(token.ARROW) {
2562+
// Parse lambda expression
2563+
lambda := p.parseExpression(LOWEST)
2564+
asterisk.Transformers = append(asterisk.Transformers, &ast.ColumnTransformer{
2565+
Position: pos,
2566+
Type: "apply",
2567+
ApplyLambda: lambda,
2568+
})
2569+
} else if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2570+
// Parse function name (can be IDENT or keyword like sum, avg, etc.)
25622571
funcName := p.current.Value
25632572
asterisk.Apply = append(asterisk.Apply, funcName)
25642573
asterisk.Transformers = append(asterisk.Transformers, &ast.ColumnTransformer{
@@ -2586,8 +2595,17 @@ func (p *Parser) parseColumnsApply(matcher *ast.ColumnsMatcher) ast.Expression {
25862595
p.nextToken() // skip (
25872596
}
25882597

2589-
// Parse function name (can be IDENT or keyword like sum, avg, etc.)
2590-
if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2598+
// Check for lambda expression: x -> expr
2599+
if p.currentIs(token.IDENT) && p.peekIs(token.ARROW) {
2600+
// Parse lambda expression
2601+
lambda := p.parseExpression(LOWEST)
2602+
matcher.Transformers = append(matcher.Transformers, &ast.ColumnTransformer{
2603+
Position: pos,
2604+
Type: "apply",
2605+
ApplyLambda: lambda,
2606+
})
2607+
} else if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
2608+
// Parse function name (can be IDENT or keyword like sum, avg, etc.)
25912609
funcName := p.current.Value
25922610
matcher.Apply = append(matcher.Apply, funcName)
25932611
matcher.Transformers = append(matcher.Transformers, &ast.ColumnTransformer{
Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
{
22
"explain_todo": {
3-
"stmt4": true,
4-
"stmt6": true,
5-
"stmt7": true
3+
"stmt4": true
64
}
75
}
Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt63": true,
4-
"stmt64": true,
5-
"stmt66": true,
6-
"stmt67": true,
7-
"stmt69": true,
8-
"stmt70": true
9-
}
10-
}
1+
{}
Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt31": true,
4-
"stmt32": true,
5-
"stmt33": true
6-
}
7-
}
1+
{}
Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,6 @@
11
{
22
"explain_todo": {
3-
"stmt101": true,
4-
"stmt105": true,
5-
"stmt109": true,
63
"stmt185": true,
7-
"stmt207": true,
8-
"stmt215": true,
9-
"stmt227": true,
10-
"stmt235": true,
11-
"stmt33": true,
12-
"stmt37": true,
13-
"stmt41": true,
14-
"stmt57": true,
15-
"stmt67": true,
16-
"stmt89": true,
17-
"stmt93": true,
18-
"stmt97": true
4+
"stmt67": true
195
}
206
}
Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt20": true,
4-
"stmt22": true,
5-
"stmt24": true
6-
}
7-
}
1+
{}
Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt3": true
4-
}
5-
}
1+
{}
Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt3": true
4-
}
5-
}
1+
{}

0 commit comments

Comments
 (0)