Skip to content

Commit bdde761

Browse files
apstndbapstndb
authored
Support MODEL DDL statements (#195)
* Implement CREATE MODEL, ALTER MODEL, DROP MODEL statements * Update testdata * Make Name Ident * Update testdata * Implement OR REPLACE and IF NOT EXISTS * Fix IF NOT EXISTS * Add newlines to reflect review comments
1 parent 951f4e6 commit bdde761

25 files changed

+1262
-11
lines changed

‎ast/ast.go

+79
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,9 @@ func (CreateSequence) isStatement() {}
9191
func (AlterSequence) isStatement() {}
9292
func (DropSequence) isStatement() {}
9393
func (AlterStatistics) isStatement() {}
94+
func (CreateModel) isStatement() {}
95+
func (AlterModel) isStatement() {}
96+
func (DropModel) isStatement() {}
9497
func (Analyze) isStatement() {}
9598
func (CreateVectorIndex) isStatement() {}
9699
func (DropVectorIndex) isStatement() {}
@@ -336,6 +339,9 @@ func (CreateSequence) isDDL() {}
336339
func (AlterSequence) isDDL() {}
337340
func (DropSequence) isDDL() {}
338341
func (AlterStatistics) isDDL() {}
342+
func (CreateModel) isDDL() {}
343+
func (AlterModel) isDDL() {}
344+
func (DropModel) isDDL() {}
339345
func (Analyze) isDDL() {}
340346
func (CreateVectorIndex) isDDL() {}
341347
func (DropVectorIndex) isDDL() {}
@@ -2910,6 +2916,79 @@ type Analyze struct {
29102916
Analyze token.Pos // position of "ANALYZE" keyword
29112917
}
29122918

2919+
// CreateModelColumn is a single column definition node in CREATE MODEL.
2920+
//
2921+
// {{.Name | sql}} {{.DataType | sql}} {{.Options | sqlOpt}}
2922+
type CreateModelColumn struct {
2923+
// pos = Name.pos
2924+
// end = (Options ?? DataType).end
2925+
2926+
Name *Ident
2927+
DataType SchemaType
2928+
Options *Options // optional
2929+
}
2930+
2931+
// CreateModelInputOutput is INPUT and OUTPUT column list node.
2932+
//
2933+
// INPUT ({{.InputColumns | sqlJoin ", "}}) OUTPUT ({{.OutputColumns | sqlJoin ", "}})
2934+
type CreateModelInputOutput struct {
2935+
// pos = Input
2936+
// end = Rparen + 1
2937+
2938+
Input token.Pos
2939+
Rparen token.Pos // position of the last ")"
2940+
2941+
InputColumns []*CreateModelColumn
2942+
OutputColumns []*CreateModelColumn
2943+
}
2944+
2945+
// CreateModel is CREATE MODEL statement node.
2946+
//
2947+
// CREATE {{if .OrReplace}}OR REPLACE{{end}} MODEL {{if .IfNotExists}}IF NOT EXISTS{{end}} {{.Name | sql}}
2948+
// {{.InputOutput | sqlOpt}}
2949+
// REMOTE
2950+
// {{.Options | sqlOpt}}
2951+
type CreateModel struct {
2952+
// pos = Create
2953+
// end = Options.end || Remote + 6
2954+
2955+
Create token.Pos // position of "CREATE" keyword
2956+
Remote token.Pos // position of "REMOTE" keyword
2957+
2958+
OrReplace bool
2959+
IfNotExists bool
2960+
Name *Ident
2961+
InputOutput *CreateModelInputOutput // optional
2962+
Options *Options // optional
2963+
}
2964+
2965+
// AlterModel is ALTER MODEL statement node.
2966+
//
2967+
// ALTER MODEL {{if .IfExists}}IF EXISTS{{end}} {{.Name | sql}} SET {{.Options | sql}}
2968+
type AlterModel struct {
2969+
// pos = Alter
2970+
// end = Options.end
2971+
2972+
Alter token.Pos
2973+
2974+
IfExists bool
2975+
Name *Ident
2976+
Options *Options
2977+
}
2978+
2979+
// DropModel is DROP MODEL statement node.
2980+
//
2981+
// DROP MODEL {{if .IfExists}}IF EXISTS{{end}} {{.Name | sql}}
2982+
type DropModel struct {
2983+
// pos = Drop
2984+
// end = Name.end
2985+
2986+
Drop token.Pos
2987+
2988+
IfExists bool
2989+
Name *Ident
2990+
}
2991+
29132992
// ================================================================================
29142993
//
29152994
// Types for Schema

‎ast/pos.go

+40
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎ast/sql.go

+29
Original file line numberDiff line numberDiff line change
@@ -1280,6 +1280,35 @@ func (s *AlterStatistics) SQL() string {
12801280
}
12811281
func (a *Analyze) SQL() string { return "ANALYZE" }
12821282

1283+
func (c *CreateModelColumn) SQL() string {
1284+
return c.Name.SQL() + " " + c.DataType.SQL() + sqlOpt(" ", c.Options, "")
1285+
}
1286+
1287+
func (c *CreateModelInputOutput) SQL() string {
1288+
return "INPUT (" + sqlJoin(c.InputColumns, ", ") + ") OUTPUT (" + sqlJoin(c.OutputColumns, ", ") + ")"
1289+
}
1290+
1291+
func (c *CreateModel) SQL() string {
1292+
return "CREATE " + strOpt(c.OrReplace, "OR REPLACE ") +
1293+
"MODEL " +
1294+
c.Name.SQL() +
1295+
strOpt(c.IfNotExists, " IF NOT EXISTS") +
1296+
sqlOpt(" ", c.InputOutput, "") +
1297+
" REMOTE" +
1298+
sqlOpt(" ", c.Options, "")
1299+
}
1300+
1301+
func (a *AlterModel) SQL() string {
1302+
return "ALTER MODEL " +
1303+
strOpt(a.IfExists, "IF EXISTS ") +
1304+
a.Name.SQL() +
1305+
" SET " + a.Options.SQL()
1306+
}
1307+
1308+
func (d *DropModel) SQL() string {
1309+
return "DROP MODEL " + strOpt(d.IfExists, "IF EXISTS ") + d.Name.SQL()
1310+
}
1311+
12831312
// ================================================================================
12841313
//
12851314
// Types for Schema

‎parser.go

+102-11
Original file line numberDiff line numberDiff line change
@@ -2333,8 +2333,8 @@ func (p *Parser) parseDDL() ast.DDL {
23332333
return p.parseCreateTable(pos)
23342334
case p.Token.IsKeywordLike("SEQUENCE"):
23352335
return p.parseCreateSequence(pos)
2336-
case p.Token.IsKeywordLike("VIEW") || p.Token.Kind == "OR":
2337-
return p.parseCreateView(pos)
2336+
case p.Token.IsKeywordLike("VIEW"):
2337+
return p.parseCreateView(pos, false)
23382338
case p.Token.IsKeywordLike("INDEX") || p.Token.IsKeywordLike("UNIQUE") || p.Token.IsKeywordLike("NULL_FILTERED"):
23392339
return p.parseCreateIndex(pos)
23402340
case p.Token.IsKeywordLike("VECTOR"):
@@ -2345,8 +2345,19 @@ func (p *Parser) parseDDL() ast.DDL {
23452345
return p.parseCreateRole(pos)
23462346
case p.Token.IsKeywordLike("CHANGE"):
23472347
return p.parseCreateChangeStream(pos)
2348+
case p.Token.IsKeywordLike("MODEL"):
2349+
return p.parseCreateModel(pos, false)
2350+
case p.Token.Kind == "OR":
2351+
p.expect("OR")
2352+
p.expectKeywordLike("REPLACE")
2353+
switch {
2354+
case p.Token.IsKeywordLike("VIEW"):
2355+
return p.parseCreateView(pos, true)
2356+
case p.Token.IsKeywordLike("MODEL"):
2357+
return p.parseCreateModel(pos, true)
2358+
}
23482359
}
2349-
p.panicfAtToken(&p.Token, "expected pseudo keyword: DATABASE, TABLE, INDEX, UNIQUE, NULL_FILTERED, ROLE, CHANGE but: %s", p.Token.AsString)
2360+
p.panicfAtToken(&p.Token, "expected pseudo keyword: DATABASE, TABLE, INDEX, UNIQUE, NULL_FILTERED, ROLE, CHANGE, OR but: %s", p.Token.AsString)
23502361
case p.Token.IsKeywordLike("ALTER"):
23512362
p.nextToken()
23522363
switch {
@@ -2366,6 +2377,8 @@ func (p *Parser) parseDDL() ast.DDL {
23662377
return p.parseAlterChangeStream(pos)
23672378
case p.Token.IsKeywordLike("STATISTICS"):
23682379
return p.parseAlterStatistics(pos)
2380+
case p.Token.IsKeywordLike("MODEL"):
2381+
return p.parseAlterModel(pos)
23692382
}
23702383
p.panicfAtToken(&p.Token, "expected pseudo keyword: TABLE, CHANGE, but: %s", p.Token.AsString)
23712384
case p.Token.IsKeywordLike("DROP"):
@@ -2391,8 +2404,10 @@ func (p *Parser) parseDDL() ast.DDL {
23912404
return p.parseDropRole(pos)
23922405
case p.Token.IsKeywordLike("CHANGE"):
23932406
return p.parseDropChangeStream(pos)
2407+
case p.Token.IsKeywordLike("MODEL"):
2408+
return p.parseDropModel(pos)
23942409
}
2395-
p.panicfAtToken(&p.Token, "expected pseudo keyword: TABLE, INDEX, ROLE, CHANGE, but: %s", p.Token.AsString)
2410+
p.panicfAtToken(&p.Token, "expected pseudo keyword: TABLE, INDEX, ROLE, CHANGE, MODEL, but: %s", p.Token.AsString)
23962411
case p.Token.IsKeywordLike("RENAME"):
23972412
p.nextToken()
23982413
return p.parseRenameTable(pos)
@@ -2603,13 +2618,7 @@ func (p *Parser) parseCreateSequence(pos token.Pos) *ast.CreateSequence {
26032618
}
26042619
}
26052620

2606-
func (p *Parser) parseCreateView(pos token.Pos) *ast.CreateView {
2607-
var orReplace bool
2608-
if p.Token.Kind == "OR" {
2609-
p.nextToken()
2610-
p.expectKeywordLike("REPLACE")
2611-
orReplace = true
2612-
}
2621+
func (p *Parser) parseCreateView(pos token.Pos, orReplace bool) *ast.CreateView {
26132622
p.expectKeywordLike("VIEW")
26142623

26152624
name := p.parsePath()
@@ -3783,6 +3792,88 @@ func (p *Parser) parseAnalyze() *ast.Analyze {
37833792
}
37843793
}
37853794

3795+
func (p *Parser) tryParseCreateModelColumn() *ast.CreateModelColumn {
3796+
name := p.parseIdent()
3797+
dataType := p.parseSchemaType()
3798+
options := p.tryParseOptions()
3799+
3800+
return &ast.CreateModelColumn{
3801+
Name: name,
3802+
DataType: dataType,
3803+
Options: options,
3804+
}
3805+
}
3806+
3807+
func (p *Parser) tryParseCreateModelInputOutput() *ast.CreateModelInputOutput {
3808+
if !p.Token.IsKeywordLike("INPUT") {
3809+
return nil
3810+
}
3811+
3812+
pos := p.expectKeywordLike("INPUT").Pos
3813+
p.expect("(")
3814+
inputColumns := parseCommaSeparatedList(p, p.tryParseCreateModelColumn)
3815+
p.expect(")")
3816+
3817+
p.expectKeywordLike("OUTPUT")
3818+
p.expect("(")
3819+
outputColumns := parseCommaSeparatedList(p, p.tryParseCreateModelColumn)
3820+
rparen := p.expect(")").Pos
3821+
3822+
return &ast.CreateModelInputOutput{
3823+
Input: pos,
3824+
Rparen: rparen,
3825+
InputColumns: inputColumns,
3826+
OutputColumns: outputColumns,
3827+
}
3828+
}
3829+
3830+
func (p *Parser) parseCreateModel(pos token.Pos, orReplace bool) *ast.CreateModel {
3831+
p.expectKeywordLike("MODEL")
3832+
name := p.parseIdent()
3833+
ifNotExists := p.parseIfNotExists()
3834+
inputOutput := p.tryParseCreateModelInputOutput()
3835+
remote := p.expectKeywordLike("REMOTE").Pos
3836+
options := p.tryParseOptions()
3837+
3838+
return &ast.CreateModel{
3839+
Create: pos,
3840+
OrReplace: orReplace,
3841+
IfNotExists: ifNotExists,
3842+
Name: name,
3843+
InputOutput: inputOutput,
3844+
Remote: remote,
3845+
Options: options,
3846+
}
3847+
3848+
}
3849+
3850+
func (p *Parser) parseAlterModel(pos token.Pos) *ast.AlterModel {
3851+
p.expectKeywordLike("MODEL")
3852+
ifExists := p.parseIfExists()
3853+
name := p.parseIdent()
3854+
p.expect("SET")
3855+
options := p.parseOptions()
3856+
3857+
return &ast.AlterModel{
3858+
Alter: pos,
3859+
IfExists: ifExists,
3860+
Name: name,
3861+
Options: options,
3862+
}
3863+
}
3864+
3865+
func (p *Parser) parseDropModel(pos token.Pos) *ast.DropModel {
3866+
p.expectKeywordLike("MODEL")
3867+
ifExists := p.parseIfExists()
3868+
name := p.parseIdent()
3869+
3870+
return &ast.DropModel{
3871+
Drop: pos,
3872+
IfExists: ifExists,
3873+
Name: name,
3874+
}
3875+
}
3876+
37863877
var scalarSchemaTypes = []string{
37873878
"BOOL",
37883879
"INT64",

‎testdata/input/ddl/alter_model.sql

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
ALTER MODEL MyClassificationModel
2+
SET OPTIONS (
3+
endpoints = [
4+
'//aiplatform.googleapis.com/projects/aaa/locations/tl/endpoints/aaa',
5+
'//aiplatform.googleapis.com/projects/aaa/locations/tl/endpoints/bbb'
6+
],
7+
default_batch_size = 100
8+
)

‎testdata/input/ddl/alter_model_if_exists.sql

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
ALTER MODEL IF EXISTS MyClassificationModel
2+
SET OPTIONS (
3+
endpoints = [
4+
'//aiplatform.googleapis.com/projects/aaa/locations/tl/endpoints/aaa',
5+
'//aiplatform.googleapis.com/projects/aaa/locations/tl/endpoints/bbb'
6+
],
7+
default_batch_size = 100
8+
)

‎testdata/input/ddl/create_model.sql

+14
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
CREATE MODEL MyClassificationModel
2+
INPUT (
3+
length FLOAT64,
4+
material STRING(MAX),
5+
tag_array ARRAY<STRING(MAX)>
6+
)
7+
OUTPUT (
8+
scores ARRAY<FLOAT64>,
9+
classes ARRAY<STRING(MAX)>
10+
)
11+
REMOTE
12+
OPTIONS (
13+
endpoint = '//aiplatform.googleapis.com/projects/PROJECT/locations/LOCATION/endpoints/ENDPOINT_ID'
14+
)

‎testdata/input/ddl/create_model_if_not_exists.sql

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
CREATE MODEL GeminiPro IF NOT EXISTS
2+
INPUT (prompt STRING(MAX))
3+
OUTPUT (content STRING(MAX))
4+
REMOTE OPTIONS (
5+
endpoint = '//aiplatform.googleapis.com/projects/fake-project/locations/asia-northeast1/publishers/google/models/gemini-pro',
6+
default_batch_size = 1
7+
)

‎testdata/input/ddl/create_or_replace_model.sql

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
CREATE OR REPLACE MODEL GeminiPro
2+
INPUT (prompt STRING(MAX))
3+
OUTPUT (content STRING(MAX))
4+
REMOTE OPTIONS (
5+
endpoint = '//aiplatform.googleapis.com/projects/fake-project/locations/asia-northeast1/publishers/google/models/gemini-pro',
6+
default_batch_size = 1
7+
)

‎testdata/input/ddl/drop_model.sql

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
DROP MODEL MyClassificationModel

‎testdata/input/ddl/drop_model_if_exists.sql

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
DROP MODEL IF EXISTS MyClassificationModel

0 commit comments

Comments
 (0)