From 51eab702d1ffee9f1d3e2fc9c774856b5dbb0981 Mon Sep 17 00:00:00 2001 From: zhenghaoz Date: Wed, 12 Aug 2026 08:27:02 +0800 Subject: [PATCH 1/2] [verified] feat(bench): support xvec and zvec backends --- cmd/vector-db-bench/README.md | 43 +++- cmd/vector-db-bench/backend.go | 58 +++++ cmd/vector-db-bench/benchmark.go | 36 +-- cmd/vector-db-bench/benchmark_test.go | 7 + cmd/vector-db-bench/config.go | 17 +- cmd/vector-db-bench/config_test.go | 14 +- cmd/vector-db-bench/report.go | 7 +- cmd/vector-db-bench/zvec_backend.go | 313 +++++++++++++++++++++++ cmd/vector-db-bench/zvec_backend_stub.go | 37 +++ cmd/vector-db-bench/zvec_backend_test.go | 29 +++ go.mod | 2 + go.sum | 4 + 12 files changed, 535 insertions(+), 32 deletions(-) create mode 100644 cmd/vector-db-bench/backend.go create mode 100644 cmd/vector-db-bench/zvec_backend.go create mode 100644 cmd/vector-db-bench/zvec_backend_stub.go create mode 100644 cmd/vector-db-bench/zvec_backend_test.go diff --git a/cmd/vector-db-bench/README.md b/cmd/vector-db-bench/README.md index b3f33fc..d51bca9 100644 --- a/cmd/vector-db-bench/README.md +++ b/cmd/vector-db-bench/README.md @@ -1,8 +1,9 @@ # vector-db-bench -`vector-db-bench` is a native Go benchmark driver for xvec. It follows the -VectorDBBench Cohere performance workload used by the Alibaba zvec benchmark -guide: +`vector-db-bench` is a native Go benchmark driver for comparing +[xvec](https://github.com/gorse-io/xvec) and +[zvec-go](https://github.com/zvec-ai/zvec-go). It follows the VectorDBBench +Cohere performance workload used by the Alibaba zvec benchmark guide: - Cohere 1M and 10M Parquet datasets; - HNSW loading and optimization; @@ -18,7 +19,20 @@ substantially more. ## Build ```bash -go build -o vector-db-bench ./cmd/vector-db-bench +CGO_ENABLED=0 go build -o vector-db-bench ./cmd/vector-db-bench +``` + +The pure-Go build keeps xvec and zvec-go in one binary without linking zvec at +build time. Running the `zvec` backend requires the zvec C API shared library. +Download the archive for your platform from the +[zvec-go v0.6.0 release](https://github.com/zvec-ai/zvec-go/releases/tag/v0.6.0), +extract it, and point `ZVEC_LIBRARY_PATH` to the extracted library or its +directory. The xvec backend does not require this library. + +For example, on Linux x86-64: + +```bash +export ZVEC_LIBRARY_PATH=/path/to/linux_amd64/libzvec_c_api.so ``` ## Cohere 1M @@ -28,7 +42,7 @@ first invocation downloads the data, recreates the collection, loads it, and runs both search phases. ```bash -./vector-db-bench \ +./vector-db-bench xvec \ --path ./Performance768D1M \ --case-type Performance768D1M \ --num-concurrency 12,14,16,18,20 \ @@ -37,10 +51,23 @@ runs both search phases. --output result-cohere-1m.json ``` +Run the same workload against zvec-go by changing only the backend and output +path: + +```bash +./vector-db-bench zvec \ + --path ./Performance768D1M-zvec \ + --case-type Performance768D1M \ + --num-concurrency 12,14,16,18,20 \ + --m 15 \ + --ef-search 180 \ + --output result-zvec-cohere-1m.json +``` + To rerun only the search phases against that collection: ```bash -./vector-db-bench \ +./vector-db-bench xvec \ --path ./Performance768D1M \ --case-type Performance768D1M \ --num-concurrency 12,14,16,18,20 \ @@ -56,7 +83,7 @@ To rerun only the search phases against that collection: This mirrors the published INT8/refiner configuration: ```bash -./vector-db-bench \ +./vector-db-bench xvec \ --path ./Performance768D10M \ --case-type Performance768D10M \ --num-concurrency 12,14,16,18,20 \ @@ -88,7 +115,7 @@ The built-in cases use the VectorDBBench schema: A local custom dataset can be exercised without downloads: ```bash -./vector-db-bench \ +./vector-db-bench xvec \ --path ./custom-collection \ --case-type Custom \ --dataset-dir ./dataset \ diff --git a/cmd/vector-db-bench/backend.go b/cmd/vector-db-bench/backend.go new file mode 100644 index 0000000..a771055 --- /dev/null +++ b/cmd/vector-db-bench/backend.go @@ -0,0 +1,58 @@ +// Copyright 2026-present the xvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "fmt" + "io" + + "github.com/gorse-io/xvec" +) + +func initializeBenchmarkBackend(backend string) (func(), error) { + if backend == backendZvec { + return initializeZvecBackend() + } + return func() {}, nil +} + +func loadBenchmarkDataset(ctx context.Context, config benchConfig, log io.Writer) (loadMetrics, error) { + switch config.Backend { + case backendXvec: + return loadXvecDataset(ctx, config, log) + case backendZvec: + return loadZvecDataset(ctx, config, log) + default: + return loadMetrics{}, fmt.Errorf("unsupported backend %q", config.Backend) + } +} + +func openBenchmarkQueryEngine(ctx context.Context, config benchConfig) (benchmarkQueryEngine, io.Closer, error) { + switch config.Backend { + case backendXvec: + collection, err := xvec.Open(ctx, config.Path, xvec.CollectionOptions{ + ReadOnly: true, EnableMmap: config.EnableMmap, MaxBufferSize: uint32(config.MaxBufferSize), + }) + if err != nil { + return nil, nil, fmt.Errorf("open xvec benchmark collection for search: %w", err) + } + return newXvecQueryEngine(collection, config), collection, nil + case backendZvec: + return openZvecQueryEngine(config) + default: + return nil, nil, fmt.Errorf("unsupported backend %q", config.Backend) + } +} diff --git a/cmd/vector-db-bench/benchmark.go b/cmd/vector-db-bench/benchmark.go index d5a7df6..477a612 100644 --- a/cmd/vector-db-bench/benchmark.go +++ b/cmd/vector-db-bench/benchmark.go @@ -31,6 +31,11 @@ import ( func runBenchmark(ctx context.Context, config benchConfig, log io.Writer) (benchmarkReport, error) { report := newBenchmarkReport(config) + shutdown, err := initializeBenchmarkBackend(config.Backend) + if err != nil { + return report, err + } + defer shutdown() needsSearch := !config.SkipSerialSearch || !config.SkipConcurrentSearch if !config.SkipLoad || needsSearch { if err := prepareDataset(ctx, config, !config.SkipLoad, log); err != nil { @@ -38,7 +43,7 @@ func runBenchmark(ctx context.Context, config benchConfig, log io.Writer) (bench } } if !config.SkipLoad { - metrics, err := loadDataset(ctx, config, log) + metrics, err := loadBenchmarkDataset(ctx, config, log) if err != nil { return report, err } @@ -53,14 +58,11 @@ func runBenchmark(ctx context.Context, config benchConfig, log io.Writer) (bench if err != nil { return report, err } - collection, err := xvec.Open(ctx, config.Path, xvec.CollectionOptions{ - ReadOnly: true, EnableMmap: config.EnableMmap, MaxBufferSize: uint32(config.MaxBufferSize), - }) + engine, closer, err := openBenchmarkQueryEngine(ctx, config) if err != nil { - return report, fmt.Errorf("open benchmark collection for search: %w", err) + return report, err } - defer func() { _ = collection.Close() }() - engine := newQueryEngine(collection, config) + defer func() { _ = closer.Close() }() if err := warmupSearch(ctx, engine, data, config.WarmupQueries); err != nil { return report, err } @@ -94,7 +96,7 @@ func runBenchmark(ctx context.Context, config benchConfig, log io.Writer) (bench return report, nil } -func loadDataset(ctx context.Context, config benchConfig, log io.Writer) (loadMetrics, error) { +func loadXvecDataset(ctx context.Context, config benchConfig, log io.Writer) (loadMetrics, error) { collection, err := writableBenchmarkCollection(ctx, config) if err != nil { return loadMetrics{}, err @@ -243,20 +245,24 @@ func parseQuantization(value string) (xvec.QuantizeType, error) { } } -type queryEngine struct { +type benchmarkQueryEngine interface { + search(context.Context, []float32) ([]int64, error) +} + +type xvecQueryEngine struct { collection *xvec.Collection params xvec.HNSWQueryParams k int } -func newQueryEngine(collection *xvec.Collection, config benchConfig) queryEngine { +func newXvecQueryEngine(collection *xvec.Collection, config benchConfig) xvecQueryEngine { params := xvec.NewHNSWQueryParams() params.EF = config.EFSearch params.UseRefiner = config.UseRefiner - return queryEngine{collection: collection, params: params, k: config.K} + return xvecQueryEngine{collection: collection, params: params, k: config.K} } -func (e queryEngine) search(ctx context.Context, vector []float32) ([]int64, error) { +func (e xvecQueryEngine) search(ctx context.Context, vector []float32) ([]int64, error) { results, err := e.collection.Query(ctx, xvec.VectorQuery{ Field: "dense", DenseVector: xvec.VectorFP32(vector), TopK: e.k, Params: e.params, Projection: xvec.Projection{OutputFields: []string{}}, @@ -274,7 +280,7 @@ func (e queryEngine) search(ctx context.Context, vector []float32) ([]int64, err return ids, nil } -func warmupSearch(ctx context.Context, engine queryEngine, data queryData, count int) error { +func warmupSearch(ctx context.Context, engine benchmarkQueryEngine, data queryData, count int) error { if count > len(data.Vectors) { count = len(data.Vectors) } @@ -286,7 +292,7 @@ func warmupSearch(ctx context.Context, engine queryEngine, data queryData, count return nil } -func runSerialSearch(ctx context.Context, engine queryEngine, data queryData, k int) (searchMetrics, error) { +func runSerialSearch(ctx context.Context, engine benchmarkQueryEngine, data queryData, k int) (searchMetrics, error) { latencies := make([]float64, len(data.Vectors)) var recall float64 started := time.Now() @@ -312,7 +318,7 @@ type concurrentWorkerResult struct { func runConcurrentSearch( ctx context.Context, - engine queryEngine, + engine benchmarkQueryEngine, data queryData, concurrency int, duration time.Duration, diff --git a/cmd/vector-db-bench/benchmark_test.go b/cmd/vector-db-bench/benchmark_test.go index 64cad6f..44be93b 100644 --- a/cmd/vector-db-bench/benchmark_test.go +++ b/cmd/vector-db-bench/benchmark_test.go @@ -39,6 +39,11 @@ func TestRecallPercentileAndSearchSummary(t *testing.T) { } func TestVectorDBBenchEndToEndCustomDataset(t *testing.T) { + testVectorDBBenchEndToEndCustomDataset(t, backendXvec) +} + +func testVectorDBBenchEndToEndCustomDataset(t *testing.T, backend string) { + t.Helper() directory := t.TempDir() datasetDir := filepath.Join(directory, "dataset") require.NoError(t, mkdir(datasetDir)) @@ -62,6 +67,7 @@ func TestVectorDBBenchEndToEndCustomDataset(t *testing.T) { var stdout, stderr bytes.Buffer err := runCLI(context.Background(), []string{ + backend, "--path", filepath.Join(directory, "collection"), "--case-type", caseCustom, "--dataset-dir", datasetDir, @@ -82,6 +88,7 @@ func TestVectorDBBenchEndToEndCustomDataset(t *testing.T) { require.NoError(t, err, stderr.String()) var report benchmarkReport require.NoError(t, json.Unmarshal(stdout.Bytes(), &report), fmt.Sprintf("stdout: %s", stdout.String())) + require.Equal(t, backend, report.Config.Backend) require.NotNil(t, report.Load) require.Equal(t, int64(32), report.Load.Rows) require.NotNil(t, report.Serial) diff --git a/cmd/vector-db-bench/config.go b/cmd/vector-db-bench/config.go index 667b279..1a7fefe 100644 --- a/cmd/vector-db-bench/config.go +++ b/cmd/vector-db-bench/config.go @@ -27,6 +27,9 @@ import ( ) const ( + backendXvec = "xvec" + backendZvec = "zvec" + casePerformance768D1M = "Performance768D1M" casePerformance768D10M = "Performance768D10M" caseCustom = "Custom" @@ -43,6 +46,7 @@ type benchmarkCase struct { } type benchConfig struct { + Backend string Path string CaseType string DatasetDir string @@ -84,9 +88,16 @@ type benchConfig struct { func parseConfig(args []string, stderr io.Writer) (benchConfig, error) { var config benchConfig + if len(args) == 0 || strings.HasPrefix(args[0], "-") { + return benchConfig{}, errors.New("backend is required: xvec or zvec") + } + config.Backend = strings.ToLower(args[0]) + if config.Backend != backendXvec && config.Backend != backendZvec { + return benchConfig{}, fmt.Errorf("unsupported backend %q: use xvec or zvec", args[0]) + } flags := flag.NewFlagSet("vector-db-bench", flag.ContinueOnError) flags.SetOutput(stderr) - flags.StringVar(&config.Path, "path", "", "xvec collection path (required)") + flags.StringVar(&config.Path, "path", "", "collection path (required)") flags.StringVar(&config.CaseType, "case-type", casePerformance768D1M, "benchmark case: Performance768D1M, Performance768D10M, or Custom") flags.StringVar(&config.DatasetDir, "dataset-dir", "", "local VectorDBBench dataset directory") flags.StringVar(&config.DatasetBaseURL, "dataset-base-url", "https://assets.zilliz.com/benchmark", "VectorDBBench dataset base URL") @@ -117,11 +128,11 @@ func parseConfig(args []string, stderr io.Writer) (benchConfig, error) { flags.BoolVar(&config.SkipConcurrentSearch, "skip-search-concurrent", false, "skip sustained concurrent search") flags.BoolVar(&config.DryRun, "dry-run", false, "validate and print configuration without downloading or running") flags.StringVar(&config.Output, "output", "", "write result JSON to this file; empty writes JSON to stdout") - flags.StringVar(&config.DBLabel, "db-label", "xvec-go", "label stored in the result") + flags.StringVar(&config.DBLabel, "db-label", config.Backend+"-go", "label stored in the result") flags.StringVar(&config.Note, "note", "", "non-sensitive run context stored in the result") operationTimeout := flags.String("operation-timeout", "0", "whole-run timeout; zero disables it") flags.Int64Var(&config.Seed, "seed", 0, "deterministic concurrent-query seed") - if err := flags.Parse(args); err != nil { + if err := flags.Parse(args[1:]); err != nil { return benchConfig{}, err } if flags.NArg() != 0 { diff --git a/cmd/vector-db-bench/config_test.go b/cmd/vector-db-bench/config_test.go index 3328352..7d0c6ac 100644 --- a/cmd/vector-db-bench/config_test.go +++ b/cmd/vector-db-bench/config_test.go @@ -24,6 +24,7 @@ import ( func TestParseConfigVectorDBBenchCases(t *testing.T) { config, err := parseConfig([]string{ + backendXvec, "--path", t.TempDir(), "--case-type", casePerformance768D10M, "--num-concurrency", "12,14,16", @@ -33,6 +34,7 @@ func TestParseConfigVectorDBBenchCases(t *testing.T) { "--is-using-refiner", }, &bytes.Buffer{}) require.NoError(t, err) + require.Equal(t, backendXvec, config.Backend) require.Equal(t, casePerformance768D10M, config.caseSpec.Name) require.Equal(t, 768, config.caseSpec.Dimension) require.Equal(t, int64(10_000_000), config.caseSpec.Size) @@ -46,20 +48,26 @@ func TestParseConfigVectorDBBenchCases(t *testing.T) { func TestParseConfigCustomAndValidation(t *testing.T) { datasetDir := t.TempDir() config, err := parseConfig([]string{ + backendZvec, "--path", t.TempDir(), "--case-type", caseCustom, "--dataset-dir", datasetDir, "--dimension", "3", "--metric", "l2", "--train-files", "part-0.parquet,part-1.parquet", }, &bytes.Buffer{}) require.NoError(t, err) + require.Equal(t, backendZvec, config.Backend) require.Equal(t, []string{"part-0.parquet", "part-1.parquet"}, config.caseSpec.TrainFiles) require.Equal(t, "l2", config.caseSpec.Metric) - _, err = parseConfig([]string{"--path", t.TempDir(), "--skip-load"}, &bytes.Buffer{}) + _, err = parseConfig([]string{backendXvec, "--path", t.TempDir(), "--skip-load"}, &bytes.Buffer{}) require.ErrorContains(t, err, "skip-load requires skip-drop-old") - _, err = parseConfig([]string{"--path", t.TempDir(), "--num-concurrency", "1,1"}, &bytes.Buffer{}) + _, err = parseConfig([]string{backendXvec, "--path", t.TempDir(), "--num-concurrency", "1,1"}, &bytes.Buffer{}) require.ErrorContains(t, err, "duplicate concurrency") - _, err = parseConfig([]string{"--path", t.TempDir(), "--serial-cooldown", "-1"}, &bytes.Buffer{}) + _, err = parseConfig([]string{backendXvec, "--path", t.TempDir(), "--serial-cooldown", "-1"}, &bytes.Buffer{}) require.ErrorContains(t, err, "serial-cooldown cannot be negative") + _, err = parseConfig([]string{"--path", t.TempDir()}, &bytes.Buffer{}) + require.ErrorContains(t, err, "backend is required") + _, err = parseConfig([]string{"unknown", "--path", t.TempDir()}, &bytes.Buffer{}) + require.ErrorContains(t, err, "unsupported backend") } func TestParseFlexibleDuration(t *testing.T) { diff --git a/cmd/vector-db-bench/report.go b/cmd/vector-db-bench/report.go index cac5f64..a185e41 100644 --- a/cmd/vector-db-bench/report.go +++ b/cmd/vector-db-bench/report.go @@ -27,6 +27,7 @@ import ( const reportSchemaVersion = "vector-db-bench/v1" type reportConfig struct { + Backend string `json:"backend"` Path string `json:"path"` DBLabel string `json:"db_label"` M int `json:"m"` @@ -85,7 +86,7 @@ type benchmarkReport struct { } // vectorDBBenchMetric mirrors the names used by VectorDBBench's Metric model -// so result consumers can ingest the core xvec performance fields directly. +// so result consumers can ingest the core performance fields directly. type vectorDBBenchMetric struct { InsertedCount int64 `json:"inserted_count"` InsertDuration float64 `json:"insert_duration"` @@ -117,13 +118,13 @@ func newBenchmarkReport(config benchConfig) benchmarkReport { } return benchmarkReport{ SchemaVersion: reportSchemaVersion, - Tool: "xvec-go/cmd/vector-db-bench", + Tool: "xvec/cmd/vector-db-bench", Timestamp: time.Now().UTC(), Case: config.caseSpec, DatasetDir: config.DatasetDir, Note: config.Note, Config: reportConfig{ - Path: config.Path, DBLabel: config.DBLabel, + Backend: config.Backend, Path: config.Path, DBLabel: config.DBLabel, M: config.M, EFConstruction: config.EFConstruction, EFSearch: config.EFSearch, QuantizeType: quantize, UseRefiner: config.UseRefiner, K: config.K, BatchSize: config.BatchSize, LoadLimit: config.LoadLimit, QueryLimit: config.QueryLimit, diff --git a/cmd/vector-db-bench/zvec_backend.go b/cmd/vector-db-bench/zvec_backend.go new file mode 100644 index 0000000..d817540 --- /dev/null +++ b/cmd/vector-db-bench/zvec_backend.go @@ -0,0 +1,313 @@ +// Copyright 2026-present the xvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build purego || !cgo + +package main + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "strconv" + "strings" + "time" + + zvec "github.com/zvec-ai/zvec-go" +) + +func initializeZvecBackend() (func(), error) { + if err := zvec.Initialize(nil); err != nil { + return nil, fmt.Errorf("initialize zvec: %w", err) + } + return func() { _ = zvec.Shutdown() }, nil +} + +func loadZvecDataset(ctx context.Context, config benchConfig, log io.Writer) (loadMetrics, error) { + collection, err := writableZvecBenchmarkCollection(config) + if err != nil { + return loadMetrics{}, err + } + closed := false + defer func() { + if !closed { + _ = collection.Close() + } + }() + + loadStarted := time.Now() + insertStarted := loadStarted + inserted := int64(0) + nextProgress := int64(100_000) + rows, err := forEachTrainingBatch( + ctx, config.DatasetDir, config.caseSpec.TrainFiles, config.BatchSize, config.LoadLimit, + func(rows []vectorParquetRow) error { + documents := make([]*zvec.Doc, len(rows)) + defer func() { + for _, document := range documents { + if document != nil { + document.Destroy() + } + } + }() + for index, row := range rows { + if len(row.Embedding) != config.caseSpec.Dimension { + return fmt.Errorf("training vector %d has dimension %d, want %d", row.ID, len(row.Embedding), config.caseSpec.Dimension) + } + document := zvec.NewDoc() + if document == nil { + return errors.New("create zvec document") + } + documents[index] = document + document.SetPK(strconv.FormatInt(row.ID, 10)) + if err := document.AddInt64Field("id", row.ID); err != nil { + return fmt.Errorf("set zvec document id %d: %w", row.ID, err) + } + if err := document.AddVectorFP32Field("dense", row.Embedding); err != nil { + return fmt.Errorf("set zvec document vector %d: %w", row.ID, err) + } + } + result, err := collection.Insert(documents) + if err != nil { + return fmt.Errorf("insert zvec batch at row %d: %w", inserted, err) + } + if result.ErrorCount != 0 || result.SuccessCount != uint64(len(documents)) { + return fmt.Errorf("insert zvec batch at row %d: %d succeeded, %d failed", inserted, result.SuccessCount, result.ErrorCount) + } + inserted += int64(len(documents)) + if inserted >= nextProgress { + _, _ = fmt.Fprintf(log, "inserted %d vectors (%.1f rows/s)\n", inserted, float64(inserted)/time.Since(insertStarted).Seconds()) + nextProgress = (inserted/100_000 + 1) * 100_000 + } + return nil + }, + ) + if err != nil { + return loadMetrics{}, err + } + insertDuration := time.Since(insertStarted) + optimizeStarted := time.Now() + if err := collection.Optimize(); err != nil { + return loadMetrics{}, fmt.Errorf("optimize zvec benchmark collection: %w", err) + } + optimizeDuration := time.Since(optimizeStarted) + if err := collection.Flush(); err != nil { + return loadMetrics{}, fmt.Errorf("flush zvec benchmark collection: %w", err) + } + if err := collection.Close(); err != nil { + return loadMetrics{}, fmt.Errorf("close loaded zvec collection: %w", err) + } + closed = true + metrics := loadMetrics{ + Rows: rows, InsertDurationSec: insertDuration.Seconds(), OptimizeDurationSec: optimizeDuration.Seconds(), + LoadDurationSec: time.Since(loadStarted).Seconds(), + } + if insertDuration > 0 { + metrics.RowsPerSecond = float64(rows) / insertDuration.Seconds() + } + return metrics, nil +} + +func writableZvecBenchmarkCollection(config benchConfig) (*zvec.Collection, error) { + options, err := newZvecCollectionOptions(config, false) + if err != nil { + return nil, err + } + defer options.Destroy() + if config.SkipDropOld { + collection, err := zvec.Open(config.Path, options) + if err != nil { + return nil, fmt.Errorf("open existing zvec benchmark collection: %w", err) + } + return collection, nil + } + if _, err := os.Stat(config.Path); err == nil { + old, err := zvec.Open(config.Path, options) + if err != nil { + return nil, fmt.Errorf("open old zvec benchmark collection before destroy: %w", err) + } + if err := old.Destroy(); err != nil { + return nil, fmt.Errorf("destroy old zvec benchmark collection: %w", err) + } + } else if !errors.Is(err, os.ErrNotExist) { + return nil, fmt.Errorf("stat zvec benchmark collection: %w", err) + } + + metric, err := parseZvecMetric(config.caseSpec.Metric) + if err != nil { + return nil, err + } + index, err := zvec.NewHNSWIndexParams(metric, config.M, config.EFConstruction) + if err != nil { + return nil, fmt.Errorf("create zvec HNSW index params: %w", err) + } + defer index.Destroy() + quantize, err := parseZvecQuantization(config.Quantize) + if err != nil { + return nil, err + } + if err := index.SetQuantizeType(quantize); err != nil { + return nil, fmt.Errorf("set zvec quantization: %w", err) + } + + schema := zvec.NewCollectionSchema("vector_bench_test") + if schema == nil { + return nil, errors.New("create zvec collection schema") + } + defer schema.Destroy() + if err := schema.SetMaxDocCountPerSegment(config.MaxDocsPerSegment); err != nil { + return nil, fmt.Errorf("set zvec maximum documents per segment: %w", err) + } + idField := zvec.NewFieldSchema("id", zvec.DataTypeInt64, false, 0) + if idField == nil { + return nil, errors.New("create zvec id field") + } + defer idField.Destroy() + if err := schema.AddField(idField); err != nil { + return nil, fmt.Errorf("add zvec id field: %w", err) + } + vectorField := zvec.NewFieldSchema("dense", zvec.DataTypeVectorFP32, false, uint32(config.caseSpec.Dimension)) + if vectorField == nil { + return nil, errors.New("create zvec vector field") + } + defer vectorField.Destroy() + if err := vectorField.SetIndexParams(index); err != nil { + return nil, fmt.Errorf("set zvec vector index: %w", err) + } + if err := schema.AddField(vectorField); err != nil { + return nil, fmt.Errorf("add zvec vector field: %w", err) + } + collection, err := zvec.CreateAndOpen(config.Path, schema, options) + if err != nil { + return nil, fmt.Errorf("create zvec benchmark collection: %w", err) + } + return collection, nil +} + +func newZvecCollectionOptions(config benchConfig, readOnly bool) (*zvec.CollectionOptions, error) { + options := zvec.NewCollectionOptions() + if options == nil { + return nil, errors.New("create zvec collection options") + } + if err := options.SetEnableMmap(config.EnableMmap); err != nil { + options.Destroy() + return nil, fmt.Errorf("set zvec mmap option: %w", err) + } + if err := options.SetMaxBufferSize(uint64(config.MaxBufferSize)); err != nil { + options.Destroy() + return nil, fmt.Errorf("set zvec maximum buffer size: %w", err) + } + if err := options.SetReadOnly(readOnly); err != nil { + options.Destroy() + return nil, fmt.Errorf("set zvec read-only option: %w", err) + } + return options, nil +} + +func parseZvecMetric(value string) (zvec.MetricType, error) { + switch strings.ToLower(value) { + case "cosine": + return zvec.MetricTypeCosine, nil + case "l2": + return zvec.MetricTypeL2, nil + case "ip": + return zvec.MetricTypeIP, nil + default: + return 0, fmt.Errorf("unsupported metric %q", value) + } +} + +func parseZvecQuantization(value string) (zvec.QuantizeType, error) { + switch strings.ToLower(value) { + case "", "none": + return zvec.QuantizeTypeUndefined, nil + case "fp16": + return zvec.QuantizeTypeFP16, nil + case "int8": + return zvec.QuantizeTypeInt8, nil + case "int4": + return zvec.QuantizeTypeInt4, nil + default: + return 0, fmt.Errorf("unsupported quantization %q", value) + } +} + +type zvecQueryEngine struct { + collection *zvec.Collection + ef int + useRefiner bool + k int +} + +func openZvecQueryEngine(config benchConfig) (benchmarkQueryEngine, io.Closer, error) { + options, err := newZvecCollectionOptions(config, true) + if err != nil { + return nil, nil, err + } + defer options.Destroy() + collection, err := zvec.Open(config.Path, options) + if err != nil { + return nil, nil, fmt.Errorf("open zvec benchmark collection for search: %w", err) + } + return zvecQueryEngine{collection: collection, ef: config.EFSearch, useRefiner: config.UseRefiner, k: config.K}, collection, nil +} + +func (e zvecQueryEngine) search(ctx context.Context, vector []float32) ([]int64, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + query := zvec.NewSearchQuery() + if query == nil { + return nil, errors.New("create zvec search query") + } + defer query.Destroy() + if err := query.SetFieldName("dense"); err != nil { + return nil, fmt.Errorf("set zvec query field: %w", err) + } + if err := query.SetQueryVector(vector); err != nil { + return nil, fmt.Errorf("set zvec query vector: %w", err) + } + if err := query.SetTopK(e.k); err != nil { + return nil, fmt.Errorf("set zvec query top K: %w", err) + } + if err := query.SetIncludeVector(false); err != nil { + return nil, fmt.Errorf("set zvec query vector projection: %w", err) + } + if err := query.SetIncludeDocID(false); err != nil { + return nil, fmt.Errorf("set zvec query document ID projection: %w", err) + } + params := zvec.NewHNSWQueryParams(e.ef, -1, false, e.useRefiner) + if params == nil { + return nil, errors.New("create zvec HNSW query params") + } + defer params.Destroy() + if err := query.SetHNSWParams(params); err != nil { + return nil, fmt.Errorf("set zvec HNSW query params: %w", err) + } + results, err := e.collection.Query(query) + if err != nil { + return nil, err + } + defer zvec.FreeDocs(results) + ids := make([]int64, len(results)) + for index, result := range results { + ids[index], err = strconv.ParseInt(result.GetPK(), 10, 64) + if err != nil { + return nil, fmt.Errorf("parse zvec result primary key %q: %w", result.GetPK(), err) + } + } + return ids, ctx.Err() +} diff --git a/cmd/vector-db-bench/zvec_backend_stub.go b/cmd/vector-db-bench/zvec_backend_stub.go new file mode 100644 index 0000000..2d53810 --- /dev/null +++ b/cmd/vector-db-bench/zvec_backend_stub.go @@ -0,0 +1,37 @@ +// Copyright 2026-present the xvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build cgo && !purego + +package main + +import ( + "context" + "errors" + "io" +) + +var errZvecBackendUnavailable = errors.New("zvec backend requires a pure-Go benchmark build: CGO_ENABLED=0 go build -o vector-db-bench ./cmd/vector-db-bench") + +func initializeZvecBackend() (func(), error) { + return nil, errZvecBackendUnavailable +} + +func loadZvecDataset(context.Context, benchConfig, io.Writer) (loadMetrics, error) { + return loadMetrics{}, errZvecBackendUnavailable +} + +func openZvecQueryEngine(benchConfig) (benchmarkQueryEngine, io.Closer, error) { + return nil, nil, errZvecBackendUnavailable +} diff --git a/cmd/vector-db-bench/zvec_backend_test.go b/cmd/vector-db-bench/zvec_backend_test.go new file mode 100644 index 0000000..9862d8f --- /dev/null +++ b/cmd/vector-db-bench/zvec_backend_test.go @@ -0,0 +1,29 @@ +// Copyright 2026-present the xvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build purego || !cgo + +package main + +import ( + "os" + "testing" +) + +func TestVectorDBBenchZvecEndToEndCustomDataset(t *testing.T) { + if os.Getenv("ZVEC_LIBRARY_PATH") == "" { + t.Skip("ZVEC_LIBRARY_PATH is not set") + } + testVectorDBBenchEndToEndCustomDataset(t, backendZvec) +} diff --git a/go.mod b/go.mod index 63233e2..ebc9acd 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ require ( github.com/gofrs/flock v0.13.0 github.com/parquet-go/parquet-go v0.30.1 github.com/stretchr/testify v1.11.1 + github.com/zvec-ai/zvec-go v0.6.1-0.20260721023313-9199195b29da golang.org/x/sync v0.22.0 golang.org/x/sys v0.47.0 ) @@ -28,6 +29,7 @@ require ( github.com/cockroachdb/swiss v0.0.0-20251224182025-b0f6560f979b // indirect github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/ebitengine/purego v0.10.1 // indirect github.com/getsentry/sentry-go v0.27.0 // indirect github.com/gogo/protobuf v1.3.2 // indirect github.com/golang/protobuf v1.5.3 // indirect diff --git a/go.sum b/go.sum index a5e063e..7edc2de 100644 --- a/go.sum +++ b/go.sum @@ -46,6 +46,8 @@ github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6N github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/ebitengine/purego v0.10.1 h1:dewVBCBT2GaMu1SrNTYxQhgQBethzfhiwvZiLGP/qyY= +github.com/ebitengine/purego v0.10.1/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ= github.com/getsentry/sentry-go v0.27.0 h1:Pv98CIbtB3LkMWmXi4Joa5OOcwbmnX88sF5qbK3r3Ps= github.com/getsentry/sentry-go v0.27.0/go.mod h1:lc76E2QywIyW8WuBnwl8Lc4bkmQH4+w1gwTf25trprY= github.com/ghemawat/stream v0.0.0-20171120220530-696b145b53b9 h1:r5GgOLGbza2wVHRzK7aAj6lWZjfbAwiu/RDCVOKjRyM= @@ -133,6 +135,8 @@ github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3i github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +github.com/zvec-ai/zvec-go v0.6.1-0.20260721023313-9199195b29da h1:4wINeawyVOYz/Rj4mDJQlSAUYLkQ76QELU1dd2IEU3k= +github.com/zvec-ai/zvec-go v0.6.1-0.20260721023313-9199195b29da/go.mod h1:xT+sd/4zDEvGRM7OLByMzHNz9xQjV+HXftliO7RmCbI= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= From d11c826c1a743073ae237f4e3e79236942e043c5 Mon Sep 17 00:00:00 2001 From: zhenghaoz Date: Wed, 12 Aug 2026 08:37:51 +0800 Subject: [PATCH 2/2] [verified] test(bench): cover backend dispatch --- cmd/vector-db-bench/backend_stub_test.go | 53 ++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 cmd/vector-db-bench/backend_stub_test.go diff --git a/cmd/vector-db-bench/backend_stub_test.go b/cmd/vector-db-bench/backend_stub_test.go new file mode 100644 index 0000000..1a39607 --- /dev/null +++ b/cmd/vector-db-bench/backend_stub_test.go @@ -0,0 +1,53 @@ +// Copyright 2026-present the xvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build cgo && !purego + +package main + +import ( + "bytes" + "context" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBenchmarkBackendDispatchAndZvecStub(t *testing.T) { + shutdown, err := initializeBenchmarkBackend(backendXvec) + require.NoError(t, err) + require.NotNil(t, shutdown) + shutdown() + + _, err = initializeBenchmarkBackend(backendZvec) + require.ErrorIs(t, err, errZvecBackendUnavailable) + + config := benchConfig{Backend: backendZvec} + _, err = loadBenchmarkDataset(context.Background(), config, &bytes.Buffer{}) + require.ErrorIs(t, err, errZvecBackendUnavailable) + _, _, err = openBenchmarkQueryEngine(context.Background(), config) + require.ErrorIs(t, err, errZvecBackendUnavailable) + + _, err = loadBenchmarkDataset(context.Background(), benchConfig{Backend: "unknown"}, &bytes.Buffer{}) + require.ErrorContains(t, err, "unsupported backend") + _, _, err = openBenchmarkQueryEngine(context.Background(), benchConfig{Backend: "unknown"}) + require.ErrorContains(t, err, "unsupported backend") + + _, _, err = openBenchmarkQueryEngine(context.Background(), benchConfig{ + Backend: backendXvec, + Path: filepath.Join(t.TempDir(), "missing"), + }) + require.ErrorContains(t, err, "open xvec benchmark collection") +}