use universal heap in Reader for IVFFlat

Author: aunjgr

pkg/objectio/types.go

@@ -64,12 +64,35 @@ func (f BlockReadFilter) DecideSearchFunc(isSortedBlk bool) ReadFilterSearchFunc
 	return nil
 }
 
+type Float64Heap []float64
+
+func (h Float64Heap) Len() int           { return len(h) }
+func (h Float64Heap) Less(i, j int) bool { return h[i] > h[j] }
+func (h Float64Heap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *Float64Heap) Push(x any) {
+	*h = append(*h, x.(float64))
+}
+
+func (h *Float64Heap) Pop() any {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[0 : n-1]
+	return x
+}
+
 type BlockReadTopOp struct {
 	Typ    types.T
 	Metric metric.MetricType
 	ColPos int32
 	NumVec []byte
 	Limit  uint64
+
+	//LowerBound float64
+	//UpperBound float64
+
+	DistHeap Float64Heap
 }
 
 type WriteOptions struct {

pkg/sql/compile/operator.go

@@ -1024,7 +1024,7 @@ func constructTableFunction(n *plan.Node, qry *plan.Query) *table_function.Table
 	arg.FuncName = n.TableDef.TblFunc.Name
 	arg.Params = n.TableDef.TblFunc.Param
 	arg.IsSingle = n.TableDef.TblFunc.IsSingle
-	arg.Limit = n.Limit
+	arg.Limit = n.BlockLimit
 	// probe side runtime filter specs
 	arg.RuntimeFilterSpecs = n.RuntimeFilterProbeList
 	return arg

pkg/sql/plan/apply_indices_ivfflat.go

@@ -182,7 +182,7 @@ func (builder *QueryBuilder) applyIndicesForSortUsingIvfflat(nodeID int32, projN
 			overFetchFactor := calculatePostFilterOverFetchFactor(originalLimit)
 
 			newLimit := max(uint64(float64(originalLimit)*overFetchFactor), originalLimit+10)
-			tableFuncNode.Limit = &Expr{
+			tableFuncNode.BlockLimit = &Expr{
 				Typ: limit.Typ,
 				Expr: &plan.Expr_Lit{
 					Lit: &plan.Literal{
@@ -195,11 +195,11 @@ func (builder *QueryBuilder) applyIndicesForSortUsingIvfflat(nodeID int32, projN
 			}
 		} else {
 			// If limit is not a constant, just copy it
-			tableFuncNode.Limit = DeepCopyExpr(limit)
+			tableFuncNode.BlockLimit = DeepCopyExpr(limit)
 		}
 	} else {
 		// No filters, use original limit
-		tableFuncNode.Limit = DeepCopyExpr(limit)
+		tableFuncNode.BlockLimit = DeepCopyExpr(limit)
 	}
 
 	// Determine join structure based on rankOption.mode:

pkg/vm/engine/readutil/reader.go

@@ -476,6 +476,11 @@ func (r *reader) SetBlockTop(orderBy []*plan.OrderBySpec, limit uint64) {
 	r.orderByLimit.ColPos = col.ColPos
 	r.orderByLimit.NumVec = []byte(numVec)
 	r.orderByLimit.Limit = limit
+
+	//r.orderByLimit.LowerBound = math.Inf(-1)
+	//r.orderByLimit.UpperBound = math.Inf(1)
+
+	r.orderByLimit.DistHeap = make(objectio.Float64Heap, 0, limit)
 }
 
 func (r *reader) GetOrderBy() []*plan.OrderBySpec {

pkg/vm/engine/tae/blockio/read.go

@@ -18,7 +18,7 @@ import (
 	"container/heap"
 	"context"
 	"fmt"
-	"sort"
+	"slices"
 	"time"
 
 	"github.com/matrixorigin/matrixone/pkg/common/moerr"
@@ -390,148 +390,90 @@ func HandleOrderByLimitOnIVFFlatIndex(
 		}
 	}
 
-	var sels []int64
-	var dists []float64
-
 	nullsBm := vecCol.GetNulls()
+	selectRows = slices.DeleteFunc(selectRows, func(row int64) bool {
+		return nullsBm.Contains(uint64(row))
+	})
 
-	if int(orderByLimit.Limit) < len(selectRows) {
-		// apply topn if needed
-		hp := make(vectorindex.SearchResultMaxHeap, 0, orderByLimit.Limit)
-
-		switch orderByLimit.Typ {
-		case types.T_array_float32:
-			distFunc, err := metric.ResolveDistanceFn[float32](orderByLimit.Metric)
-			if err != nil {
-				return nil, nil, err
-			}
-
-			rhs := types.BytesToArray[float32](orderByLimit.NumVec)
+	searchResults := make([]vectorindex.SearchResult, 0, len(selectRows))
 
-			for _, row := range selectRows {
-				if nullsBm.Contains(uint64(row)) {
-					continue
-				}
-				dist, err := distFunc(types.BytesToArray[float32](vecCol.GetBytesAt(int(row))), rhs)
-				if err != nil {
-					return nil, nil, err
-				}
-				dist64 := float64(dist)
+	switch orderByLimit.Typ {
+	case types.T_array_float32:
+		distFunc, err := metric.ResolveDistanceFn[float32](orderByLimit.Metric)
+		if err != nil {
+			return nil, nil, err
+		}
 
-				heapItem := &vectorindex.SearchResult{
-					Id:       row,
-					Distance: dist64,
-				}
-				if len(hp) >= int(orderByLimit.Limit) {
-					if dist64 < hp[0].GetDistance() {
-						hp[0] = heapItem
-						heap.Fix(&hp, 0)
-					}
-				} else {
-					heap.Push(&hp, heapItem)
-				}
-			}
+		rhs := types.BytesToArray[float32](orderByLimit.NumVec)
 
-		case types.T_array_float64:
-			distFunc, err := metric.ResolveDistanceFn[float64](orderByLimit.Metric)
+		for _, row := range selectRows {
+			dist, err := distFunc(types.BytesToArray[float32](vecCol.GetBytesAt(int(row))), rhs)
 			if err != nil {
 				return nil, nil, err
 			}
+			dist64 := float64(dist)
 
-			rhs := types.BytesToArray[float64](orderByLimit.NumVec)
-
-			for _, row := range selectRows {
-				if nullsBm.Contains(uint64(row)) {
-					continue
-				}
-				dist, err := distFunc(types.BytesToArray[float64](vecCol.GetBytesAt(int(row))), rhs)
-				if err != nil {
-					return nil, nil, err
-				}
-
-				heapItem := &vectorindex.SearchResult{
-					Id:       row,
-					Distance: dist,
-				}
-				if len(hp) >= int(orderByLimit.Limit) {
-					if dist < hp[0].GetDistance() {
-						hp[0] = heapItem
-						heap.Fix(&hp, 0)
-					}
+			if len(orderByLimit.DistHeap) >= int(orderByLimit.Limit) {
+				if dist64 < orderByLimit.DistHeap[0] {
+					orderByLimit.DistHeap[0] = dist64
+					heap.Fix(&orderByLimit.DistHeap, 0)
 				} else {
-					heap.Push(&hp, heapItem)
+					continue
 				}
+			} else {
+				heap.Push(&orderByLimit.DistHeap, dist64)
 			}
 
-		default:
-			return nil, nil, moerr.NewInternalError(ctx, fmt.Sprintf("only support float32/float64 type for topn: %s", orderByLimit.Typ))
+			searchResults = append(searchResults, vectorindex.SearchResult{
+				Id:       row,
+				Distance: dist64,
+			})
 		}
 
-		sRes := make([]vectorindex.SearchResult, len(hp))
-		for i := range sRes {
-			sRes[i] = *hp[i].(*vectorindex.SearchResult)
+	case types.T_array_float64:
+		distFunc, err := metric.ResolveDistanceFn[float64](orderByLimit.Metric)
+		if err != nil {
+			return nil, nil, err
 		}
-		sort.Slice(sRes, func(i, j int) bool {
-			return sRes[i].Id < sRes[j].Id
-		})
 
-		sels = make([]int64, len(sRes))
-		dists = make([]float64, len(sRes))
-
-		for i := range sRes {
-			sels[i] = sRes[i].Id
-			dists[i] = sRes[i].Distance
-		}
-	} else {
-		sels = make([]int64, 0, len(selectRows))
-		dists = make([]float64, 0, len(selectRows))
+		rhs := types.BytesToArray[float64](orderByLimit.NumVec)
 
-		switch orderByLimit.Typ {
-		case types.T_array_float32:
-			distFunc, err := metric.ResolveDistanceFn[float32](orderByLimit.Metric)
+		for _, row := range selectRows {
+			dist64, err := distFunc(types.BytesToArray[float64](vecCol.GetBytesAt(int(row))), rhs)
 			if err != nil {
 				return nil, nil, err
 			}
 
-			rhs := types.BytesToArray[float32](orderByLimit.NumVec)
-
-			for _, row := range selectRows {
-				if nullsBm.Contains(uint64(row)) {
+			if len(orderByLimit.DistHeap) >= int(orderByLimit.Limit) {
+				if dist64 < orderByLimit.DistHeap[0] {
+					orderByLimit.DistHeap[0] = dist64
+					heap.Fix(&orderByLimit.DistHeap, 0)
+				} else {
 					continue
 				}
-				dist, err := distFunc(types.BytesToArray[float32](vecCol.GetBytesAt(int(row))), rhs)
-				if err != nil {
-					return nil, nil, err
-				}
-
-				sels = append(sels, row)
-				dists = append(dists, float64(dist))
-			}
-
-		case types.T_array_float64:
-			distFunc, err := metric.ResolveDistanceFn[float64](orderByLimit.Metric)
-			if err != nil {
-				return nil, nil, err
+			} else {
+				heap.Push(&orderByLimit.DistHeap, dist64)
 			}
 
-			rhs := types.BytesToArray[float64](orderByLimit.NumVec)
+			searchResults = append(searchResults, vectorindex.SearchResult{
+				Id:       row,
+				Distance: dist64,
+			})
+		}
 
-			for _, row := range selectRows {
-				if nullsBm.Contains(uint64(row)) {
-					continue
-				}
-				dist, err := distFunc(types.BytesToArray[float64](vecCol.GetBytesAt(int(row))), rhs)
-				if err != nil {
-					return nil, nil, err
-				}
+	default:
+		return nil, nil, moerr.NewInternalError(ctx, fmt.Sprintf("only support float32/float64 type for topn: %s", orderByLimit.Typ))
+	}
 
-				sels = append(sels, row)
-				dists = append(dists, float64(dist))
-			}
+	searchResults = slices.DeleteFunc(searchResults, func(res vectorindex.SearchResult) bool {
+		return res.Distance > orderByLimit.DistHeap[0]
+	})
 
-		default:
-			return nil, nil, moerr.NewInternalError(ctx, fmt.Sprintf("only support float32/float64 type for topn: %s", orderByLimit.Typ))
-		}
+	sels := make([]int64, len(searchResults))
+	dists := make([]float64, len(searchResults))
+	for i, res := range searchResults {
+		sels[i] = res.Id
+		dists[i] = res.Distance
 	}
 
 	return sels, dists, nil