github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/vector/hnsw/compress_integration_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 integrationTest
    13  // +build integrationTest
    14  
    15  package hnsw
    16  
    17  import (
    18  	"context"
    19  	"testing"
    20  
    21  	"github.com/stretchr/testify/assert"
    22  	"github.com/stretchr/testify/require"
    23  	"github.com/weaviate/weaviate/adapters/repos/db/vector/common"
    24  	"github.com/weaviate/weaviate/adapters/repos/db/vector/compressionhelpers"
    25  	"github.com/weaviate/weaviate/adapters/repos/db/vector/hnsw/distancer"
    26  	"github.com/weaviate/weaviate/adapters/repos/db/vector/testinghelpers"
    27  	"github.com/weaviate/weaviate/entities/cyclemanager"
    28  	ent "github.com/weaviate/weaviate/entities/vectorindex/hnsw"
    29  )
    30  
    31  func Test_NoRaceCompressAdaptsSegments(t *testing.T) {
    32  	ctx := context.Background()
    33  
    34  	efConstruction := 64
    35  	ef := 32
    36  	maxNeighbors := 32
    37  	vectorsCount := 1000
    38  	dimensions := 6
    39  	expectedSegments := 3
    40  
    41  	vectors, _ := testinghelpers.RandomVecs(vectorsCount, 0, dimensions)
    42  	distancer := distancer.NewL2SquaredProvider()
    43  
    44  	uc := ent.UserConfig{}
    45  	uc.MaxConnections = maxNeighbors
    46  	uc.EFConstruction = efConstruction
    47  	uc.EF = ef
    48  	uc.VectorCacheMaxObjects = 10e12
    49  	uc.PQ = ent.PQConfig{
    50  		Enabled: true,
    51  		Encoder: ent.PQEncoder{
    52  			Type:         ent.PQEncoderTypeKMeans,
    53  			Distribution: ent.PQEncoderDistributionNormal,
    54  		},
    55  	}
    56  
    57  	store := testinghelpers.NewDummyStore(t)
    58  	defer func() {
    59  		err := store.Shutdown(ctx)
    60  		require.NoError(t, err)
    61  	}()
    62  
    63  	index, err := New(
    64  		Config{
    65  			RootPath:              t.TempDir(),
    66  			ID:                    "recallbenchmark",
    67  			MakeCommitLoggerThunk: MakeNoopCommitLogger,
    68  			DistanceProvider:      distancer,
    69  			VectorForIDThunk: func(ctx context.Context, id uint64) ([]float32, error) {
    70  				return vectors[int(id)], nil
    71  			},
    72  			TempVectorForIDThunk: func(ctx context.Context, id uint64, container *common.VectorSlice) ([]float32, error) {
    73  				copy(container.Slice, vectors[int(id)])
    74  				return container.Slice, nil
    75  			},
    76  		}, uc,
    77  		cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), cyclemanager.NewCallbackGroupNoop(), store)
    78  	require.NoError(t, err)
    79  	defer func() {
    80  		err := index.Shutdown(ctx)
    81  		require.NoError(t, err)
    82  	}()
    83  
    84  	compressionhelpers.Concurrently(uint64(len(vectors)), func(id uint64) {
    85  		index.Add(uint64(id), vectors[id])
    86  	})
    87  	uc.PQ = ent.PQConfig{
    88  		Enabled: true,
    89  		Encoder: ent.PQEncoder{
    90  			Type:         ent.PQEncoderTypeKMeans,
    91  			Distribution: ent.PQEncoderDistributionNormal,
    92  		},
    93  		Segments:  0,
    94  		Centroids: 256,
    95  	}
    96  	index.compress(uc)
    97  	assert.Equal(t, expectedSegments, int(index.compressor.ExposeFields().M))
    98  	assert.Equal(t, expectedSegments, index.pqConfig.Segments)
    99  }