*: Small refine for Vector Search related utils (#56468)

ref pingcap/tidb#54245
This commit is contained in:
Wenxuan
2024-10-08 21:19:45 +08:00
committed by GitHub
parent 55fadff27a
commit f31b09149d
4 changed files with 21 additions and 16 deletions

View File

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

View File

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

View File

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

View File

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