Skip to content

Commit 6216399

Browse files
kyleconroyclaude
andcommitted
Support parameterized functions in APPLY column transformers
Add support for parsing APPLY(quantiles(0.5)) and similar parameterized function calls within column transformers. Previously, the parser only handled simple function names like APPLY(sum), but not functions with parameters. Changes: - Add ApplyParams field to ColumnTransformer struct in AST - Update parseColumnsApply and parseAsteriskApply to handle nested parentheses for parameterized functions - Fixes 01470_columns_transformers stmt41, stmt42 - Also fixes 01710_projection_with_column_transformers Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 734ba84 commit 6216399

File tree

4 files changed

+54
-22
lines changed

4 files changed

+54
-22
lines changed

ast/ast.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1428,6 +1428,7 @@ type ColumnTransformer struct {
14281428
Position token.Position `json:"-"`
14291429
Type string `json:"type"` // "apply", "except", "replace"
14301430
Apply string `json:"apply,omitempty"` // function name for APPLY
1431+
ApplyParams []Expression `json:"apply_params,omitempty"` // parameters for parameterized APPLY functions like quantiles(0.5)
14311432
ApplyLambda Expression `json:"apply_lambda,omitempty"` // lambda expression for APPLY x -> expr
14321433
Except []string `json:"except,omitempty"` // column names for EXCEPT
14331434
Pattern string `json:"pattern,omitempty"` // regex pattern for EXCEPT('pattern')

parser/expression.go

Lines changed: 51 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2957,20 +2957,40 @@ func (p *Parser) parseAsteriskApply(asterisk *ast.Asterisk) ast.Expression {
29572957
// Parse lambda expression
29582958
lambda := p.parseExpression(LOWEST)
29592959
asterisk.Transformers = append(asterisk.Transformers, &ast.ColumnTransformer{
2960-
Position: pos,
2961-
Type: "apply",
2962-
ApplyLambda: lambda,
2960+
Position: pos,
2961+
Type: "apply",
2962+
ApplyLambda: lambda,
29632963
})
29642964
} else if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
29652965
// Parse function name (can be IDENT or keyword like sum, avg, etc.)
29662966
funcName := p.current.Value
2967+
p.nextToken()
2968+
2969+
// Check for parameterized function: APPLY(quantiles(0.5))
2970+
var params []ast.Expression
2971+
if p.currentIs(token.LPAREN) {
2972+
p.nextToken() // skip (
2973+
for !p.currentIs(token.RPAREN) && !p.currentIs(token.EOF) {
2974+
expr := p.parseExpression(LOWEST)
2975+
if expr != nil {
2976+
params = append(params, expr)
2977+
}
2978+
if p.currentIs(token.COMMA) {
2979+
p.nextToken()
2980+
} else {
2981+
break
2982+
}
2983+
}
2984+
p.expect(token.RPAREN)
2985+
}
2986+
29672987
asterisk.Apply = append(asterisk.Apply, funcName)
29682988
asterisk.Transformers = append(asterisk.Transformers, &ast.ColumnTransformer{
2969-
Position: pos,
2970-
Type: "apply",
2971-
Apply: funcName,
2989+
Position: pos,
2990+
Type: "apply",
2991+
Apply: funcName,
2992+
ApplyParams: params,
29722993
})
2973-
p.nextToken()
29742994
}
29752995

29762996
if hasParens {
@@ -3002,13 +3022,33 @@ func (p *Parser) parseColumnsApply(matcher *ast.ColumnsMatcher) ast.Expression {
30023022
} else if p.currentIs(token.IDENT) || p.current.Token.IsKeyword() {
30033023
// Parse function name (can be IDENT or keyword like sum, avg, etc.)
30043024
funcName := p.current.Value
3025+
p.nextToken()
3026+
3027+
// Check for parameterized function: APPLY(quantiles(0.5))
3028+
var params []ast.Expression
3029+
if p.currentIs(token.LPAREN) {
3030+
p.nextToken() // skip (
3031+
for !p.currentIs(token.RPAREN) && !p.currentIs(token.EOF) {
3032+
expr := p.parseExpression(LOWEST)
3033+
if expr != nil {
3034+
params = append(params, expr)
3035+
}
3036+
if p.currentIs(token.COMMA) {
3037+
p.nextToken()
3038+
} else {
3039+
break
3040+
}
3041+
}
3042+
p.expect(token.RPAREN)
3043+
}
3044+
30053045
matcher.Apply = append(matcher.Apply, funcName)
30063046
matcher.Transformers = append(matcher.Transformers, &ast.ColumnTransformer{
3007-
Position: pos,
3008-
Type: "apply",
3009-
Apply: funcName,
3047+
Position: pos,
3048+
Type: "apply",
3049+
Apply: funcName,
3050+
ApplyParams: params,
30103051
})
3011-
p.nextToken()
30123052
}
30133053

30143054
if hasParens {
Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt41": true,
4-
"stmt42": true
5-
}
6-
}
1+
{}
Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1 @@
1-
{
2-
"explain_todo": {
3-
"stmt2": true
4-
}
5-
}
1+
{}

0 commit comments

Comments
 (0)