github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/vector/compressionhelpers/kmeans_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  //go:build !race
    13  
    14  package compressionhelpers_test
    15  
    16  import (
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/stretchr/testify/assert"
    21  	"github.com/weaviate/weaviate/adapters/repos/db/vector/compressionhelpers"
    22  	"github.com/weaviate/weaviate/adapters/repos/db/vector/hnsw/distancer"
    23  	testinghelpers "github.com/weaviate/weaviate/adapters/repos/db/vector/testinghelpers"
    24  )
    25  
    26  func Test_NoRaceKMeansNNearest(t *testing.T) {
    27  	distanceProvider := distancer.NewL2SquaredProvider()
    28  	vectors := [][]float32{
    29  		{0, 5},
    30  		{0.1, 4.9},
    31  		{0.01, 5.1},
    32  		{10.1, 7},
    33  		{5.1, 2},
    34  		{5.0, 2.1},
    35  	}
    36  	kmeans := compressionhelpers.NewKMeans(
    37  		3,
    38  		2,
    39  		0,
    40  	)
    41  	kmeans.Fit(vectors)
    42  	centers := make([]byte, 6)
    43  	for i := range centers {
    44  		centers[i] = byte(kmeans.Nearest(vectors[i]))
    45  	}
    46  	for v := range vectors {
    47  		min, _, _ := distanceProvider.SingleDist(vectors[v], kmeans.Centroid(centers[v]))
    48  		for c := range centers {
    49  			dist, _, _ := distanceProvider.SingleDist(vectors[v], kmeans.Centroid(centers[c]))
    50  			assert.True(t, dist >= min)
    51  		}
    52  	}
    53  }
    54  
    55  func Test_NoRaceRandomData(t *testing.T) {
    56  	vectorsSize := 10000
    57  	vectors, _ := testinghelpers.RandomVecs(vectorsSize, 0, 128)
    58  	before := time.Now()
    59  	kmeans := compressionhelpers.NewKMeans(
    60  		256,
    61  		1,
    62  		10,
    63  	)
    64  	kmeans.Fit(vectors)
    65  	assert.True(t, time.Since(before).Seconds() < 50)
    66  }