enable to use l2 index for queries using l2sq

Author: aunjgr

pkg/sql/plan/apply_indices_hnsw.go

@@ -72,7 +72,8 @@ func (builder *QueryBuilder) applyIndicesForSortUsingHnsw(nodeID int32, projNode
 		return nodeID, nil
 	}
 
-	if opType != metric.DistFuncOpTypes[distFnExpr.Func.ObjName] {
+	origFuncName := distFnExpr.Func.ObjName
+	if opType != metric.DistFuncOpTypes[origFuncName] {
 		return nodeID, nil
 	}
 
@@ -93,12 +94,13 @@ func (builder *QueryBuilder) applyIndicesForSortUsingHnsw(nodeID int32, projNode
 	}
 
 	// generate JSON by fmt.Sprintf instead of sonic.Marshal for performance
-	tblCfgStr := fmt.Sprintf(`{"db": "%s", "src": "%s", "metadata":"%s", "index":"%s", "threads_search": %d}`,
+	tblCfgStr := fmt.Sprintf(`{"db": "%s", "src": "%s", "metadata":"%s", "index":"%s", "threads_search": %d, "orig_func_name": "%s"}`,
 		scanNode.ObjRef.SchemaName,
 		scanNode.TableDef.Name,
 		metaDef.IndexTableName,
 		idxDef.IndexTableName,
-		nThread.(int64))
+		nThread.(int64),
+		origFuncName)
 
 	// JOIN between source table and hnsw_search table function
 	tableFuncTag := builder.genNewTag()

pkg/sql/plan/apply_indices_ivfflat.go

@@ -73,7 +73,9 @@ func (builder *QueryBuilder) applyIndicesForSortUsingIvfflat(nodeID int32, projN
 		return nodeID, nil
 	}
 
-	if opType != metric.DistFuncOpTypes[distFnExpr.Func.ObjName] {
+	origFuncName := distFnExpr.Func.ObjName
+	if opType != metric.DistFuncOpTypes[origFuncName] {
+		fmt.Println("IVFFlat: Unsupported distance function")
 		return nodeID, nil
 	}
 
@@ -108,7 +110,7 @@ func (builder *QueryBuilder) applyIndicesForSortUsingIvfflat(nodeID int32, projN
 	params := idxDef.IndexAlgoParams
 
 	tblCfgStr := fmt.Sprintf(`{"db": "%s", "src": "%s", "metadata":"%s", "index":"%s", "threads_search": %d,
-			"entries": "%s", "nprobe" : %d, "pktype" : %d, "pkey" : "%s", "part" : "%s", "parttype" : %d}`,
+			"entries": "%s", "nprobe" : %d, "pktype" : %d, "pkey" : "%s", "part" : "%s", "parttype" : %d, "orig_func_name": "%s"}`,
 		scanNode.ObjRef.SchemaName,
 		scanNode.TableDef.Name,
 		metaDef.IndexTableName,
@@ -119,7 +121,8 @@ func (builder *QueryBuilder) applyIndicesForSortUsingIvfflat(nodeID int32, projN
 		pkType.Id,
 		scanNode.TableDef.Pkey.PkeyColName,
 		keyPart,
-		partType.Id)
+		partType.Id,
+		origFuncName)
 
 	// JOIN between source table and hnsw_search table function
 	tableFuncTag := builder.genNewTag()

pkg/vectorindex/hnsw/search.go

@@ -133,7 +133,7 @@ func (s *HnswSearch[T]) Search(sqlproc *sqlexec.SqlProcess, anyquery any, rt vec
 			return nil, nil, moerr.NewInternalError(sqlproc.GetContext(), "heap return key is not int64")
 		}
 		reskeys = append(reskeys, sr.Id)
-		sr.Distance = metric.DistanceTransformHnsw(sr.Distance, s.Idxcfg.OpType, s.Idxcfg.Usearch.Metric)
+		sr.Distance = metric.DistanceTransformHnsw(sr.Distance, metric.DistFuncNameToMetricType[s.Tblcfg.OrigFuncName], s.Idxcfg.Usearch.Metric)
 		resdistances = append(resdistances, sr.Distance)
 	}
 
@@ -247,6 +247,6 @@ func (s *HnswSearch[T]) Load(sqlproc *sqlexec.SqlProcess) error {
 
 // check config and update some parameters such as ef_search
 func (s *HnswSearch[T]) UpdateConfig(newalgo cache.VectorIndexSearchIf) error {
-
+	s.Tblcfg.OrigFuncName = newalgo.(*HnswSearch[T]).Tblcfg.OrigFuncName
 	return nil
 }

pkg/vectorindex/ivfflat/search.go

@@ -207,7 +207,7 @@ func (idx *IvfflatSearchIndex[T]) Search(
 			resid = append(resid, pk)
 
 			dist := vector.GetFixedAtNoTypeCheck[float64](bat.Vecs[1], i)
-			dist = metric.DistanceTransformIvfflat(dist, idxcfg.OpType, metric.MetricType(idxcfg.Ivfflat.Metric))
+			dist = metric.DistanceTransformIvfflat(dist, metric.DistFuncNameToMetricType[tblcfg.OrigFuncName], metric.MetricType(idxcfg.Ivfflat.Metric))
 			distances = append(distances, dist)
 		}
 	}
@@ -269,6 +269,6 @@ func (s *IvfflatSearch[T]) Load(sqlproc *sqlexec.SqlProcess) error {
 
 // check config and update some parameters such as ef_search
 func (s *IvfflatSearch[T]) UpdateConfig(newalgo cache.VectorIndexSearchIf) error {
-
+	s.Tblcfg.OrigFuncName = newalgo.(*IvfflatSearch[T]).Tblcfg.OrigFuncName
 	return nil
 }

pkg/vectorindex/metric/types.go

@@ -55,7 +55,7 @@ const (
 var (
 	DistFuncOpTypes = map[string]string{
 		DistFn_L2Distance:     OpType_L2Distance,
-		DistFn_L2sqDistance:   OpType_L2sqDistance,
+		DistFn_L2sqDistance:   OpType_L2Distance,
 		DistFn_InnerProduct:   OpType_InnerProduct,
 		DistFn_CosineDistance: OpType_CosineDistance,
 	}
@@ -120,17 +120,17 @@ func MaxFloat[T types.RealNumbers]() T {
 	}
 }
 
-func DistanceTransformHnsw(dist float64, optype string, metric usearch.Metric) float64 {
-	if optype == OpType_L2Distance && metric == usearch.L2sq {
-		// metric is l2sq but optype is l2_distance
+func DistanceTransformHnsw(dist float64, origMetricType MetricType, metricType usearch.Metric) float64 {
+	if origMetricType == Metric_L2Distance && metricType == usearch.L2sq {
+		// metric is l2sq but origin is l2_distance
 		return math.Sqrt(dist)
 	}
 	return dist
 }
 
-func DistanceTransformIvfflat(dist float64, optype string, metric MetricType) float64 {
-	if optype == OpType_L2Distance && metric == Metric_L2sqDistance {
-		// metric is l2sq but optype is l2_distance
+func DistanceTransformIvfflat(dist float64, origMetricType, metricType MetricType) float64 {
+	if origMetricType == Metric_L2Distance && metricType == Metric_L2sqDistance {
+		// metric is l2sq but origin is l2_distance
 		return math.Sqrt(dist)
 	}
 	return dist

pkg/vectorindex/types.go

@@ -54,6 +54,7 @@ type IndexTableConfig struct {
 	IndexTable    string `json:"index"`
 	PKey          string `json:"pkey"`
 	KeyPart       string `json:"part"`
+	OrigFuncName  string `json:"orig_func_name"`
 	ThreadsBuild  int64  `json:"threads_build"`
 	ThreadsSearch int64  `json:"threads_search"`
 	IndexCapacity int64  `json:"index_capacity"`