Merge cockroachdb#136163

136163: randgen: limit random vector dimensions to 50 in race builds r=yuzefovich a=yuzefovich 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). Fixes: cockroachdb#136086. Release note: None Co-authored-by: Yahor Yuzefovich <[email protected]>
sambhav-jain-16 · Nov 26, 2024 · 67caf19 · 67caf19
2 parents 93371b5 + 10f2a47
commit 67caf19
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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)