From 10f2a4798f28c490555d916084ed1cdca395c947 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Mon, 25 Nov 2024 13:53:15 -0800
Subject: [PATCH] randgen: limit random vector dimensions to 50 in race builds

Previously, we would use 1000 as the maximum number of dimensions for
random vectors. However, we just saw a test failure under race where
vectorized cast between vector and collated string type took
prohibitively long, leading to a timeout, so this commit reduces the max
number of dimensions to 50 under race (significantly speeding up that
test).

Release note: None
---
 pkg/sql/randgen/datum.go       | 9 ++++++++-
 pkg/util/vector/vector.go      | 7 ++++---
 pkg/util/vector/vector_test.go | 2 +-
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/pkg/sql/randgen/datum.go b/pkg/sql/randgen/datum.go
index d445d0b3d482..ae47ed414cbc 100644
--- a/pkg/sql/randgen/datum.go
+++ b/pkg/sql/randgen/datum.go
@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgrepl/lsn"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/bitarray"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/ipaddr"
@@ -329,7 +330,13 @@ func RandDatumWithNullChance(
 	case types.TSQueryFamily:
 		return tree.NewDTSQuery(tsearch.RandomTSQuery(rng))
 	case types.PGVectorFamily:
-		return tree.NewDPGVector(vector.Random(rng))
+		var maxDim = 1000
+		if util.RaceEnabled {
+			// Some tests might be significantly slower under race, so we reduce
+			// the dimensionality.
+			maxDim = 50
+		}
+		return tree.NewDPGVector(vector.Random(rng, maxDim))
 	default:
 		panic(errors.AssertionFailedf("invalid type %v", typ.DebugString()))
 	}
diff --git a/pkg/util/vector/vector.go b/pkg/util/vector/vector.go
index f342a76fd3af..9880387a03f2 100644
--- a/pkg/util/vector/vector.go
+++ b/pkg/util/vector/vector.go
@@ -271,9 +271,10 @@ func Mult(t T, t2 T) (T, error) {
 	return ret, nil
 }
 
-// Random returns a random vector.
-func Random(rng *rand.Rand) T {
-	n := 1 + rng.Intn(1000)
+// Random returns a random vector with the number of dimensions in [1, maxDim]
+// range.
+func Random(rng *rand.Rand, maxDim int) T {
+	n := 1 + rng.Intn(maxDim)
 	v := make(T, n)
 	for i := range v {
 		for {
diff --git a/pkg/util/vector/vector_test.go b/pkg/util/vector/vector_test.go
index 5bd221cc2581..3c337d101aa0 100644
--- a/pkg/util/vector/vector_test.go
+++ b/pkg/util/vector/vector_test.go
@@ -66,7 +66,7 @@ func TestParseVector(t *testing.T) {
 func TestRoundtripRandomPGVector(t *testing.T) {
 	rng, _ := randutil.NewTestRand()
 	for i := 0; i < 1000; i++ {
-		v := Random(rng)
+		v := Random(rng, 1000 /* maxDim */)
 		encoded, err := Encode(nil, v)
 		assert.NoError(t, err)
 		roundtripped, err := Decode(encoded)