Skip to content

Commit 8516d55

Browse files
ti-chi-botti-chi-bot
authored
planner: support vector index in planner (#56204) (#56546)
ref #54245, close #56510
1 parent 1e24d39 commit 8516d55

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

46 files changed

+1414
-157
lines changed

‎DEPS.bzl

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5815,13 +5815,13 @@ def go_deps():
58155815
name = "com_github_pingcap_tipb",
58165816
build_file_proto_mode = "disable_global",
58175817
importpath = "github.com/pingcap/tipb",
5818-
sha256 = "119f73cede33a7d3712a9fedeb80f806228ecf1ec5258112b2b62af88731a304",
5819-
strip_prefix = "github.com/pingcap/[email protected]20240919023442-cf70966bef25",
5818+
sha256 = "b39e154272ba36d145c6049947a012a76be740b32a44a46d7253caa145c56cc9",
5819+
strip_prefix = "github.com/pingcap/[email protected]20241008083645-0bcddae67837",
58205820
urls = [
5821-
"http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
5822-
"http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
5823-
"https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
5824-
"https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
5821+
"http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
5822+
"http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
5823+
"https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
5824+
"https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
58255825
],
58265826
)
58275827
go_repository(

‎go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ require (
8989
github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d
9090
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5
9191
github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e
92-
github.com/pingcap/tipb v0.0.0-20240919023442-cf70966bef25
92+
github.com/pingcap/tipb v0.0.0-20241008083645-0bcddae67837
9393
github.com/prometheus/client_golang v1.20.4
9494
github.com/prometheus/client_model v0.6.1
9595
github.com/prometheus/common v0.57.0

‎go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -676,8 +676,8 @@ github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d h1:y3EueKVfVykdpTyfU
676676
github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d/go.mod h1:ORfBOFp1eteu2odzsyaxI+b8TzJwgjwyQcGhI+9SfEA=
677677
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeAhmOQHNPRRogMYxP13V7BkYw3ZsoSfE=
678678
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530=
679-
github.com/pingcap/tipb v0.0.0-20240919023442-cf70966bef25 h1:qH261jtqDvqWhhsnle+eURWV+PmFdqEGP/aHDqPFJAs=
680-
github.com/pingcap/tipb v0.0.0-20240919023442-cf70966bef25/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
679+
github.com/pingcap/tipb v0.0.0-20241008083645-0bcddae67837 h1:tyIymn821fB8gUmqafdvLlcFkVOpgyJXImoYJ8n9oJE=
680+
github.com/pingcap/tipb v0.0.0-20241008083645-0bcddae67837/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
681681
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
682682
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
683683
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=

‎pkg/ddl/BUILD.bazel

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -325,7 +325,6 @@ go_test(
325325
"//pkg/store/gcworker",
326326
"//pkg/store/helper",
327327
"//pkg/store/mockstore",
328-
"//pkg/store/mockstore/unistore",
329328
"//pkg/table",
330329
"//pkg/table/tables",
331330
"//pkg/tablecodec",
@@ -353,7 +352,6 @@ go_test(
353352
"@com_github_ngaut_pools//:pools",
354353
"@com_github_pingcap_errors//:errors",
355354
"@com_github_pingcap_failpoint//:failpoint",
356-
"@com_github_pingcap_kvproto//pkg/metapb",
357355
"@com_github_stretchr_testify//assert",
358356
"@com_github_stretchr_testify//require",
359357
"@com_github_tikv_client_go_v2//oracle",

‎pkg/ddl/cancel_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"github.com/pingcap/tidb/pkg/domain/infosync"
2828
"github.com/pingcap/tidb/pkg/errno"
2929
"github.com/pingcap/tidb/pkg/meta/model"
30+
"github.com/pingcap/tidb/pkg/store/mockstore"
3031
"github.com/pingcap/tidb/pkg/testkit"
3132
"github.com/pingcap/tidb/pkg/testkit/external"
3233
"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
@@ -223,7 +224,7 @@ func TestCancelVariousJobs(t *testing.T) {
223224
return enterCnt.Load() == exitCnt.Load()
224225
}, 10*time.Second, 10*time.Millisecond)
225226
}
226-
store := testkit.CreateMockStoreWithSchemaLease(t, 100*time.Millisecond, withMockTiFlash(2))
227+
store := testkit.CreateMockStoreWithSchemaLease(t, 100*time.Millisecond, mockstore.WithMockTiFlash(2))
227228
tk := testkit.NewTestKit(t, store)
228229
tkCancel := testkit.NewTestKit(t, store)
229230

‎pkg/ddl/db_change_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ import (
3737
sessiontypes "github.com/pingcap/tidb/pkg/session/types"
3838
"github.com/pingcap/tidb/pkg/sessionctx"
3939
"github.com/pingcap/tidb/pkg/sessiontxn"
40+
"github.com/pingcap/tidb/pkg/store/mockstore"
4041
"github.com/pingcap/tidb/pkg/testkit"
4142
"github.com/pingcap/tidb/pkg/testkit/external"
4243
"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
@@ -1159,7 +1160,7 @@ func TestParallelAlterAddIndex(t *testing.T) {
11591160
}
11601161

11611162
func TestParallelAlterAddVectorIndex(t *testing.T) {
1162-
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(2))
1163+
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
11631164
tk := testkit.NewTestKit(t, store)
11641165
tk.MustExec("create database test_db_state default charset utf8 default collate utf8_bin")
11651166
tk.MustExec("use test_db_state")

‎pkg/ddl/index.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -385,9 +385,9 @@ func buildVectorInfoWithCheck(indexPartSpecifications []*ast.IndexPartSpecificat
385385
if !ok {
386386
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs(fmt.Sprintf("unsupported function: %v", idxPart.Expr))
387387
}
388-
distanceMetric, ok := model.FnNameToDistanceMetric[f.FnName.L]
388+
distanceMetric, ok := model.IndexableFnNameToDistanceMetric[f.FnName.L]
389389
if !ok {
390-
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs("unsupported function")
390+
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs("currently only L2 and Cosine distance is indexable")
391391
}
392392
colExpr, ok := f.Args[0].(*ast.ColumnNameExpr)
393393
if !ok {

‎pkg/ddl/index_modify_test.go

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@ import (
2626
"time"
2727

2828
"github.com/pingcap/errors"
29-
"github.com/pingcap/kvproto/pkg/metapb"
3029
"github.com/pingcap/tidb/pkg/ddl"
3130
testddlutil "github.com/pingcap/tidb/pkg/ddl/testutil"
3231
"github.com/pingcap/tidb/pkg/domain/infosync"
@@ -42,7 +41,6 @@ import (
4241
"github.com/pingcap/tidb/pkg/sessionctx/variable"
4342
"github.com/pingcap/tidb/pkg/sessiontxn"
4443
"github.com/pingcap/tidb/pkg/store/mockstore"
45-
"github.com/pingcap/tidb/pkg/store/mockstore/unistore"
4644
"github.com/pingcap/tidb/pkg/table"
4745
"github.com/pingcap/tidb/pkg/table/tables"
4846
"github.com/pingcap/tidb/pkg/tablecodec"
@@ -56,7 +54,6 @@ import (
5654
"github.com/pingcap/tidb/pkg/util/dbterror"
5755
"github.com/pingcap/tidb/pkg/util/sqlexec"
5856
"github.com/stretchr/testify/require"
59-
"github.com/tikv/client-go/v2/testutils"
6057
)
6158

6259
const indexModifyLease = 600 * time.Millisecond
@@ -1066,26 +1063,6 @@ func TestAddIndexUniqueFailOnDuplicate(t *testing.T) {
10661063
ddl.ResultCounterForTest = nil
10671064
}
10681065

1069-
// withMockTiFlash sets the mockStore to have N TiFlash stores (naming as tiflash0, tiflash1, ...).
1070-
func withMockTiFlash(nodes int) mockstore.MockTiKVStoreOption {
1071-
return mockstore.WithMultipleOptions(
1072-
mockstore.WithClusterInspector(func(c testutils.Cluster) {
1073-
mockCluster := c.(*unistore.Cluster)
1074-
_, _, region1 := mockstore.BootstrapWithSingleStore(c)
1075-
tiflashIdx := 0
1076-
for tiflashIdx < nodes {
1077-
store2 := c.AllocID()
1078-
peer2 := c.AllocID()
1079-
addr2 := fmt.Sprintf("tiflash%d", tiflashIdx)
1080-
mockCluster.AddStore(store2, addr2, &metapb.StoreLabel{Key: "engine", Value: "tiflash"})
1081-
mockCluster.AddPeer(region1, store2, peer2)
1082-
tiflashIdx++
1083-
}
1084-
}),
1085-
mockstore.WithStoreType(mockstore.EmbedUnistore),
1086-
)
1087-
}
1088-
10891066
func getJobsBySQL(se sessiontypes.Session, tbl, condition string) ([]*model.Job, error) {
10901067
rs, err := se.Execute(context.Background(), fmt.Sprintf("select job_meta from mysql.%s %s", tbl, condition))
10911068
if err != nil {
@@ -1140,7 +1117,7 @@ func TestCreateTableWithVectorIndex(t *testing.T) {
11401117

11411118
// test TiFlash store count is 2
11421119
mockTiflashStoreCnt := uint64(2)
1143-
store, dom = testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(int(mockTiflashStoreCnt)), mockstore.WithDDLChecker())
1120+
store, dom = testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(int(mockTiflashStoreCnt)), mockstore.WithDDLChecker())
11441121
tk = testkit.NewTestKit(t, store)
11451122
tk.MustExec("use test")
11461123
checkCreateTableWithVectorIdx(1)
@@ -1160,7 +1137,7 @@ func TestCreateTableWithVectorIndex(t *testing.T) {
11601137
}
11611138

11621139
func TestAddVectorIndexSimple(t *testing.T) {
1163-
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(2))
1140+
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
11641141
tk := testkit.NewTestKit(t, store)
11651142
tk.MustExec("use test")
11661143
tk.MustExec("drop table if exists t, pt;")
@@ -1345,7 +1322,7 @@ func TestAddVectorIndexSimple(t *testing.T) {
13451322
}
13461323

13471324
func TestAddVectorIndexRollback(t *testing.T) {
1348-
store, _ := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(2))
1325+
store, _ := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
13491326
tk := testkit.NewTestKit(t, store)
13501327
tk.MustExec("use test")
13511328
tk.MustExec("drop table if exists t;")

‎pkg/executor/show.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1230,7 +1230,7 @@ func constructResultOfShowCreateTable(ctx sessionctx.Context, dbName *pmodel.CIS
12301230
cols = append(cols, colInfo)
12311231
}
12321232
if idxInfo.VectorInfo != nil {
1233-
funcName := model.DistanceMetricToFnName[idxInfo.VectorInfo.DistanceMetric]
1233+
funcName := model.IndexableDistanceMetricToFnName[idxInfo.VectorInfo.DistanceMetric]
12341234
fmt.Fprintf(buf, "((%s(%s)))", strings.ToUpper(funcName), strings.Join(cols, ","))
12351235
} else {
12361236
fmt.Fprintf(buf, "(%s)", strings.Join(cols, ","))

‎pkg/expression/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ go_library(
6969
"simple_rewriter.go",
7070
"util.go",
7171
"vectorized.go",
72+
"vs_helper.go",
7273
],
7374
importpath = "github.com/pingcap/tidb/pkg/expression",
7475
visibility = ["//visibility:public"],

0 commit comments

Comments
 (0)