Fix hack/benchmark search interface change (#2129)

vdaas · Jul 25, 2023 · 1b918b5 · 1b918b5
1 parent 13143ab
commit 1b918b5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/hack/benchmark/core/benchmark/strategy/search.go b/hack/benchmark/core/benchmark/strategy/search.go
@@ -43,7 +43,7 @@ func NewSearch(size int, epsilon, radius float32, opts ...StrategyOption) benchm
 				}
 				b.StartTimer()
 				defer b.StopTimer()
-				return c.Search(v.([]float32), size, epsilon, radius)
+				return c.Search(ctx, v.([]float32), size, epsilon, radius)
 			},
 		),
 		WithPreProp64(
@@ -59,7 +59,7 @@ func NewSearch(size int, epsilon, radius float32, opts ...StrategyOption) benchm
 				}
 				b.StartTimer()
 				defer b.StopTimer()
-				return c.Search(float32To64(v.([]float32)), size, epsilon, radius)
+				return c.Search(ctx, float32To64(v.([]float32)), size, epsilon, radius)
 			},
 		),
 	}, opts...)...)

diff --git a/hack/benchmark/internal/core/algorithm/algorithm.go b/hack/benchmark/internal/core/algorithm/algorithm.go
@@ -17,6 +17,8 @@
 // Package algorithm provides core interface
 package algorithm
 
+import "context"
+
 type Mode uint32
 
 const (
@@ -29,7 +31,7 @@ type Closer interface {
 }
 
 type Bit32 interface {
-	Search(vec []float32, size int, epsilon, radius float32) (interface{}, error)
+	Search(ctx context.Context, vec []float32, size int, epsilon, radius float32) (interface{}, error)
 	Insert(vec []float32) (uint, error)
 	InsertCommit(vec []float32, poolSize uint32) (uint, error)
 	BulkInsert(vecs [][]float32) ([]uint, []error)
@@ -44,7 +46,7 @@ type Bit32 interface {
 }
 
 type Bit64 interface {
-	Search(vec []float64, size int, epsilon, radius float32) (interface{}, error)
+	Search(ctx context.Context, vec []float64, size int, epsilon, radius float32) (interface{}, error)
 	Insert(vec []float64) (uint, error)
 	InsertCommit(vec []float64, poolSize uint32) (uint, error)
 	BulkInsert(vecs [][]float64) ([]uint, []error)

diff --git a/hack/benchmark/internal/core/algorithm/ngt/ngt.go b/hack/benchmark/internal/core/algorithm/ngt/ngt.go
@@ -18,6 +18,7 @@
 package ngt
 
 import (
+	"context"
 	"os"
 
 	c "github.com/vdaas/vald/hack/benchmark/internal/core/algorithm"
@@ -74,8 +75,8 @@ func New(opts ...Option) (c.Bit32, error) {
 	return c, nil
 }
 
-func (c *core) Search(vec []float32, size int, epsilon, radius float32) (interface{}, error) {
-	return c.NGT.Search(vec, size, epsilon, radius)
+func (c *core) Search(ctx context.Context, vec []float32, size int, epsilon, radius float32) (interface{}, error) {
+	return c.NGT.Search(ctx, vec, size, epsilon, radius)
 }
 
 func (c *core) Close() {