Skip to content

Commit d977e68

Browse files
jenniferspjennifersp
committed
Merge branch 'main' into jennifer/doltgres
2 parents 260794c + 81b13e8 commit d977e68

27 files changed

+486
-133
lines changed

‎enginetest/enginetests.go‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5870,6 +5870,18 @@ func TestIndexes(t *testing.T, h Harness) {
58705870
}
58715871
}
58725872

5873+
func TestVectorIndexes(t *testing.T, h Harness) {
5874+
for _, tt := range queries.VectorIndexQueries {
5875+
TestScript(t, h, tt)
5876+
}
5877+
}
5878+
5879+
func TestVectorFunctions(t *testing.T, h Harness) {
5880+
for _, tt := range queries.VectorFunctionQueries {
5881+
TestScript(t, h, tt)
5882+
}
5883+
}
5884+
58735885
func TestIndexPrefix(t *testing.T, h Harness) {
58745886
for _, tt := range queries.IndexPrefixQueries {
58755887
TestScript(t, h, tt)

‎enginetest/memory_engine_test.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -890,6 +890,14 @@ func TestIndexes(t *testing.T) {
890890
enginetest.TestIndexes(t, enginetest.NewDefaultMemoryHarness())
891891
}
892892

893+
func TestVectorIndexes(t *testing.T) {
894+
enginetest.TestVectorIndexes(t, enginetest.NewDefaultMemoryHarness())
895+
}
896+
897+
func TestVectorFunctions(t *testing.T) {
898+
enginetest.TestVectorFunctions(t, enginetest.NewDefaultMemoryHarness())
899+
}
900+
893901
func TestIndexPrefix(t *testing.T) {
894902
enginetest.TestIndexPrefix(t, enginetest.NewDefaultMemoryHarness())
895903
}

‎enginetest/queries/priv_auth_queries.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2268,6 +2268,24 @@ FROM ((SELECT 1 as found FROM information_schema.tables WHERE table_schema = 'te
22682268
},
22692269
},
22702270
},
2271+
{
2272+
Name: "Test user creation with hashed password",
2273+
SetUpScript: []string{
2274+
"CREATE USER 'lol'@'%' IDENTIFIED WITH mysql_native_password AS '*91D9861DFC07DD967611B8C96953474EF270AD5E';",
2275+
},
2276+
Assertions: []UserPrivilegeTestAssertion{
2277+
{
2278+
Query: "SELECT User, plugin, authentication_string FROM mysql.user WHERE User = 'lol';",
2279+
Expected: []sql.Row{
2280+
{
2281+
"lol", // User
2282+
"mysql_native_password", // plugin
2283+
"*91D9861DFC07DD967611B8C96953474EF270AD5E", // authentication_string
2284+
},
2285+
},
2286+
},
2287+
},
2288+
},
22712289
}
22722290

22732291
// NoopPlaintextPlugin is used to authenticate plaintext user plugins

‎enginetest/queries/vector_function_queries.go‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
// Copyright 2024 Dolthub, Inc.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package queries
16+
17+
import (
18+
"github.com/dolthub/go-mysql-server/sql"
19+
"github.com/dolthub/go-mysql-server/sql/types"
20+
)
21+
22+
var VectorFunctionQueries = []ScriptTest{
23+
{
24+
Name: "basic usage of VEC_DISTANCE without index",
25+
SetUpScript: []string{
26+
"create table vectors (id int primary key, v json);",
27+
`insert into vectors values (1, '[3.0,4.0]'), (2, '[0.0,0.0]'), (3, '[1.0,-1.0]'), (4, '[-2.0,0.0]');`,
28+
},
29+
Assertions: []ScriptTestAssertion{
30+
{
31+
Query: "select VEC_DISTANCE('[10.0]', '[20.0]');",
32+
Expected: []sql.Row{{100.0}},
33+
},
34+
{
35+
Query: "select VEC_DISTANCE_L2_SQUARED('[1.0, 2.0]', '[5.0, 5.0]');",
36+
Expected: []sql.Row{{25.0}},
37+
},
38+
{
39+
Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
40+
Expected: []sql.Row{
41+
{2, types.MustJSON(`[0.0, 0.0]`)},
42+
{3, types.MustJSON(`[1.0, -1.0]`)},
43+
{4, types.MustJSON(`[-2.0, 0.0]`)},
44+
{1, types.MustJSON(`[3.0, 4.0]`)},
45+
},
46+
},
47+
{
48+
Query: "select * from vectors order by VEC_DISTANCE_L2_SQUARED('[-2.0,0.0]', v)",
49+
Expected: []sql.Row{
50+
{4, types.MustJSON(`[-2.0, 0.0]`)},
51+
{2, types.MustJSON(`[0.0, 0.0]`)},
52+
{3, types.MustJSON(`[1.0, -1.0]`)},
53+
{1, types.MustJSON(`[3.0, 4.0]`)},
54+
},
55+
},
56+
},
57+
},
58+
}

‎enginetest/queries/vector_index_queries.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
// Copyright 2024 Dolthub, Inc.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package queries
16+
17+
import (
18+
"github.com/dolthub/go-mysql-server/sql"
19+
"github.com/dolthub/go-mysql-server/sql/types"
20+
)
21+
22+
var VectorIndexQueries = []ScriptTest{
23+
{
24+
Name: "basic vector index",
25+
SetUpScript: []string{
26+
"create table vectors (id int primary key, v json);",
27+
`insert into vectors values (1, '[4.0,3.0]'), (2, '[0.0,0.0]'), (3, '[-1.0,1.0]'), (4, '[0.0,-2.0]');`,
28+
`create vector index v_idx on vectors(v);`,
29+
},
30+
Assertions: []ScriptTestAssertion{
31+
{
32+
Query: "show create table vectors",
33+
Expected: []sql.Row{
34+
{"vectors", "CREATE TABLE `vectors` (\n `id` int NOT NULL,\n `v` json,\n PRIMARY KEY (`id`),\n VECTOR KEY `v_idx` (`v`)\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"},
35+
},
36+
},
37+
{
38+
Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v) limit 4",
39+
Expected: []sql.Row{
40+
{2, types.MustJSON(`[0.0, 0.0]`)},
41+
{3, types.MustJSON(`[-1.0, 1.0]`)},
42+
{4, types.MustJSON(`[0.0, -2.0]`)},
43+
{1, types.MustJSON(`[4.0, 3.0]`)},
44+
},
45+
ExpectedIndexes: []string{"v_idx"},
46+
},
47+
{
48+
// Only queries with a limit can use a vector index.
49+
Query: "select * from vectors order by VEC_DISTANCE('[0.0,0.0]', v)",
50+
Expected: []sql.Row{
51+
{2, types.MustJSON(`[0.0, 0.0]`)},
52+
{3, types.MustJSON(`[-1.0, 1.0]`)},
53+
{4, types.MustJSON(`[0.0, -2.0]`)},
54+
{1, types.MustJSON(`[4.0, 3.0]`)},
55+
},
56+
ExpectedIndexes: nil,
57+
},
58+
{
59+
Query: "select * from vectors order by VEC_DISTANCE_L2_SQUARED('[0.0,-2.0]', v) limit 4",
60+
Expected: []sql.Row{
61+
{4, types.MustJSON(`[0.0, -2.0]`)},
62+
{2, types.MustJSON(`[0.0, 0.0]`)},
63+
{3, types.MustJSON(`[-1.0, 1.0]`)},
64+
{1, types.MustJSON(`[4.0, 3.0]`)},
65+
},
66+
ExpectedIndexes: []string{"v_idx"},
67+
},
68+
{
69+
// Ensure vector index is not used for range lookups.
70+
Query: "select * from vectors order by v limit 4",
71+
Expected: []sql.Row{
72+
{3, types.MustJSON(`[-1.0, 1.0]`)},
73+
{4, types.MustJSON(`[0.0, -2.0]`)},
74+
{2, types.MustJSON(`[0.0, 0.0]`)},
75+
{1, types.MustJSON(`[4.0, 3.0]`)},
76+
},
77+
ExpectedIndexes: []string{},
78+
},
79+
},
80+
},
81+
}

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ require (
66
github.com/dolthub/go-icu-regex v0.0.0-20240916130659-0118adc6b662
77
github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71
88
github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81
9-
github.com/dolthub/vitess v0.0.0-20241104125316-860772ba6683
9+
github.com/dolthub/vitess v0.0.0-20241111235433-a20a5ab9d7c9
1010
github.com/go-kit/kit v0.10.0
1111
github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d
1212
github.com/gocraft/dbr/v2 v2.7.2

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,8 @@ github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71 h1:bMGS25NWAGTE
5858
github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71/go.mod h1:2/2zjLQ/JOOSbbSboojeg+cAwcRV0fDLzIiWch/lhqI=
5959
github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81 h1:7/v8q9XGFa6q5Ap4Z/OhNkAMBaK5YeuEzwJt+NZdhiE=
6060
github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81/go.mod h1:siLfyv2c92W1eN/R4QqG/+RjjX5W2+gCTRjZxBjI3TY=
61-
github.com/dolthub/vitess v0.0.0-20241104125316-860772ba6683 h1:2/RJeUfNAXS7mbBnEr9C36htiCJHk5XldDPzhxtEsME=
62-
github.com/dolthub/vitess v0.0.0-20241104125316-860772ba6683/go.mod h1:uBvlRluuL+SbEWTCZ68o0xvsdYZER3CEG/35INdzfJM=
61+
github.com/dolthub/vitess v0.0.0-20241111235433-a20a5ab9d7c9 h1:s36zDuLPuZRWC0nBCJs2Z8joP19eKEtcsIsuE8K9Kx0=
62+
github.com/dolthub/vitess v0.0.0-20241111235433-a20a5ab9d7c9/go.mod h1:uBvlRluuL+SbEWTCZ68o0xvsdYZER3CEG/35INdzfJM=
6363
github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
6464
github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
6565
github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=

‎memory/index.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020

2121
"github.com/dolthub/go-mysql-server/sql"
2222
"github.com/dolthub/go-mysql-server/sql/expression"
23+
"github.com/dolthub/go-mysql-server/sql/expression/function/vector"
2324
"github.com/dolthub/go-mysql-server/sql/fulltext"
2425
"github.com/dolthub/go-mysql-server/sql/types"
2526
)
@@ -38,7 +39,7 @@ type Index struct {
3839
Fulltext bool
3940
// If SupportedVectorFunction is non-nil, this index can be used to optimize ORDER BY
4041
// expressions on this type of distance function.
41-
SupportedVectorFunction expression.DistanceType
42+
SupportedVectorFunction vector.DistanceType
4243
CommentStr string
4344
PrefixLens []uint16
4445
fulltextInfo
@@ -121,11 +122,15 @@ func (idx *Index) IsFullText() bool {
121122
return idx.Fulltext
122123
}
123124

125+
func (idx *Index) IsVector() bool {
126+
return idx.SupportedVectorFunction != nil
127+
}
128+
124129
func (idx *Index) CanSupportOrderBy(expr sql.Expression) bool {
125130
if idx.SupportedVectorFunction == nil {
126131
return false
127132
}
128-
dist, isDist := expr.(*expression.Distance)
133+
dist, isDist := expr.(*vector.Distance)
129134
return isDist && idx.SupportedVectorFunction.CanEval(dist.DistanceType)
130135
}
131136

‎memory/table.go‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import (
2929
"github.com/dolthub/go-mysql-server/sql"
3030
"github.com/dolthub/go-mysql-server/sql/analyzer/analyzererrors"
3131
"github.com/dolthub/go-mysql-server/sql/expression"
32+
"github.com/dolthub/go-mysql-server/sql/expression/function/vector"
3233
"github.com/dolthub/go-mysql-server/sql/fulltext"
3334
"github.com/dolthub/go-mysql-server/sql/iters"
3435
"github.com/dolthub/go-mysql-server/sql/transform"
@@ -1597,7 +1598,6 @@ func (t *IndexedTable) LookupPartitions(ctx *sql.Context, lookup sql.IndexLookup
15971598

15981599
if lookup.VectorOrderAndLimit.OrderBy != nil {
15991600
return &vectorPartitionIter{
1600-
Column: lookup.Index.(*Index).Exprs[0],
16011601
OrderAndLimit: lookup.VectorOrderAndLimit,
16021602
}, nil
16031603
}
@@ -2011,6 +2011,11 @@ func (t *Table) createIndex(data *TableData, name string, columns []sql.IndexCol
20112011
}
20122012
}
20132013

2014+
var vectorFunction vector.DistanceType
2015+
if constraint == sql.IndexConstraint_Vector {
2016+
vectorFunction = vector.DistanceL2Squared{}
2017+
}
2018+
20142019
return &Index{
20152020
DB: t.dbName(),
20162021
DriverName: "",
@@ -2021,7 +2026,7 @@ func (t *Table) createIndex(data *TableData, name string, columns []sql.IndexCol
20212026
Unique: constraint == sql.IndexConstraint_Unique,
20222027
Spatial: constraint == sql.IndexConstraint_Spatial,
20232028
Fulltext: constraint == sql.IndexConstraint_Fulltext,
2024-
SupportedVectorFunction: nil,
2029+
SupportedVectorFunction: vectorFunction,
20252030
CommentStr: comment,
20262031
PrefixLens: prefixLengths,
20272032
}, nil
@@ -2121,7 +2126,7 @@ func (t *Table) CreateFulltextIndex(ctx *sql.Context, indexDef sql.IndexDef, key
21212126
return nil
21222127
}
21232128

2124-
func (t *Table) CreateVectorIndex(ctx *sql.Context, idx sql.IndexDef, distanceType expression.DistanceType) error {
2129+
func (t *Table) CreateVectorIndex(ctx *sql.Context, idx sql.IndexDef, distanceType vector.DistanceType) error {
21252130
if len(idx.Columns) > 1 {
21262131
return fmt.Errorf("vector indexes must have exactly one column")
21272132
}

‎sql/analyzer/index_analyzer_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,7 @@ func (i *dummyIdx) Table() string { return i.table }
148148
func (i *dummyIdx) IsUnique() bool { return false }
149149
func (i *dummyIdx) IsSpatial() bool { return false }
150150
func (i *dummyIdx) IsFullText() bool { return false }
151+
func (i *dummyIdx) IsVector() bool { return false }
151152
func (i *dummyIdx) Comment() string { return "" }
152153
func (i *dummyIdx) IsGenerated() bool { return false }
153154
func (i *dummyIdx) CanSupportOrderBy(sql.Expression) bool { return false }

0 commit comments

Comments
 (0)