From f31b09149d2f962d61cd27d0da64341d87f478da Mon Sep 17 00:00:00 2001
From: Wenxuan <breezewish@outlook.com>
Date: Tue, 8 Oct 2024 21:19:45 +0800
Subject: [PATCH] *: Small refine for Vector Search related utils (#56468)

ref pingcap/tidb#54245
---
 pkg/ddl/index.go                    |  2 +-
 pkg/executor/show.go                |  2 +-
 pkg/meta/model/index.go             | 20 +++++++++++++++++++-
 pkg/sessionctx/variable/varsutil.go | 13 -------------
 4 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go
index f62906d334..a13178994b 100644
--- a/pkg/ddl/index.go
+++ b/pkg/ddl/index.go
@@ -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")
 	}
diff --git a/pkg/executor/show.go b/pkg/executor/show.go
index 2717338a88..ed5f9f1f29 100644
--- a/pkg/executor/show.go
+++ b/pkg/executor/show.go
@@ -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, ","))
diff --git a/pkg/meta/model/index.go b/pkg/meta/model/index.go
index 7e21e852a9..f6af916a62 100644
--- a/pkg/meta/model/index.go
+++ b/pkg/meta/model/index.go
@@ -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.
diff --git a/pkg/sessionctx/variable/varsutil.go b/pkg/sessionctx/variable/varsutil.go
index dd19c3a08d..17146c1a95 100644
--- a/pkg/sessionctx/variable/varsutil.go
+++ b/pkg/sessionctx/variable/varsutil.go
@@ -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,
-}