planner: support vector index by CBO

pkg/ddl/BUILD.bazel

@@ -320,7 +320,6 @@ go_test(
         "//pkg/store/gcworker",
         "//pkg/store/helper",
         "//pkg/store/mockstore",
-        "//pkg/store/mockstore/unistore",
         "//pkg/table",
         "//pkg/table/tables",
         "//pkg/tablecodec",
@@ -349,7 +348,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",

pkg/ddl/cancel_test.go

@@ -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"
@@ -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)
 

pkg/ddl/index_modify_test.go

@@ -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"
@@ -41,7 +40,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"
@@ -55,7 +53,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
@@ -1071,26 +1068,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 {
@@ -1118,7 +1095,7 @@ func getJobsBySQL(se sessiontypes.Session, tbl, condition string) ([]*model.Job,
 }
 
 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;")
@@ -1267,7 +1244,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;")

pkg/planner/core/BUILD.bazel

@@ -273,7 +273,6 @@ go_test(
     deps = [
         "//pkg/config",
         "//pkg/domain",
-        "//pkg/domain/infosync",
         "//pkg/expression",
         "//pkg/expression/aggregation",
         "//pkg/expression/context",
@@ -307,12 +306,12 @@ go_test(
         "//pkg/sessionctx/variable",
         "//pkg/sessiontxn",
         "//pkg/statistics",
+        "//pkg/store/mockstore",
         "//pkg/table",
         "//pkg/testkit",
         "//pkg/testkit/ddlhelper",
         "//pkg/testkit/external",
         "//pkg/testkit/testdata",
-        "//pkg/testkit/testfailpoint",
         "//pkg/testkit/testmain",
         "//pkg/testkit/testsetup",
         "//pkg/testkit/testutil",

pkg/planner/core/access_object.go

@@ -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
 }
 

pkg/planner/core/casetest/enforcempp/BUILD.bazel

@@ -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",

pkg/planner/core/casetest/enforcempp/enforce_mpp_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/pkg/meta/model"
 	pmodel "github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/planner/util/coretestsdk"
+	"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/testdata"
@@ -461,7 +462,7 @@ func TestMPPSingleDistinct3Stage(t *testing.T) {
 //
 //	since it doesn't change the schema out (index ref is still the right), so by now it's fine. SEE case: EXPLAIN select count(distinct a), count(distinct b), sum(c) from t.
 func TestMPPMultiDistinct3Stage(t *testing.T) {
-	store := testkit.CreateMockStore(t, coretestsdk.WithMockTiFlash(2))
+	store := testkit.CreateMockStore(t, mockstore.WithMockTiFlash(2))
 	tk := testkit.NewTestKit(t, store)
 
 	// test table
@@ -519,7 +520,7 @@ func TestMPPMultiDistinct3Stage(t *testing.T) {
 
 // Test null-aware semi join push down for MPP mode
 func TestMPPNullAwareSemiJoinPushDown(t *testing.T) {
-	store := testkit.CreateMockStore(t, coretestsdk.WithMockTiFlash(2))
+	store := testkit.CreateMockStore(t, mockstore.WithMockTiFlash(2))
 	tk := testkit.NewTestKit(t, store)
 
 	// test table
@@ -567,7 +568,7 @@ func TestMPPNullAwareSemiJoinPushDown(t *testing.T) {
 }
 
 func TestMPPSharedCTEScan(t *testing.T) {
-	store := testkit.CreateMockStore(t, coretestsdk.WithMockTiFlash(2))
+	store := testkit.CreateMockStore(t, mockstore.WithMockTiFlash(2))
 	tk := testkit.NewTestKit(t, store)
 
 	// test table
@@ -615,7 +616,7 @@ func TestMPPSharedCTEScan(t *testing.T) {
 }
 
 func TestRollupMPP(t *testing.T) {
-	store := testkit.CreateMockStore(t, coretestsdk.WithMockTiFlash(2))
+	store := testkit.CreateMockStore(t, mockstore.WithMockTiFlash(2))
 	tk := testkit.NewTestKit(t, store)
 
 	tk.MustExec("use test")

pkg/planner/core/casetest/hint/BUILD.bazel

@@ -17,6 +17,7 @@ go_test(
         "//pkg/parser/model",
         "//pkg/planner/util/coretestsdk",
         "//pkg/sessionctx/variable",
+        "//pkg/store/mockstore",
         "//pkg/testkit",
         "//pkg/testkit/testdata",
         "//pkg/testkit/testmain",

pkg/planner/core/casetest/hint/hint_test.go

@@ -23,6 +23,7 @@ import (
 	pmodel "github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/planner/util/coretestsdk"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/store/mockstore"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/testkit/testdata"
 	"github.com/stretchr/testify/require"
@@ -69,7 +70,7 @@ func TestReadFromStorageHint(t *testing.T) {
 }
 
 func TestAllViewHintType(t *testing.T) {
-	store := testkit.CreateMockStore(t, coretestsdk.WithMockTiFlash(2))
+	store := testkit.CreateMockStore(t, mockstore.WithMockTiFlash(2))
 	tk := testkit.NewTestKit(t, store)
 
 	tk.MustExec("use test")
@@ -130,7 +131,7 @@ func TestAllViewHintType(t *testing.T) {
 }
 
 func TestJoinHintCompatibility(t *testing.T) {
-	store := testkit.CreateMockStore(t, coretestsdk.WithMockTiFlash(2))
+	store := testkit.CreateMockStore(t, mockstore.WithMockTiFlash(2))
 	tk := testkit.NewTestKit(t, store)
 
 	tk.MustExec("use test")

pkg/planner/core/casetest/index/BUILD.bazel

@@ -9,13 +9,19 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 6,
+    shard_count = 8,
     deps = [
+        "//pkg/domain",
+        "//pkg/domain/infosync",
+        "//pkg/parser/model",
+        "//pkg/store/mockstore",
         "//pkg/testkit",
         "//pkg/testkit/testdata",
+        "//pkg/testkit/testfailpoint",
         "//pkg/testkit/testmain",
         "//pkg/testkit/testsetup",
         "//pkg/util",
+        "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",
     ],
 )

pkg/planner/core/casetest/index/index_test.go

@@ -15,12 +15,20 @@
 package index
 
 import (
+	"context"
 	"fmt"
 	"testing"
+	"time"
 
+	"github.com/pingcap/tidb/pkg/domain"
+	"github.com/pingcap/tidb/pkg/domain/infosync"
+	"github.com/pingcap/tidb/pkg/parser/model"
+	"github.com/pingcap/tidb/pkg/store/mockstore"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/testkit/testdata"
+	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
 	"github.com/pingcap/tidb/pkg/util"
+	"github.com/stretchr/testify/require"
 )
 
 func TestNullConditionForPrefixIndex(t *testing.T) {
@@ -221,3 +229,83 @@ func TestOrderedIndexWithIsNull(t *testing.T) {
 		"  └─IndexRangeScan_11 10.00 cop[tikv] table:t1, index:b(b, c) range:[NULL,NULL], keep order:true, stats:pseudo",
 	))
 }
+
+const tiflashReplicaLease = 600 * time.Millisecond
+
+func TestVectorIndex(t *testing.T) {
+	store := testkit.CreateMockStoreWithSchemaLease(t, tiflashReplicaLease, mockstore.WithMockTiFlash(2))
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+
+	tiflash := infosync.NewMockTiFlash()
+	infosync.SetMockTiFlash(tiflash)
+	defer func() {
+		tiflash.Lock()
+		tiflash.StatusServer.Close()
+		tiflash.Unlock()
+	}()
+
+	tk.MustExec("create table t (a int, b vector, c vector(3), d vector(4));")
+	tk.MustExec("alter table t set tiflash replica 1;")
+	tk.MustExec("alter table t add vector index vecIdx1((vec_cosine_distance(d))) USING HNSW;")
+	tk.MustUseIndex("select * from t use index(vecIdx1) order by vec_cosine_distance(d) limit 1", "vecIdx1")
+}
+
+func TestAnalyzeVectorIndex(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomainWithSchemaLease(t, 200*time.Millisecond, mockstore.WithMockTiFlash(2))
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t;")
+
+	tiflash := infosync.NewMockTiFlash()
+	infosync.SetMockTiFlash(tiflash)
+	defer func() {
+		tiflash.Lock()
+		tiflash.StatusServer.Close()
+		tiflash.Unlock()
+	}()
+	tk.MustExec(`create table t(a int, b vector(2), c vector(3), j json, index(a))`)
+	tk.MustExec("alter table t set tiflash replica 2 location labels 'a','b';")
+	tblInfo, err := dom.InfoSchema().TableByName(context.Background(), model.NewCIStr("test"), model.NewCIStr("t"))
+	require.NoError(t, err)
+	err = domain.GetDomain(tk.Session()).DDLExecutor().UpdateTableReplicaInfo(tk.Session(), tblInfo.Meta().ID, true)
+	require.NoError(t, err)
+
+	testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/ddl/MockCheckVectorIndexProcess", `return(0)`)
+	tk.MustExec("alter table t add vector index idx((VEC_COSINE_DISTANCE(b))) USING HNSW")
+	tk.MustExec("alter table t add vector index idx2((VEC_COSINE_DISTANCE(c))) USING HNSW")
+
+	tk.MustUseIndex("select * from t use index(idx) order by vec_cosine_distance(b, '[1, 0]') limit 1", "idx")
+	tk.MustUseIndex("select * from t order by vec_cosine_distance(b, '[1, 0]') limit 1", "idx")
+	tk.MustNoIndexUsed("select * from t ignore index(idx) order by vec_cosine_distance(b, '[1, 0]') limit 1")
+
+	tk.MustExec("set tidb_analyze_version=2")
+	tk.MustExec("analyze table t")
+	tk.MustQuery("show warnings").Sort().Check(testkit.Rows(
+		"Note 1105 Analyze use auto adjusted sample rate 1.000000 for table test.t, reason to use this rate is \"use min(1, 110000/10000) as the sample-rate=1\"",
+		"Warning 1105 No predicate column has been collected yet for table test.t, so only indexes and the columns composing the indexes will be analyzed",
+		"Warning 1105 analyzing vector index is not supported, skip idx",
+		"Warning 1105 analyzing vector index is not supported, skip idx2"))
+	tk.MustExec("analyze table t index idx")
+	tk.MustQuery("show warnings").Sort().Check(testkit.Rows(
+		"Note 1105 Analyze use auto adjusted sample rate 1.000000 for table test.t, reason to use this rate is \"TiDB assumes that the table is empty, use sample-rate=1\"",
+		"Warning 1105 No predicate column has been collected yet for table test.t, so only indexes and the columns composing the indexes will be analyzed",
+		"Warning 1105 The version 2 would collect all statistics not only the selected indexes",
+		"Warning 1105 analyzing vector index is not supported, skip idx",
+		"Warning 1105 analyzing vector index is not supported, skip idx2"))
+
+	tk.MustExec("set tidb_analyze_version=1")
+	tk.MustExec("analyze table t")
+	tk.MustQuery("show warnings").Sort().Check(testkit.Rows(
+		"Warning 1105 analyzing vector index is not supported, skip idx",
+		"Warning 1105 analyzing vector index is not supported, skip idx2"))
+	tk.MustExec("analyze table t index idx")
+	tk.MustQuery("show warnings").Sort().Check(testkit.Rows(
+		"Warning 1105 analyzing vector index is not supported, skip idx"))
+	tk.MustExec("analyze table t index a")
+	tk.MustQuery("show warnings").Sort().Check(testkit.Rows())
+	tk.MustExec("analyze table t index a, idx, idx2")
+	tk.MustQuery("show warnings").Sort().Check(testkit.Rows(
+		"Warning 1105 analyzing vector index is not supported, skip idx",
+		"Warning 1105 analyzing vector index is not supported, skip idx2"))
+}

pkg/planner/core/casetest/physicalplantest/BUILD.bazel

@@ -23,6 +23,7 @@ go_test(
         "//pkg/planner/core",
         "//pkg/planner/core/resolve",
         "//pkg/planner/util/coretestsdk",
+        "//pkg/store/mockstore",
         "//pkg/testkit",
         "//pkg/testkit/external",
         "//pkg/testkit/testdata",