Merge cockroachdb#155998

155998: Create a vector indexing roachtest r=mw5h a=mw5h

#### workload/vecann: add shared utility functions for vector workloads

Add DeriveDistanceMetric and CalculateRecall functions to the vecann
package to be shared between vecbench and the upcoming vector index
roachtest. These functions extract the distance metric from dataset
naming conventions and compute recall for search result validation.

Refactor vecbench to use these shared implementations, removing
duplicate code.

Informs: cockroachdb#154590
Release note: None

#### roachtest/vecindex: add a vector index test stub

Add minimal test infrastructure for vector index roachtests:
- Create vecindex.go with vecIndexOptions struct and empty
  registerVectorIndex function
- Register the test in registry.go
- Add helper functions for test naming, key generation, and distance
  operator mapping.
- Stubs out the phases of the test.

Add test configurations and registration for 5 vector index test
variants:
- vecindex/dbpedia-100k/nodes=3 (standard, no prefix)
- vecindex/dbpedia-100k/nodes=3/prefix=3 (with prefix columns)
- vecindex/dbpedia-1m/nodes=6 (large-scale)
- vecindex/random-s/nodes=1 (local development)
- vecindex/random-s/nodes=1/prefix=2 (local with prefix)

Informs: cockroachdb#154590
Release note: None

#### roachtest/vecindex: implement a test of backfill and merge

Add a test phase that loads a dataset using a pool of workers and, at a
test-specified percentage of table population, kicks off a create vector
index for the data. This allows us to test both backfill (pre-create)
and merge (post-create starting). Times are reported for both but are
not used as a pass criteria (yet).

Informs: cockroachdb#154590
Release note: None

#### roachtest/vecindex: test recall of vector ann data

This addition to the vecindex roachtest tests the recall of nearest
neighbors from the test data provided with each data set. Each test has
a configurable set of beam sizes to test and a minimum recall
correctneess that is acceptable for each beam size. Tests that load
multiple prefixes test each prefix.

Informs: cockroachdb#154590
Release note: None

#### roachtest/vecindex: add a concurrent reader/writer subtest

This subtest spins up a configurable number of readers and writers to
drive vector search load to the database.

Each writer inserts rows from the first train data file in single
row batches until it has inserted all of the rows in that file, at
which point it switches into delete mode and starts deleting rows
in 10 row batches. When all rows have been deleted, the writer once
again becomes an inserter and the process repeats.

Each reader randomly selects a beam size from the sizes configured for
the test and then runs searches for random vectors in the test data for
the dataset. The reader ignores rows inserted by the writer threads to
avoid too heavily skewing results. To do this, it searches for more
vectors than called for and then filters the output to remove vectors
written by the insert workers. When the read worker exits, it validates
its recall rate against the expected rate for the number of searches it
performed.

For multi-prefix tests, this subtest only reads and writes to the first
prefix to ensure the maximum amount of contention.

Fixes: cockroachdb#154590
Release note: None


Co-authored-by: Matt White <[email protected]>

Author: craig[bot]

build/teamcity/cockroach/nightlies/roachtest_nightly_gce_vecindex.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+# Copyright 2024 The Cockroach Authors.
+#
+# Use of this software is governed by the CockroachDB Software License
+# included in the /LICENSE file.
+
+
+set -exuo pipefail
+
+dir="$(dirname $(dirname $(dirname $(dirname "${0}"))))"
+
+source "$dir/teamcity-support.sh"  # For $root
+source "$dir/teamcity-bazel-support.sh"  # For run_bazel
+
+BAZEL_SUPPORT_EXTRA_DOCKER_ARGS="-e LITERAL_ARTIFACTS_DIR=$root/artifacts -e BUILD_VCS_NUMBER -e CLOUD -e COCKROACH_DEV_LICENSE -e TESTS -e COUNT -e GITHUB_API_TOKEN -e GITHUB_ORG -e GITHUB_REPO -e GOOGLE_EPHEMERAL_CREDENTIALS -e GOOGLE_KMS_KEY_A -e GOOGLE_KMS_KEY_B -e GOOGLE_CREDENTIALS_ASSUME_ROLE -e GOOGLE_SERVICE_ACCOUNT -e SLACK_TOKEN -e TC_BUILDTYPE_ID -e TC_BUILD_BRANCH -e TC_BUILD_ID -e TC_SERVER_URL -e SELECT_PROBABILITY -e COCKROACH_RANDOM_SEED -e ROACHTEST_ASSERTIONS_ENABLED_SEED -e ROACHTEST_FORCE_RUN_INVALID_RELEASE_BRANCH -e GRAFANA_SERVICE_ACCOUNT_JSON -e GRAFANA_SERVICE_ACCOUNT_AUDIENCE -e USE_SPOT -e SNOWFLAKE_USER -e SNOWFLAKE_PVT_KEY -e COCKROACH_EA_PROBABILITY" \
+			       run_bazel build/teamcity/cockroach/nightlies/vecindex_nightly_impl.sh

build/teamcity/cockroach/nightlies/vecindex_nightly_impl.sh

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+
+# Copyright 2024 The Cockroach Authors.
+#
+# Use of this software is governed by the CockroachDB Software License
+# included in the /LICENSE file.
+
+set -exuo pipefail
+
+dir="$(dirname $(dirname $(dirname $(dirname "${0}"))))"
+set -a
+source "$dir/teamcity-support.sh"
+set +a
+
+if [[ ! -f ~/.ssh/id_rsa.pub ]]; then
+  ssh-keygen -q -C "roachtest-nightly-bazel $(date)" -N "" -f ~/.ssh/id_rsa
+fi
+
+source $root/build/teamcity/util/roachtest_util.sh
+
+artifacts=/artifacts
+
+arch=amd64
+$root/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh $arch
+
+build/teamcity-roachtest-invoke.sh \
+  --suite vecindex \
+  --cloud "gce" \
+  --cluster-id "${TC_BUILD_ID}" \
+  --artifacts=/artifacts \
+  --artifacts-literal="${LITERAL_ARTIFACTS_DIR:-}" \
+  --parallelism="${PARALLELISM}" \
+  --cpu-quota="${CPUQUOTA}" \
+  --use-spot="${USE_SPOT:-auto}" \
+  --slack-token="${SLACK_TOKEN}" \

pkg/cmd/roachtest/registry/test_spec.go

@@ -459,12 +459,13 @@ const (
 	Acceptance            = "acceptance"
 	Perturbation          = "perturbation"
 	MixedVersion          = "mixedversion"
+	VecIndex              = "vecindex"
 )
 
 var allSuites = []string{
 	Nightly, Weekly, ReleaseQualification, ORM, Driver, Tool, Quick, Fixtures,
 	Pebble, PebbleNightlyWrite, PebbleNightlyYCSB, PebbleNightlyYCSBRace, Roachtest, Acceptance,
-	Perturbation, MixedVersion,
+	Perturbation, MixedVersion, VecIndex,
 }
 
 // SuiteSet represents a set of suites.

pkg/cmd/roachtest/registry/testdata/filter/errors

@@ -6,11 +6,11 @@ error: invalid owner "badowner"
 
 filter suite=badsuite
 ----
-error: invalid suite "badsuite"; valid suites are nightly,weekly,release_qualification,orm,driver,tool,quick,fixtures,pebble,pebble_nightly_write,pebble_nightly_ycsb,pebble_nightly_ycsb_race,roachtest,acceptance,perturbation,mixedversion
+error: invalid suite "badsuite"; valid suites are nightly,weekly,release_qualification,orm,driver,tool,quick,fixtures,pebble,pebble_nightly_write,pebble_nightly_ycsb,pebble_nightly_ycsb_race,roachtest,acceptance,perturbation,mixedversion,vecindex
 
 filter owner=badowner suite=badsuite
 ----
-error: invalid suite "badsuite"; valid suites are nightly,weekly,release_qualification,orm,driver,tool,quick,fixtures,pebble,pebble_nightly_write,pebble_nightly_ycsb,pebble_nightly_ycsb_race,roachtest,acceptance,perturbation,mixedversion
+error: invalid suite "badsuite"; valid suites are nightly,weekly,release_qualification,orm,driver,tool,quick,fixtures,pebble,pebble_nightly_write,pebble_nightly_ycsb,pebble_nightly_ycsb_race,roachtest,acceptance,perturbation,mixedversion,vecindex
 
 # Filters with one field leading to no matches.
 

pkg/cmd/roachtest/tests/BUILD.bazel

@@ -209,6 +209,7 @@ go_library(
         "typeorm.go",
         "unoptimized_query_oracle.go",
         "validate_system_schema_after_version_upgrade.go",
+        "vecindex.go",
         "versionupgrade.go",
         "ycsb.go",
     ],
@@ -284,6 +285,8 @@ go_library(
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/sem/tree",
         "//pkg/sql/ttl/ttlbase",
+        "//pkg/sql/vecindex/cspann",
+        "//pkg/sql/vecindex/vecpb",
         "//pkg/storage/enginepb",
         "//pkg/storage/fs",
         "//pkg/testutils",
@@ -310,6 +313,7 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/uuid",
+        "//pkg/util/vector",
         "//pkg/workload",
         "//pkg/workload/debug",
         "//pkg/workload/histogram",
@@ -318,6 +322,7 @@ go_library(
         "//pkg/workload/tpcc",
         "//pkg/workload/tpcds",
         "//pkg/workload/tpch",
+        "//pkg/workload/vecann",
         "@com_github_aws_aws_sdk_go//aws",
         "@com_github_aws_aws_sdk_go_v2//aws",
         "@com_github_aws_aws_sdk_go_v2_config//:config",
@@ -340,8 +345,10 @@ go_library(
         "@com_github_google_go_cmp//cmp",
         "@com_github_google_pprof//profile",
         "@com_github_ibm_sarama//:sarama",
+        "@com_github_jackc_pgconn//:pgconn",
         "@com_github_jackc_pgtype//:pgtype",
         "@com_github_jackc_pgx_v5//:pgx",
+        "@com_github_jackc_pgx_v5//pgxpool",
         "@com_github_kr_pretty//:pretty",
         "@com_github_lib_pq//:pq",
         "@com_github_montanaflynn_stats//:stats",

pkg/cmd/roachtest/tests/registry.go

@@ -169,6 +169,7 @@ func RegisterTests(r registry.Registry) {
 	registerTypeORM(r)
 	registerUnoptimizedQueryOracle(r)
 	registerValidateSystemSchemaAfterVersionUpgradeSeparateProcess(r)
+	registerVectorIndex(r)
 	registerYCSB(r)
 	registerDeclarativeSchemaChangerJobCompatibilityInMixedVersion(r)
 	registerMultiRegionMixedVersion(r)