Skip to content

Commit

Permalink
planner: support vector index in planner (#56204)
Browse files Browse the repository at this point in the history
ref #54245, close #56510
  • Loading branch information
winoros authored Oct 10, 2024
1 parent 5e8bb8e commit 7aefb08
Show file tree
Hide file tree
Showing 46 changed files with 1,414 additions and 157 deletions.
12 changes: 6 additions & 6 deletions DEPS.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -5815,13 +5815,13 @@ def go_deps():
name = "com_github_pingcap_tipb",
build_file_proto_mode = "disable_global",
importpath = "github.com/pingcap/tipb",
sha256 = "119f73cede33a7d3712a9fedeb80f806228ecf1ec5258112b2b62af88731a304",
strip_prefix = "github.com/pingcap/tipb@v0.0.0-20240919023442-cf70966bef25",
sha256 = "b39e154272ba36d145c6049947a012a76be740b32a44a46d7253caa145c56cc9",
strip_prefix = "github.com/pingcap/tipb@v0.0.0-20241008083645-0bcddae67837",
urls = [
"http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
"http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
"https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
"https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240919023442-cf70966bef25.zip",
"http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
"http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
"https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
"https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241008083645-0bcddae67837.zip",
],
)
go_repository(
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ require (
github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5
github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e
github.com/pingcap/tipb v0.0.0-20240919023442-cf70966bef25
github.com/pingcap/tipb v0.0.0-20241008083645-0bcddae67837
github.com/prometheus/client_golang v1.20.4
github.com/prometheus/client_model v0.6.1
github.com/prometheus/common v0.57.0
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -676,8 +676,8 @@ github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d h1:y3EueKVfVykdpTyfU
github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d/go.mod h1:ORfBOFp1eteu2odzsyaxI+b8TzJwgjwyQcGhI+9SfEA=
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeAhmOQHNPRRogMYxP13V7BkYw3ZsoSfE=
github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530=
github.com/pingcap/tipb v0.0.0-20240919023442-cf70966bef25 h1:qH261jtqDvqWhhsnle+eURWV+PmFdqEGP/aHDqPFJAs=
github.com/pingcap/tipb v0.0.0-20240919023442-cf70966bef25/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
github.com/pingcap/tipb v0.0.0-20241008083645-0bcddae67837 h1:tyIymn821fB8gUmqafdvLlcFkVOpgyJXImoYJ8n9oJE=
github.com/pingcap/tipb v0.0.0-20241008083645-0bcddae67837/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
Expand Down
2 changes: 0 additions & 2 deletions pkg/ddl/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -325,7 +325,6 @@ go_test(
"//pkg/store/gcworker",
"//pkg/store/helper",
"//pkg/store/mockstore",
"//pkg/store/mockstore/unistore",
"//pkg/table",
"//pkg/table/tables",
"//pkg/tablecodec",
Expand Down Expand Up @@ -353,7 +352,6 @@ go_test(
"@com_github_ngaut_pools//:pools",
"@com_github_pingcap_errors//:errors",
"@com_github_pingcap_failpoint//:failpoint",
"@com_github_pingcap_kvproto//pkg/metapb",
"@com_github_stretchr_testify//assert",
"@com_github_stretchr_testify//require",
"@com_github_tikv_client_go_v2//oracle",
Expand Down
3 changes: 2 additions & 1 deletion pkg/ddl/cancel_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"github.com/pingcap/tidb/pkg/domain/infosync"
"github.com/pingcap/tidb/pkg/errno"
"github.com/pingcap/tidb/pkg/meta/model"
"github.com/pingcap/tidb/pkg/store/mockstore"
"github.com/pingcap/tidb/pkg/testkit"
"github.com/pingcap/tidb/pkg/testkit/external"
"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
Expand Down Expand Up @@ -223,7 +224,7 @@ func TestCancelVariousJobs(t *testing.T) {
return enterCnt.Load() == exitCnt.Load()
}, 10*time.Second, 10*time.Millisecond)
}
store := testkit.CreateMockStoreWithSchemaLease(t, 100*time.Millisecond, withMockTiFlash(2))
store := testkit.CreateMockStoreWithSchemaLease(t, 100*time.Millisecond, mockstore.WithMockTiFlash(2))
tk := testkit.NewTestKit(t, store)
tkCancel := testkit.NewTestKit(t, store)

Expand Down
3 changes: 2 additions & 1 deletion pkg/ddl/db_change_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ import (
sessiontypes "github.com/pingcap/tidb/pkg/session/types"
"github.com/pingcap/tidb/pkg/sessionctx"
"github.com/pingcap/tidb/pkg/sessiontxn"
"github.com/pingcap/tidb/pkg/store/mockstore"
"github.com/pingcap/tidb/pkg/testkit"
"github.com/pingcap/tidb/pkg/testkit/external"
"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
Expand Down Expand Up @@ -1159,7 +1160,7 @@ func TestParallelAlterAddIndex(t *testing.T) {
}

func TestParallelAlterAddVectorIndex(t *testing.T) {
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(2))
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
tk := testkit.NewTestKit(t, store)
tk.MustExec("create database test_db_state default charset utf8 default collate utf8_bin")
tk.MustExec("use test_db_state")
Expand Down
4 changes: 2 additions & 2 deletions pkg/ddl/index.go
Original file line number Diff line number Diff line change
Expand Up @@ -385,9 +385,9 @@ func buildVectorInfoWithCheck(indexPartSpecifications []*ast.IndexPartSpecificat
if !ok {
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs(fmt.Sprintf("unsupported function: %v", idxPart.Expr))
}
distanceMetric, ok := model.FnNameToDistanceMetric[f.FnName.L]
distanceMetric, ok := model.IndexableFnNameToDistanceMetric[f.FnName.L]
if !ok {
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs("unsupported function")
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs("currently only L2 and Cosine distance is indexable")
}
colExpr, ok := f.Args[0].(*ast.ColumnNameExpr)
if !ok {
Expand Down
29 changes: 3 additions & 26 deletions pkg/ddl/index_modify_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@ import (
"time"

"github.com/pingcap/errors"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/pingcap/tidb/pkg/ddl"
testddlutil "github.com/pingcap/tidb/pkg/ddl/testutil"
"github.com/pingcap/tidb/pkg/domain/infosync"
Expand All @@ -42,7 +41,6 @@ import (
"github.com/pingcap/tidb/pkg/sessionctx/variable"
"github.com/pingcap/tidb/pkg/sessiontxn"
"github.com/pingcap/tidb/pkg/store/mockstore"
"github.com/pingcap/tidb/pkg/store/mockstore/unistore"
"github.com/pingcap/tidb/pkg/table"
"github.com/pingcap/tidb/pkg/table/tables"
"github.com/pingcap/tidb/pkg/tablecodec"
Expand All @@ -56,7 +54,6 @@ import (
"github.com/pingcap/tidb/pkg/util/dbterror"
"github.com/pingcap/tidb/pkg/util/sqlexec"
"github.com/stretchr/testify/require"
"github.com/tikv/client-go/v2/testutils"
)

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

// withMockTiFlash sets the mockStore to have N TiFlash stores (naming as tiflash0, tiflash1, ...).
func withMockTiFlash(nodes int) mockstore.MockTiKVStoreOption {
return mockstore.WithMultipleOptions(
mockstore.WithClusterInspector(func(c testutils.Cluster) {
mockCluster := c.(*unistore.Cluster)
_, _, region1 := mockstore.BootstrapWithSingleStore(c)
tiflashIdx := 0
for tiflashIdx < nodes {
store2 := c.AllocID()
peer2 := c.AllocID()
addr2 := fmt.Sprintf("tiflash%d", tiflashIdx)
mockCluster.AddStore(store2, addr2, &metapb.StoreLabel{Key: "engine", Value: "tiflash"})
mockCluster.AddPeer(region1, store2, peer2)
tiflashIdx++
}
}),
mockstore.WithStoreType(mockstore.EmbedUnistore),
)
}

func getJobsBySQL(se sessiontypes.Session, tbl, condition string) ([]*model.Job, error) {
rs, err := se.Execute(context.Background(), fmt.Sprintf("select job_meta from mysql.%s %s", tbl, condition))
if err != nil {
Expand Down Expand Up @@ -1140,7 +1117,7 @@ func TestCreateTableWithVectorIndex(t *testing.T) {

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

func TestAddVectorIndexSimple(t *testing.T) {
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(2))
store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
tk := testkit.NewTestKit(t, store)
tk.MustExec("use test")
tk.MustExec("drop table if exists t, pt;")
Expand Down Expand Up @@ -1345,7 +1322,7 @@ func TestAddVectorIndexSimple(t *testing.T) {
}

func TestAddVectorIndexRollback(t *testing.T) {
store, _ := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, withMockTiFlash(2))
store, _ := testkit.CreateMockStoreAndDomainWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
tk := testkit.NewTestKit(t, store)
tk.MustExec("use test")
tk.MustExec("drop table if exists t;")
Expand Down
2 changes: 1 addition & 1 deletion pkg/executor/show.go
Original file line number Diff line number Diff line change
Expand Up @@ -1230,7 +1230,7 @@ func constructResultOfShowCreateTable(ctx sessionctx.Context, dbName *pmodel.CIS
cols = append(cols, colInfo)
}
if idxInfo.VectorInfo != nil {
funcName := model.DistanceMetricToFnName[idxInfo.VectorInfo.DistanceMetric]
funcName := model.IndexableDistanceMetricToFnName[idxInfo.VectorInfo.DistanceMetric]
fmt.Fprintf(buf, "((%s(%s)))", strings.ToUpper(funcName), strings.Join(cols, ","))
} else {
fmt.Fprintf(buf, "(%s)", strings.Join(cols, ","))
Expand Down
1 change: 1 addition & 0 deletions pkg/expression/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ go_library(
"simple_rewriter.go",
"util.go",
"vectorized.go",
"vs_helper.go",
],
importpath = "github.com/pingcap/tidb/pkg/expression",
visibility = ["//visibility:public"],
Expand Down
32 changes: 16 additions & 16 deletions pkg/expression/integration_test/integration_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -327,24 +327,24 @@ func TestVectorConstantExplain(t *testing.T) {
tk := testkit.NewTestKit(t, store)
tk.MustExec("use test")
tk.MustExec("CREATE TABLE t(c VECTOR);")
tk.MustQuery(`EXPLAIN SELECT VEC_COSINE_DISTANCE(c, '[1,2,3,4,5,6,7,8,9,10,11]') FROM t;`).Check(testkit.Rows(
"Projection_3 10000.00 root vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#3",
"└─TableReader_5 10000.00 root data:TableFullScan_4",
" └─TableFullScan_4 10000.00 cop[tikv] table:t keep order:false, stats:pseudo",
tk.MustQuery(`EXPLAIN format='brief' SELECT VEC_COSINE_DISTANCE(c, '[1,2,3,4,5,6,7,8,9,10,11]') FROM t;`).Check(testkit.Rows(
"Projection 10000.00 root vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#3",
"└─TableReader 10000.00 root data:TableFullScan",
" └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo",
))
tk.MustQuery(`EXPLAIN SELECT VEC_COSINE_DISTANCE(c, VEC_FROM_TEXT('[1,2,3,4,5,6,7,8,9,10,11]')) FROM t;`).Check(testkit.Rows(
"Projection_3 10000.00 root vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#3",
"└─TableReader_5 10000.00 root data:TableFullScan_4",
" └─TableFullScan_4 10000.00 cop[tikv] table:t keep order:false, stats:pseudo",
tk.MustQuery(`EXPLAIN format='brief' SELECT VEC_COSINE_DISTANCE(c, VEC_FROM_TEXT('[1,2,3,4,5,6,7,8,9,10,11]')) FROM t;`).Check(testkit.Rows(
"Projection 10000.00 root vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#3",
"└─TableReader 10000.00 root data:TableFullScan",
" └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo",
))
tk.MustQuery(`EXPLAIN SELECT VEC_COSINE_DISTANCE(c, '[1,2,3,4,5,6,7,8,9,10,11]') AS d FROM t ORDER BY d LIMIT 10;`).Check(testkit.Rows(
"Projection_6 10.00 root vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#3",
"└─Projection_13 10.00 root test.t.c",
" └─TopN_7 10.00 root Column#4, offset:0, count:10",
" └─Projection_14 10.00 root test.t.c, vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#4",
" └─TableReader_12 10.00 root data:TopN_11",
" └─TopN_11 10.00 cop[tikv] vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...]), offset:0, count:10",
" └─TableFullScan_10 10000.00 cop[tikv] table:t keep order:false, stats:pseudo",
tk.MustQuery(`EXPLAIN format = 'brief' SELECT VEC_COSINE_DISTANCE(c, '[1,2,3,4,5,6,7,8,9,10,11]') AS d FROM t ORDER BY d LIMIT 10;`).Check(testkit.Rows(
"Projection 10.00 root vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#3",
"└─Projection 10.00 root test.t.c",
" └─TopN 10.00 root Column#4, offset:0, count:10",
" └─Projection 10.00 root test.t.c, vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...])->Column#4",
" └─TableReader 10.00 root data:TopN",
" └─TopN 10.00 cop[tikv] vec_cosine_distance(test.t.c, [1,2,3,4,5,(6 more)...]), offset:0, count:10",
" └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo",
))

// Prepare a large Vector string
Expand Down
91 changes: 91 additions & 0 deletions pkg/expression/vs_helper.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
// Copyright 2024 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package expression

import (
"strings"

"github.com/pingcap/tidb/pkg/parser/ast"
"github.com/pingcap/tidb/pkg/parser/model"
"github.com/pingcap/tidb/pkg/parser/mysql"
"github.com/pingcap/tidb/pkg/types"
"github.com/pingcap/tidb/pkg/util/intest"
"github.com/pingcap/tipb/go-tipb"
)

var (
vsDistanceFnNamesLower = map[string]struct{}{
strings.ToLower(ast.VecL1Distance): {},
strings.ToLower(ast.VecL2Distance): {},
strings.ToLower(ast.VecCosineDistance): {},
strings.ToLower(ast.VecNegativeInnerProduct): {},
}
)

// VectorHelper is a helper struct for vector indexes.
type VectorHelper struct {
DistanceFnName model.CIStr
FnPbCode tipb.ScalarFuncSig
Vec types.VectorFloat32
Column *Column
}

// ExtractVectorHelper extracts a VectorSearchExpr from an expression.
// NOTE: not all VectorSearch functions are supported by the index. The caller
// needs to check the distance function name.
func ExtractVectorHelper(expr Expression) *VectorHelper {
x, ok := expr.(*ScalarFunction)
if !ok {
return nil
}

if _, isVecFn := vsDistanceFnNamesLower[x.FuncName.L]; !isVecFn {
return nil
}

args := x.GetArgs()
// One arg must be a vector column ref, and one arg must be a vector constant.
var vectorConstant *Constant = nil
var vectorColumn *Column = nil
nVectorColumns := 0
nVectorConstants := 0
for _, arg := range args {
if v, ok := arg.(*Column); ok {
if v.RetType.GetType() != mysql.TypeTiDBVectorFloat32 {
return nil
}
vectorColumn = v
nVectorColumns++
} else if v, ok := arg.(*Constant); ok {
if v.RetType.GetType() != mysql.TypeTiDBVectorFloat32 {
return nil
}
vectorConstant = v
nVectorConstants++
}
}
if nVectorColumns != 1 || nVectorConstants != 1 {
return nil
}

intest.Assert(vectorConstant.Value.Kind() == types.KindVectorFloat32, "internal: expect vectorFloat32 constant, but got %s", vectorConstant.Value.String())

return &VectorHelper{
DistanceFnName: x.FuncName,
FnPbCode: x.Function.PbCode(),
Vec: vectorConstant.Value.GetVectorFloat32(),
Column: vectorColumn,
}
}
10 changes: 5 additions & 5 deletions pkg/meta/model/index.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,23 +37,23 @@ const (
DistanceMetricInnerProduct DistanceMetric = "INNER_PRODUCT"
)

// FnNameToDistanceMetric maps a distance function name to the distance metric.
// IndexableFnNameToDistanceMetric maps a distance function name to the distance metric.
// Only indexable distance functions should be listed here!
var FnNameToDistanceMetric = map[string]DistanceMetric{
var IndexableFnNameToDistanceMetric = map[string]DistanceMetric{
ast.VecCosineDistance: DistanceMetricCosine,
ast.VecL2Distance: DistanceMetricL2,
}

// DistanceMetricToFnName maps a distance metric to the distance function name.
var DistanceMetricToFnName = map[DistanceMetric]string{
// IndexableDistanceMetricToFnName maps a distance metric to the distance function name.
var IndexableDistanceMetricToFnName = map[DistanceMetric]string{
DistanceMetricCosine: ast.VecCosineDistance,
DistanceMetricL2: ast.VecL2Distance,
}

// VectorIndexInfo is the information of vector index of a column.
type VectorIndexInfo struct {
// Dimension is the dimension of the vector.
Dimension uint64 `json:"dimension"` // Set to 0 when initially parsed from comment. Will be assigned to flen later.
Dimension uint64 `json:"dimension"`
// DistanceMetric is the distance metric used by the index.
DistanceMetric DistanceMetric `json:"distance_metric"`
}
Expand Down
1 change: 1 addition & 0 deletions pkg/planner/core/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,7 @@ go_test(
"//pkg/sessionctx/variable",
"//pkg/sessiontxn",
"//pkg/statistics",
"//pkg/store/mockstore",
"//pkg/table",
"//pkg/testkit",
"//pkg/testkit/ddlhelper",
Expand Down
13 changes: 13 additions & 0 deletions pkg/planner/core/access_object.go
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,19 @@ func (p *PhysicalTableScan) AccessObject() base.AccessObject {
res.Partitions = []string{partitionName}
}
}
if p.AnnIndexExtra != nil {
index := IndexAccess{
Name: p.AnnIndexExtra.IndexInfo.Name.O,
}
for _, idxCol := range p.AnnIndexExtra.IndexInfo.Columns {
if tblCol := p.Table.Columns[idxCol.Offset]; tblCol.Hidden {
index.Cols = append(index.Cols, tblCol.GeneratedExprString)
} else {
index.Cols = append(index.Cols, idxCol.Name.O)
}
}
res.Indexes = []IndexAccess{index}
}
return res
}

Expand Down
1 change: 1 addition & 0 deletions pkg/planner/core/casetest/enforcempp/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ go_test(
"//pkg/meta/model",
"//pkg/parser/model",
"//pkg/planner/util/coretestsdk",
"//pkg/store/mockstore",
"//pkg/testkit",
"//pkg/testkit/external",
"//pkg/testkit/testdata",
Expand Down
Loading

0 comments on commit 7aefb08

Please sign in to comment.