Skip to content

Commit

Permalink
Support vector_length (#222)
Browse files Browse the repository at this point in the history 
* Support vector_length

* Update testdata
  • Loading branch information
@apstndb
apstndb authored Dec 18, 2024
1 parent eccf53d commit bf9af9a
Show file tree
Hide file tree
Showing 9 changed files with 170 additions and 54 deletions.
12 changes: 7 additions & 5 deletions ast/ast.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3089,15 +3089,17 @@ type SizedSchemaType struct {

// ArraySchemaType is array type node in schema.
//
// ARRAY<{{.Item | sql}}>
// ARRAY<{{.Item | sql}}>{{if .NamedArgs}}({{.NamedArgs | sqlJoin ", "}}){{end}}
type ArraySchemaType struct {
// pos = Array
// end = Gt + 1
// end = Rparen + 1 || Gt + 1

Array token.Pos // position of "ARRAY" keyword
Gt token.Pos // position of ">"
Array token.Pos // position of "ARRAY" keyword
Gt token.Pos // position of ">"
Rparen token.Pos // position of ")" when len(NamedArgs) > 0

Item SchemaType // ScalarSchemaType or SizedSchemaType
Item SchemaType // ScalarSchemaType or SizedSchemaType
NamedArgs []*NamedArg
}

// ================================================================================
Expand Down
2 changes: 1 addition & 1 deletion ast/pos.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion ast/sql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1335,7 +1335,7 @@ func (s *SizedSchemaType) SQL() string {
}

func (a *ArraySchemaType) SQL() string {
return "ARRAY<" + a.Item.SQL() + ">"
return "ARRAY<" + a.Item.SQL() + ">" + strOpt(len(a.NamedArgs) > 0, "("+sqlJoin(a.NamedArgs, ", ")+")")
}

// ================================================================================
Expand Down
35 changes: 26 additions & 9 deletions parser.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1577,19 +1577,25 @@ func (p *Parser) lookaheadNamedArg() bool {
return p.Token.Kind == "=>"
}

func (p *Parser) tryParseNamedArg() *ast.NamedArg {
if !p.lookaheadNamedArg() {
return nil
}
func (p *Parser) parseNamedArg() *ast.NamedArg {
name := p.parseIdent()
p.expect("=>")
value := p.parseExpr()

return &ast.NamedArg{
Name: name,
Value: value,
}
}

func (p *Parser) tryParseNamedArg() *ast.NamedArg {
if !p.lookaheadNamedArg() {
return nil
}

return p.parseNamedArg()
}

func (p *Parser) lookaheadLambdaArg() bool {
lexer := p.Lexer.Clone()
defer func() {
Expand Down Expand Up @@ -3916,10 +3922,21 @@ func (p *Parser) parseSchemaType() ast.SchemaType {
p.expect("<")
t := p.parseScalarSchemaType()
end := p.expect(">").Pos

var namedArgs []*ast.NamedArg
rparen := token.InvalidPos
if p.Token.Kind == "(" {
p.nextToken()
namedArgs = parseCommaSeparatedList(p, p.parseNamedArg)
rparen = p.expect(")").Pos
}

return &ast.ArraySchemaType{
Array: pos,
Gt: end,
Item: t,
Array: pos,
Gt: end,
Item: t,
NamedArgs: namedArgs,
Rparen: rparen,
}
}

Expand Down Expand Up @@ -4134,9 +4151,9 @@ func (p *Parser) tryParseWithAction() *ast.WithAction {
alias := p.tryParseAsAlias(withRequiredAs)

return &ast.WithAction{
With: with,
With: with,
Action: action,
Alias: alias,
Alias: alias,
}
}

Expand Down
1 change: 1 addition & 0 deletions testdata/input/ddl/create_table_types.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,5 @@ create table types (
ab array<bool>,
abs array<bytes(max)>,
p examples.ProtoType,
af32vl array<float32>(vector_length=>128),
) primary key (i)
24 changes: 15 additions & 9 deletions testdata/result/ddl/create_model.sql.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,15 +62,17 @@ OPTIONS (
Name: "tag_array",
},
DataType: &ast.ArraySchemaType{
Array: 97,
Gt: 114,
Item: &ast.SizedSchemaType{
Array: 97,
Gt: 114,
Rparen: -1,
Item: &ast.SizedSchemaType{
NamePos: 103,
Rparen: 113,
Name: "STRING",
Max: true,
Size: nil,
},
NamedArgs: []*ast.NamedArg(nil),
},
Options: (*ast.Options)(nil),
},
Expand All @@ -83,12 +85,14 @@ OPTIONS (
Name: "scores",
},
DataType: &ast.ArraySchemaType{
Array: 136,
Gt: 149,
Item: &ast.ScalarSchemaType{
Array: 136,
Gt: 149,
Rparen: -1,
Item: &ast.ScalarSchemaType{
NamePos: 142,
Name: "FLOAT64",
},
NamedArgs: []*ast.NamedArg(nil),
},
Options: (*ast.Options)(nil),
},
Expand All @@ -99,15 +103,17 @@ OPTIONS (
Name: "classes",
},
DataType: &ast.ArraySchemaType{
Array: 162,
Gt: 179,
Item: &ast.SizedSchemaType{
Array: 162,
Gt: 179,
Rparen: -1,
Item: &ast.SizedSchemaType{
NamePos: 168,
Rparen: 178,
Name: "STRING",
Max: true,
Size: nil,
},
NamedArgs: []*ast.NamedArg(nil),
},
Options: (*ast.Options)(nil),
},
Expand Down
62 changes: 52 additions & 10 deletions testdata/result/ddl/create_table_types.sql.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,13 @@ create table types (
ab array<bool>,
abs array<bytes(max)>,
p examples.ProtoType,
af32vl array<float32>(vector_length=>128),
) primary key (i)

--- AST
&ast.CreateTable{
Create: 0,
Rparen: 252,
Rparen: 297,
IfNotExists: false,
Name: &ast.Path{
Idents: []*ast.Ident{
Expand Down Expand Up @@ -230,12 +231,14 @@ create table types (
Name: "ab",
},
Type: &ast.ArraySchemaType{
Array: 174,
Gt: 184,
Item: &ast.ScalarSchemaType{
Array: 174,
Gt: 184,
Rparen: -1,
Item: &ast.ScalarSchemaType{
NamePos: 180,
Name: "BOOL",
},
NamedArgs: []*ast.NamedArg(nil),
},
NotNull: false,
DefaultExpr: (*ast.ColumnDefaultExpr)(nil),
Expand All @@ -251,15 +254,17 @@ create table types (
Name: "abs",
},
Type: &ast.ArraySchemaType{
Array: 193,
Gt: 209,
Item: &ast.SizedSchemaType{
Array: 193,
Gt: 209,
Rparen: -1,
Item: &ast.SizedSchemaType{
NamePos: 199,
Rparen: 208,
Name: "BYTES",
Max: true,
Size: nil,
},
NamedArgs: []*ast.NamedArg(nil),
},
NotNull: false,
DefaultExpr: (*ast.ColumnDefaultExpr)(nil),
Expand Down Expand Up @@ -294,14 +299,51 @@ create table types (
Hidden: -1,
Options: (*ast.Options)(nil),
},
&ast.ColumnDef{
Null: -1,
Name: &ast.Ident{
NamePos: 238,
NameEnd: 244,
Name: "af32vl",
},
Type: &ast.ArraySchemaType{
Array: 245,
Gt: 258,
Rparen: 278,
Item: &ast.ScalarSchemaType{
NamePos: 251,
Name: "FLOAT32",
},
NamedArgs: []*ast.NamedArg{
&ast.NamedArg{
Name: &ast.Ident{
NamePos: 260,
NameEnd: 273,
Name: "vector_length",
},
Value: &ast.IntLiteral{
ValuePos: 275,
ValueEnd: 278,
Base: 10,
Value: "128",
},
},
},
},
NotNull: false,
DefaultExpr: (*ast.ColumnDefaultExpr)(nil),
GeneratedExpr: (*ast.GeneratedColumnExpr)(nil),
Hidden: -1,
Options: (*ast.Options)(nil),
},
},
TableConstraints: []*ast.TableConstraint(nil),
PrimaryKeys: []*ast.IndexKey{
&ast.IndexKey{
DirPos: -1,
Name: &ast.Ident{
NamePos: 251,
NameEnd: 252,
NamePos: 296,
NameEnd: 297,
Name: "i",
},
Dir: "",
Expand All @@ -313,4 +355,4 @@ create table types (
}

--- SQL
CREATE TABLE types (b BOOL, i INT64, f32 FLOAT32, f FLOAT64, d DATE, t TIMESTAMP, s STRING(256), smax STRING(MAX), bs BYTES(256), bsmax BYTES(MAX), ab ARRAY<BOOL>, abs ARRAY<BYTES(MAX)>, p examples.ProtoType) PRIMARY KEY (i)
CREATE TABLE types (b BOOL, i INT64, f32 FLOAT32, f FLOAT64, d DATE, t TIMESTAMP, s STRING(256), smax STRING(MAX), bs BYTES(256), bsmax BYTES(MAX), ab ARRAY<BOOL>, abs ARRAY<BYTES(MAX)>, p examples.ProtoType, af32vl ARRAY<FLOAT32>(vector_length => 128)) PRIMARY KEY (i)
24 changes: 15 additions & 9 deletions testdata/result/statement/create_model.sql.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,15 +62,17 @@ OPTIONS (
Name: "tag_array",
},
DataType: &ast.ArraySchemaType{
Array: 97,
Gt: 114,
Item: &ast.SizedSchemaType{
Array: 97,
Gt: 114,
Rparen: -1,
Item: &ast.SizedSchemaType{
NamePos: 103,
Rparen: 113,
Name: "STRING",
Max: true,
Size: nil,
},
NamedArgs: []*ast.NamedArg(nil),
},
Options: (*ast.Options)(nil),
},
Expand All @@ -83,12 +85,14 @@ OPTIONS (
Name: "scores",
},
DataType: &ast.ArraySchemaType{
Array: 136,
Gt: 149,
Item: &ast.ScalarSchemaType{
Array: 136,
Gt: 149,
Rparen: -1,
Item: &ast.ScalarSchemaType{
NamePos: 142,
Name: "FLOAT64",
},
NamedArgs: []*ast.NamedArg(nil),
},
Options: (*ast.Options)(nil),
},
Expand All @@ -99,15 +103,17 @@ OPTIONS (
Name: "classes",
},
DataType: &ast.ArraySchemaType{
Array: 162,
Gt: 179,
Item: &ast.SizedSchemaType{
Array: 162,
Gt: 179,
Rparen: -1,
Item: &ast.SizedSchemaType{
NamePos: 168,
Rparen: 178,
Name: "STRING",
Max: true,
Size: nil,
},
NamedArgs: []*ast.NamedArg(nil),
},
Options: (*ast.Options)(nil),
},
Expand Down
Loading

0 comments on commit bf9af9a

Please sign in to comment.