Skip to content

Commit

Permalink
*: Small refine for Vector Search related utils (#56468)
Browse files Browse the repository at this point in the history
ref #54245
  • Loading branch information
breezewish authored Oct 8, 2024
1 parent 55fadff commit f31b091
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 16 deletions.
2 changes: 1 addition & 1 deletion pkg/ddl/index.go
Original file line number Diff line number Diff line change
Expand Up @@ -385,7 +385,7 @@ func buildVectorInfoWithCheck(indexPartSpecifications []*ast.IndexPartSpecificat
if !ok {
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs(fmt.Sprintf("unsupported function: %v", idxPart.Expr))
}
distanceMetric, ok := variable.DistanceMetric4VectorIndex[f.FnName.L]
distanceMetric, ok := model.FnNameToDistanceMetric[f.FnName.L]
if !ok {
return nil, "", dbterror.ErrUnsupportedAddVectorIndex.FastGenByArgs("unsupported function")
}
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 := variable.Function4VectorIndex[idxInfo.VectorInfo.DistanceMetric]
funcName := model.DistanceMetricToFnName[idxInfo.VectorInfo.DistanceMetric]
fmt.Fprintf(buf, "((%s(%s)))", strings.ToUpper(funcName), strings.Join(cols, ","))
} else {
fmt.Fprintf(buf, "(%s)", strings.Join(cols, ","))
Expand Down
20 changes: 19 additions & 1 deletion pkg/meta/model/index.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,14 @@
package model

import (
"github.com/pingcap/tidb/pkg/parser/ast"
"github.com/pingcap/tidb/pkg/parser/model"
"github.com/pingcap/tidb/pkg/parser/types"
)

// DistanceMetric is the distance metric used by the vector index.
// `DistanceMetric` is actually vector functions in ast package. Use `DistanceMetric` to avoid cycle dependency
// Note that not all distance functions are indexable.
// See FnNameToDistanceMetric for a list of indexable distance functions.
type DistanceMetric string

// Note: tipb.VectorDistanceMetric's enum names must be aligned with these constant values.
Expand All @@ -29,9 +31,25 @@ const (
// DistanceMetricCosine is cosine distance.
DistanceMetricCosine DistanceMetric = "COSINE"
// DistanceMetricInnerProduct is inner product.
// Currently this distance metric is not supported. It is placed here only for
// reminding what's the desired naming convension (UPPER_UNDER_SCORE) if this
// is going to be implemented.
DistanceMetricInnerProduct DistanceMetric = "INNER_PRODUCT"
)

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

// DistanceMetricToFnName maps a distance metric to the distance function name.
var DistanceMetricToFnName = 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.
Expand Down
13 changes: 0 additions & 13 deletions pkg/sessionctx/variable/varsutil.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ import (

"github.com/docker/go-units"
"github.com/pingcap/errors"
"github.com/pingcap/tidb/pkg/meta/model"
"github.com/pingcap/tidb/pkg/parser/ast"
"github.com/pingcap/tidb/pkg/parser/charset"
"github.com/pingcap/tidb/pkg/parser/mysql"
Expand Down Expand Up @@ -643,15 +642,3 @@ func parseSchemaCacheSize(s *SessionVars, normalizedValue string, originalValue

return 0, "", ErrTruncatedWrongValue.GenWithStackByArgs(TiDBSchemaCacheSize, originalValue)
}

// DistanceMetric4VectorIndex stores distance metrics for the vector index.
var DistanceMetric4VectorIndex = map[string]model.DistanceMetric{
ast.VecCosineDistance: model.DistanceMetricCosine,
ast.VecL2Distance: model.DistanceMetricL2,
}

// Function4VectorIndex stores functions for the vector index.
var Function4VectorIndex = map[model.DistanceMetric]string{
model.DistanceMetricCosine: ast.VecCosineDistance,
model.DistanceMetricL2: ast.VecL2Distance,
}

0 comments on commit f31b091

Please sign in to comment.