Skip to content

planner: fix that vector index output empty result when pk is non-int type #57629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Nov 27, 2024
Merged

planner: fix that vector index output empty result when pk is non-int type #57629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
38e74df
planner: fix that vector index output empty result when pk is non-int…
winoros Nov 22, 2024
6c41ea2
fix tests
winoros Nov 22, 2024
ec3fb7b
planner/core: fix a wrong privilege check for CTE & UPDATE statement …
bb7133 Nov 22, 2024
2ba6fda
planner, runtime_filter: Remove redundant logs whose meaning can be d…
elsa0520 Nov 22, 2024
eca8aca
add unit test
winoros Nov 26, 2024
ccb08e7
fix static check
winoros Nov 26, 2024
d08472b
Merge remote-tracking branch 'origin/master' into vectorindex-nonintpk
winoros Nov 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion pkg/planner/core/casetest/vectorsearch/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,21 +9,25 @@ go_test(
],
data = glob(["testdata/**"]),
flaky = True,
shard_count = 4,
shard_count = 5,
deps = [
"//pkg/config",
"//pkg/domain",
"//pkg/domain/infosync",
"//pkg/meta/model",
"//pkg/parser/model",
"//pkg/planner",
"//pkg/planner/core",
"//pkg/planner/core/base",
"//pkg/planner/core/resolve",
"//pkg/session",
"//pkg/store/mockstore",
"//pkg/testkit",
"//pkg/testkit/testdata",
"//pkg/testkit/testfailpoint",
"//pkg/testkit/testmain",
"//pkg/testkit/testsetup",
"//pkg/types",
"//pkg/util/plancodec",
"@com_github_pingcap_tipb//go-tipb",
"@com_github_stretchr_testify//require",
Expand Down
68 changes: 68 additions & 0 deletions pkg/planner/core/casetest/vectorsearch/vector_index_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,16 @@ import (
"github.com/pingcap/tidb/pkg/domain/infosync"
"github.com/pingcap/tidb/pkg/meta/model"
pmodel "github.com/pingcap/tidb/pkg/parser/model"
"github.com/pingcap/tidb/pkg/planner"
"github.com/pingcap/tidb/pkg/planner/core"
"github.com/pingcap/tidb/pkg/planner/core/base"
"github.com/pingcap/tidb/pkg/planner/core/resolve"
"github.com/pingcap/tidb/pkg/session"
"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/types"
"github.com/pingcap/tidb/pkg/util/plancodec"
"github.com/pingcap/tipb/go-tipb"
"github.com/stretchr/testify/require"
Expand Down Expand Up @@ -244,3 +248,67 @@ func TestANNInexWithSimpleCBO(t *testing.T) {
testkit.SetTiFlashReplica(t, dom, "test", "t1")
tk.MustUseIndex("select * from t1 order by vec_cosine_distance(vec, '[1,1,1]') limit 1", "vector_index")
}

func TestANNIndexWithNonIntClusteredPk(t *testing.T) {
store := testkit.CreateMockStoreWithSchemaLease(t, 1*time.Second, mockstore.WithMockTiFlash(2))

tk := testkit.NewTestKit(t, store)

tiflash := infosync.NewMockTiFlash()
infosync.SetMockTiFlash(tiflash)
defer func() {
tiflash.Lock()
tiflash.StatusServer.Close()
tiflash.Unlock()
}()

testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/ddl/MockCheckVectorIndexProcess", `return(1)`)

tk.MustExec("use test")
tk.MustExec("drop table if exists t1")
tk.MustExec(`
create table t1 (
vec vector(3),
a int,
b int,
c vector(3),
d vector,
primary key (a, b)
)
`)
tk.MustExec("alter table t1 set tiflash replica 1;")
tk.MustExec("alter table t1 add vector index ((vec_cosine_distance(vec))) USING HNSW;")
tk.MustExec("insert into t1 values ('[1,1,1]', 1, 1, '[1,1,1]', '[1,1,1]')")
dom := domain.GetDomain(tk.Session())
testkit.SetTiFlashReplica(t, dom, "test", "t1")
sctx := tk.Session()
stmts, err := session.Parse(sctx, "select * from t1 use index(vector_index) order by vec_cosine_distance(vec, '[1,1,1]') limit 1")
require.NoError(t, err)
require.Len(t, stmts, 1)
stmt := stmts[0]
ret := &core.PreprocessorReturn{}
nodeW := resolve.NewNodeW(stmt)
err = core.Preprocess(context.Background(), sctx, nodeW, core.WithPreprocessorReturn(ret))
require.NoError(t, err)
var finalPlanTree base.Plan
finalPlanTree, _, err = planner.Optimize(context.Background(), sctx, nodeW, ret.InfoSchema)
require.NoError(t, err)
physicalTree, ok := finalPlanTree.(base.PhysicalPlan)
require.True(t, ok)
// Find the PhysicalTableReader node.
tableReader := physicalTree
for ; len(tableReader.Children()) > 0; tableReader = tableReader.Children()[0] {
}
castedTableReader, ok := tableReader.(*core.PhysicalTableReader)
require.True(t, ok)
tableScan, err := castedTableReader.GetTableScan()
require.NoError(t, err)
// Check that it has the extra vector index information.
require.NotNil(t, tableScan.AnnIndexExtra)
require.Len(t, tableScan.Ranges, 1)
// Check that it's full scan.
require.Equal(t, "[-inf,+inf]", tableScan.Ranges[0].String())
// Check that the -inf and +inf are the correct types.
require.Equal(t, types.KindMinNotNull, tableScan.Ranges[0].LowVal[0].Kind())
require.Equal(t, types.KindMaxValue, tableScan.Ranges[0].HighVal[0].Kind())
}
26 changes: 13 additions & 13 deletions pkg/planner/core/find_best_task.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -761,6 +761,19 @@ func compareCandidates(sctx base.PlanContext, prop *property.PhysicalProperty, l
}

func isMatchProp(ds *logicalop.DataSource, path *util.AccessPath, prop *property.PhysicalProperty) bool {
if prop.VectorProp.VectorHelper != nil && path.Index != nil && path.Index.VectorInfo != nil {
if path.Index == nil || path.Index.VectorInfo == nil {
return false
}
if ds.TableInfo.Columns[path.Index.Columns[0].Offset].ID != prop.VectorProp.Column.ID {
return false
}

if model.IndexableFnNameToDistanceMetric[prop.VectorProp.DistanceFnName.L] != path.Index.VectorInfo.DistanceMetric {
return false
}
return true
}
Comment on lines +764 to +776
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I changed it to the highest priority.
Because when for the vector index path, the IsIntHandlePath can also be true. We need to use a new way to check the real pk path.

var isMatchProp bool
if path.IsIntHandlePath {
pkCol := ds.GetPKIsHandleCol()
Expand Down Expand Up @@ -808,19 +821,6 @@ func isMatchProp(ds *logicalop.DataSource, path *util.AccessPath, prop *property
}
}
}
if prop.VectorProp.VectorHelper != nil && path.Index.VectorInfo != nil {
if path.Index == nil || path.Index.VectorInfo == nil {
return false
}
if ds.TableInfo.Columns[path.Index.Columns[0].Offset].ID != prop.VectorProp.Column.ID {
return false
}

if model.IndexableFnNameToDistanceMetric[prop.VectorProp.DistanceFnName.L] != path.Index.VectorInfo.DistanceMetric {
return false
}
return true
}
return isMatchProp
}

Expand Down
6 changes: 5 additions & 1 deletion pkg/planner/core/planbuilder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1194,15 +1194,19 @@ func getPossibleAccessPaths(ctx base.PlanContext, tableHints *hint.PlanHints, in
continue
}
}
path := &util.AccessPath{Index: index}
if index.VectorInfo != nil {
// Because the value of `TiFlashReplica.Available` changes as the user modify replica, it is not ideal if the state of index changes accordingly.
// So the current way to use the vector indexes is to require the TiFlash Replica to be available.
if !tblInfo.TiFlashReplica.Available {
continue
}
path := genTiFlashPath(tblInfo)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel that the original logic is different here set up "tablePath. IsCommonHandlePath = true". What does this setting do? Why is the original logic not needed？

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A normal index will create the index's ranges.
But the vector index still uses the table range. So we need to correctly set the pk type. So it can generate the correct table range.
For int pk, it will use [MinInt64, MaxInt64] as full range. But for other cases, [MinNotNull, MaxValue] will be full range.

Copy link
Contributor

@AilinKid AilinKid Nov 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image

i think if we don't set it here, we will step into here right? which will derive a int range for common handle tiflash path

path.StoreType = kv.TiFlash
path.Index = index
publicPaths = append(publicPaths, path)
continue
}
path := &util.AccessPath{Index: index}
publicPaths = append(publicPaths, path)
}
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/planner/util/path.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,7 @@ func (path *AccessPath) IsTiKVTablePath() bool {

// IsTiFlashSimpleTablePath returns true if it's a TiFlash path and will not use any special indexes like vector index.
func (path *AccessPath) IsTiFlashSimpleTablePath() bool {
return (path.IsIntHandlePath || path.IsCommonHandlePath) && path.StoreType == kv.TiFlash
return path.StoreType == kv.TiFlash && path.Index == nil
}

// SplitCorColAccessCondFromFilters move the necessary filter in the form of index_col = corrlated_col to access conditions.
Expand Down