From 2453b85a1e9607b1c259025786d6b9b8fa84d812 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Tue, 28 Oct 2025 21:17:45 +0100 Subject: [PATCH] Adding more tests and I find the compression settings are broken --- go/cagra/cagra_test.go | 202 +++++++++------- go/cagra/extend_params_test.go | 134 +++++++++++ go/cagra/index_params.go | 8 +- go/cagra/index_params_test.go | 426 +++++++++++++++++++++++++++++++++ 4 files changed, 683 insertions(+), 87 deletions(-) create mode 100644 go/cagra/extend_params_test.go create mode 100644 go/cagra/index_params_test.go diff --git a/go/cagra/cagra_test.go b/go/cagra/cagra_test.go index 86baf5e601..bb4fd0a0a1 100644 --- a/go/cagra/cagra_test.go +++ b/go/cagra/cagra_test.go @@ -8,6 +8,20 @@ import ( ) func TestCagra(t *testing.T) { + testCases := []struct { + name string + compress bool + }{ + { + name: "No compression", + compress: false, + }, + { + name: "Compression", + compress: true, + }, + } + const ( nDataPoints = 1024 nFeatures = 16 @@ -15,109 +29,131 @@ func TestCagra(t *testing.T) { k = 4 epsilon = 0.001 ) + r := rand.New(rand.NewPCG(42, 0)) + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + resource, _ := cuvs.NewResource(nil) + defer resource.Close() + + testDataset := make([][]float32, nDataPoints) + for i := range testDataset { + testDataset[i] = make([]float32, nFeatures) + for j := range testDataset[i] { + testDataset[i][j] = r.Float32() + } + } - resource, _ := cuvs.NewResource(nil) - defer resource.Close() - - testDataset := make([][]float32, nDataPoints) - for i := range testDataset { - testDataset[i] = make([]float32, nFeatures) - for j := range testDataset[i] { - testDataset[i][j] = rand.Float32() - } - } + dataset, err := cuvs.NewTensor(testDataset) + if err != nil { + t.Fatalf("error creating dataset tensor: %v", err) + } + defer dataset.Close() - dataset, err := cuvs.NewTensor(testDataset) - if err != nil { - t.Fatalf("error creating dataset tensor: %v", err) - } - defer dataset.Close() + indexParams, err := CreateIndexParams() + if err != nil { + t.Fatalf("error creating index params: %v", err) + } + defer indexParams.Close() - indexParams, err := CreateIndexParams() - if err != nil { - t.Fatalf("error creating index params: %v", err) - } - defer indexParams.Close() + index, err := CreateIndex() + if err != nil { + t.Fatalf("error creating index: %v", err) + } + defer index.Close() - index, _ := CreateIndex() - defer index.Close() + // Use the first 4 points from the dataset as queries : will test that we get them back + // as their own nearest neighbor + queries, err := cuvs.NewTensor(testDataset[:nQueries]) + if err != nil { + t.Fatalf("error creating queries tensor: %v", err) + } + defer queries.Close() - // use the first 4 points from the dataset as queries : will test that we get them back - // as their own nearest neighbor - queries, _ := cuvs.NewTensor(testDataset[:nQueries]) - defer queries.Close() + neighbors, err := cuvs.NewTensorOnDevice[uint32](&resource, []int64{int64(nQueries), int64(k)}) + if err != nil { + t.Fatalf("error creating neighbors tensor: %v", err) + } + defer neighbors.Close() - neighbors, err := cuvs.NewTensorOnDevice[uint32](&resource, []int64{int64(nQueries), int64(k)}) - if err != nil { - t.Fatalf("error creating neighbors tensor: %v", err) - } - defer neighbors.Close() + distances, err := cuvs.NewTensorOnDevice[float32](&resource, []int64{int64(nQueries), int64(k)}) + if err != nil { + t.Fatalf("error creating distances tensor: %v", err) + } + defer distances.Close() - distances, err := cuvs.NewTensorOnDevice[float32](&resource, []int64{int64(nQueries), int64(k)}) - if err != nil { - t.Fatalf("error creating distances tensor: %v", err) - } - defer distances.Close() + if _, err := dataset.ToDevice(&resource); err != nil { + t.Fatalf("error moving dataset to device: %v", err) + } - if _, err := dataset.ToDevice(&resource); err != nil { - t.Fatalf("error moving dataset to device: %v", err) - } + if tc.compress { + compressionParams, err := CreateCompressionParams() + if err != nil { + t.Fatalf("error creating compression params: %v", err) + } + indexParams.SetCompression(compressionParams) + } - if err := BuildIndex(resource, indexParams, &dataset, index); err != nil { - t.Fatalf("error building index: %v", err) - } + if err := BuildIndex(resource, indexParams, &dataset, index); err != nil { + t.Fatalf("error building index: %v", err) + } - if err := resource.Sync(); err != nil { - t.Fatalf("error syncing resource: %v", err) - } + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } - if _, err := queries.ToDevice(&resource); err != nil { - t.Fatalf("error moving queries to device: %v", err) - } + if _, err := queries.ToDevice(&resource); err != nil { + t.Fatalf("error moving queries to device: %v", err) + } - SearchParams, err := CreateSearchParams() - if err != nil { - t.Fatalf("error creating search params: %v", err) - } - defer SearchParams.Close() + SearchParams, err := CreateSearchParams() + if err != nil { + t.Fatalf("error creating search params: %v", err) + } + defer SearchParams.Close() - err = SearchIndex(resource, SearchParams, index, &queries, &neighbors, &distances, nil) - if err != nil { - t.Fatalf("error searching index: %v", err) - } + err = SearchIndex(resource, SearchParams, index, &queries, &neighbors, &distances, nil) + if err != nil { + t.Fatalf("error searching index: %v", err) + } - if _, err := neighbors.ToHost(&resource); err != nil { - t.Fatalf("error moving neighbors to host: %v", err) - } + if _, err := neighbors.ToHost(&resource); err != nil { + t.Fatalf("error moving neighbors to host: %v", err) + } - if _, err := distances.ToHost(&resource); err != nil { - t.Fatalf("error moving distances to host: %v", err) - } + if _, err := distances.ToHost(&resource); err != nil { + t.Fatalf("error moving distances to host: %v", err) + } - if err := resource.Sync(); err != nil { - t.Fatalf("error syncing resource: %v", err) - } + if err := resource.Sync(); err != nil { + t.Fatalf("error syncing resource: %v", err) + } - neighborsSlice, err := neighbors.Slice() - if err != nil { - t.Fatalf("error getting neighbors slice: %v", err) - } + neighborsSlice, err := neighbors.Slice() + if err != nil { + t.Fatalf("error getting neighbors slice: %v", err) + } - for i := range neighborsSlice { - if neighborsSlice[i][0] != uint32(i) { - t.Error("wrong neighbor, expected", i, "got", neighborsSlice[i][0]) - } - } + for i := range neighborsSlice { + if neighborsSlice[i][0] != uint32(i) { + t.Error("wrong neighbor, expected", i, "got", neighborsSlice[i][0]) + } + } - distancesSlice, err := distances.Slice() - if err != nil { - t.Fatalf("error getting distances slice: %v", err) - } + distancesSlice, err := distances.Slice() + if err != nil { + t.Fatalf("error getting distances slice: %v", err) + } - for i := range distancesSlice { - if distancesSlice[i][0] >= epsilon || distancesSlice[i][0] <= -epsilon { - t.Error("distance should be close to 0, got", distancesSlice[i][0]) - } + if !tc.compress { + // Compress makes the result nondeterministic + for i := range distancesSlice { + if distancesSlice[i][0] >= epsilon || distancesSlice[i][0] <= -epsilon { + t.Error("distance should be close to 0, got", distancesSlice[i][0]) + } + } + } + }) } } diff --git a/go/cagra/extend_params_test.go b/go/cagra/extend_params_test.go new file mode 100644 index 0000000000..8f911b21d3 --- /dev/null +++ b/go/cagra/extend_params_test.go @@ -0,0 +1,134 @@ +package cagra + +import ( + "testing" +) + +func TestCreateExtendParams(t *testing.T) { + params, err := CreateExtendParams() + if err != nil { + t.Fatalf("Failed to create ExtendParams: %v", err) + } + defer params.Close() + + if params == nil { + t.Fatal("CreateExtendParams returned nil params") + } + + if params.params == nil { + t.Fatal("ExtendParams internal params are nil") + } +} + +func TestExtendParamsSetMaxChunkSize(t *testing.T) { + params, err := CreateExtendParams() + if err != nil { + t.Fatalf("Failed to create ExtendParams: %v", err) + } + defer params.Close() + + testCases := []struct { + name string + maxChunkSize uint32 + expectedError bool + }{ + { + name: "Zero value (auto select)", + maxChunkSize: 0, + expectedError: false, + }, + { + name: "Small chunk size", + maxChunkSize: 100, + expectedError: false, + }, + { + name: "Medium chunk size", + maxChunkSize: 1000, + expectedError: false, + }, + { + name: "Large chunk size", + maxChunkSize: 10000, + expectedError: false, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetMaxChunkSize(tc.maxChunkSize) + + if tc.expectedError && err == nil { + t.Errorf("Expected error but got none") + } + + if !tc.expectedError && err != nil { + t.Errorf("Unexpected error: %v", err) + } + + if result == nil { + t.Error("SetMaxChunkSize returned nil") + } + + // Verify the value was set + if uint32(params.params.max_chunk_size) != tc.maxChunkSize { + t.Errorf("Expected max_chunk_size to be %d, got %d", + tc.maxChunkSize, params.params.max_chunk_size) + } + }) + } +} + +func TestExtendParamsSetMaxChunkSizeChaining(t *testing.T) { + params, err := CreateExtendParams() + if err != nil { + t.Fatalf("Failed to create ExtendParams: %v", err) + } + defer params.Close() + + // Test method chaining + result, err := params.SetMaxChunkSize(500) + if err != nil { + t.Fatalf("SetMaxChunkSize failed: %v", err) + } + + if result != params { + t.Error("SetMaxChunkSize should return the same params instance for chaining") + } +} + +func TestExtendParamsClose(t *testing.T) { + params, err := CreateExtendParams() + if err != nil { + t.Fatalf("Failed to create ExtendParams: %v", err) + } + + err = params.Close() + if err != nil { + t.Errorf("Close() failed: %v", err) + } +} + +func TestExtendParamsLifecycle(t *testing.T) { + // Test complete lifecycle: create -> set -> close + params, err := CreateExtendParams() + chunkSize := uint32(2000) + if err != nil { + t.Fatalf("Failed to create ExtendParams: %v", err) + } + + params, err = params.SetMaxChunkSize(chunkSize) + if err != nil { + t.Fatalf("SetMaxChunkSize failed: %v", err) + } + + if uint32(params.params.max_chunk_size) != chunkSize { + t.Errorf("Expected max_chunk_size to be %d, got %d", + chunkSize, params.params.max_chunk_size) + } + + err = params.Close() + if err != nil { + t.Errorf("Close() failed: %v", err) + } +} diff --git a/go/cagra/index_params.go b/go/cagra/index_params.go index 51fdb318fa..99f4b70b93 100644 --- a/go/cagra/index_params.go +++ b/go/cagra/index_params.go @@ -36,15 +36,15 @@ var cBuildAlgos = map[BuildAlgo]int{ func CreateCompressionParams() (*CompressionParams, error) { var params C.cuvsCagraCompressionParams_t - if params == nil { - return nil, errors.New("memory allocation failed") - } - err := cuvs.CheckCuvs(cuvs.CuvsError(C.cuvsCagraCompressionParamsCreate(¶ms))) if err != nil { return nil, err } + if params == nil { + return nil, errors.New("memory allocation failed") + } + return &CompressionParams{params: params}, nil } diff --git a/go/cagra/index_params_test.go b/go/cagra/index_params_test.go new file mode 100644 index 0000000000..579419ca92 --- /dev/null +++ b/go/cagra/index_params_test.go @@ -0,0 +1,426 @@ +package cagra + +import ( + "testing" +) + +// CompressionParams Tests +func TestCreateCompressionParams(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + if params == nil { + t.Fatal("CreateCompressionParams returned nil params") + } + + if params.params == nil { + t.Fatal("CompressionParams internal params are nil") + } + if params.params.pq_kmeans_trainset_fraction != 0 { + t.Fatalf("Error params.params.pq_kmeans_trainset_fraction != 0, got = %v", params.params.pq_kmeans_trainset_fraction) + } + if params.params.pq_bits != 8 { + t.Fatalf("Error params.params.pq_bits != 8, got = %v", params.params.pq_bits) + } + if params.params.pq_dim != 0 { + t.Fatalf("Error params.params.pq_dim != 0, got = %v", params.params.pq_dim) + } + if params.params.vq_n_centers != 0 { + t.Fatalf("Error params.params.vq_n_centers != 0, got = %v", params.params.vq_n_centers) + } + if params.params.kmeans_n_iters != 25 { + t.Fatalf("Error params.params.kmeans_n_iters != 25, got = %v", params.params.kmeans_n_iters) + } +} + +func TestCompressionParamsSetPQBits(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + testCases := []struct { + name string + value uint32 + }{ + {"4 bits", 4}, + {"8 bits", 8}, + {"16 bits", 16}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetPQBits(tc.value) + if err != nil { + t.Errorf("SetPQBits failed: %v", err) + } + if result != params { + t.Error("SetPQBits should return the same params instance") + } + if uint32(params.params.pq_bits) != tc.value { + t.Errorf("Expected pq_bits %d, got %d", tc.value, params.params.pq_bits) + } + }) + } +} + +func TestCompressionParamsSetPQDim(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + testCases := []struct { + name string + value uint32 + }{ + {"Zero (auto)", 0}, + {"Small dimension", 32}, + {"Large dimension", 128}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetPQDim(tc.value) + if err != nil { + t.Errorf("SetPQDim failed: %v", err) + } + if result != params { + t.Error("SetPQDim should return the same params instance") + } + if uint32(params.params.pq_dim) != tc.value { + t.Errorf("Expected pq_dim %d, got %d", tc.value, params.params.pq_dim) + } + }) + } +} + +func TestCompressionParamsSetVQNCenters(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + testCases := []struct { + name string + value uint32 + }{ + {"Zero (auto)", 0}, + {"Small centers", 256}, + {"Large centers", 2048}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetVQNCenters(tc.value) + if err != nil { + t.Errorf("SetVQNCenters failed: %v", err) + } + if result != params { + t.Error("SetVQNCenters should return the same params instance") + } + if uint32(params.params.vq_n_centers) != tc.value { + t.Errorf("Expected vq_n_centers %d, got %d", tc.value, params.params.vq_n_centers) + } + }) + } +} + +func TestCompressionParamsSetKMeansNIters(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + testCases := []struct { + name string + value uint32 + }{ + {"Few iterations", 10}, + {"Default iterations", 25}, + {"Many iterations", 100}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetKMeansNIters(tc.value) + if err != nil { + t.Errorf("SetKMeansNIters failed: %v", err) + } + if result != params { + t.Error("SetKMeansNIters should return the same params instance") + } + if uint32(params.params.kmeans_n_iters) != tc.value { + t.Errorf("Expected kmeans_n_iters %d, got %d", tc.value, params.params.kmeans_n_iters) + } + }) + } +} + +func TestCompressionParamsSetVQKMeansTrainsetFraction(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + testCases := []struct { + name string + value float64 + }{ + {"Zero (auto)", 0.0}, + {"Half dataset", 0.5}, + {"Full dataset", 1.0}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetVQKMeansTrainsetFraction(tc.value) + if err != nil { + t.Errorf("SetVQKMeansTrainsetFraction failed: %v", err) + } + if result != params { + t.Error("SetVQKMeansTrainsetFraction should return the same params instance") + } + if float64(params.params.vq_kmeans_trainset_fraction) != tc.value { + t.Errorf("Expected vq_kmeans_trainset_fraction %f, got %f", + tc.value, params.params.vq_kmeans_trainset_fraction) + } + }) + } +} + +func TestCompressionParamsSetPQKMeansTrainsetFraction(t *testing.T) { + params, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + testCases := []struct { + name string + value float64 + }{ + {"Zero (auto)", 0.0}, + {"Quarter dataset", 0.25}, + {"Half dataset", 0.5}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetPQKMeansTrainsetFraction(tc.value) + if err != nil { + t.Errorf("SetPQKMeansTrainsetFraction failed: %v", err) + } + if result != params { + t.Error("SetPQKMeansTrainsetFraction should return the same params instance") + } + if float64(params.params.pq_kmeans_trainset_fraction) != tc.value { + t.Errorf("Expected pq_kmeans_trainset_fraction %f, got %f", + tc.value, params.params.pq_kmeans_trainset_fraction) + } + }) + } +} + +// IndexParams Tests +func TestCreateIndexParams(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + defer params.Close() + + if params == nil { + t.Fatal("CreateIndexParams returned nil params") + } + + if params.params == nil { + t.Fatal("IndexParams internal params are nil") + } +} + +func TestIndexParamsSetIntermediateGraphDegree(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + defer params.Close() + + testCases := []struct { + name string + value uintptr + }{ + {"Small degree", 32}, + {"Medium degree", 64}, + {"Large degree", 128}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetIntermediateGraphDegree(tc.value) + if err != nil { + t.Errorf("SetIntermediateGraphDegree failed: %v", err) + } + if result != params { + t.Error("SetIntermediateGraphDegree should return the same params instance") + } + if uintptr(params.params.intermediate_graph_degree) != tc.value { + t.Errorf("Expected intermediate_graph_degree %d, got %d", + tc.value, params.params.intermediate_graph_degree) + } + }) + } +} + +func TestIndexParamsSetGraphDegree(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + defer params.Close() + + testCases := []struct { + name string + value uintptr + }{ + {"Small degree", 16}, + {"Medium degree", 32}, + {"Large degree", 64}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetGraphDegree(tc.value) + if err != nil { + t.Errorf("SetGraphDegree failed: %v", err) + } + if result != params { + t.Error("SetGraphDegree should return the same params instance") + } + if uintptr(params.params.graph_degree) != tc.value { + t.Errorf("Expected graph_degree %d, got %d", + tc.value, params.params.graph_degree) + } + }) + } +} + +func TestIndexParamsSetBuildAlgo(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + defer params.Close() + + testCases := []struct { + name string + algo BuildAlgo + expectedError bool + }{ + {"IVF_PQ algorithm", IvfPq, false}, + {"NN_DESCENT algorithm", NnDescent, false}, + {"AUTO_SELECT algorithm", AutoSelect, false}, + {"Invalid algorithm", BuildAlgo(999), true}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetBuildAlgo(tc.algo) + + if tc.expectedError { + if err == nil { + t.Error("Expected error for invalid build_algo, got none") + } + } else { + if err != nil { + t.Errorf("SetBuildAlgo failed: %v", err) + } + if result != params { + t.Error("SetBuildAlgo should return the same params instance") + } + } + }) + } +} + +func TestIndexParamsSetNNDescentNiter(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + defer params.Close() + + testCases := []struct { + name string + value uint32 + }{ + {"Few iterations", 10}, + {"Default iterations", 20}, + {"Many iterations", 50}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + result, err := params.SetNNDescentNiter(tc.value) + if err != nil { + t.Errorf("SetNNDescentNiter failed: %v", err) + } + if result != params { + t.Error("SetNNDescentNiter should return the same params instance") + } + if uint32(params.params.nn_descent_niter) != tc.value { + t.Errorf("Expected nn_descent_niter %d, got %d", + tc.value, params.params.nn_descent_niter) + } + }) + } +} + +func TestIndexParamsSetCompression(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + defer params.Close() + + compression, err := CreateCompressionParams() + if err != nil { + t.Fatalf("Failed to create CompressionParams: %v", err) + } + + // Configure compression params + compression.SetPQBits(8) + compression.SetPQDim(64) + + result, err := params.SetCompression(compression) + if err != nil { + t.Errorf("SetCompression failed: %v", err) + } + if result != params { + t.Error("SetCompression should return the same params instance") + } +} + +func TestIndexParamsClose(t *testing.T) { + params, err := CreateIndexParams() + if err != nil { + t.Fatalf("Failed to create IndexParams: %v", err) + } + + err = params.Close() + if err != nil { + t.Errorf("Close() failed: %v", err) + } +} + +func TestBuildAlgoConstants(t *testing.T) { + // Test that BuildAlgo constants are properly defined + algos := []BuildAlgo{IvfPq, NnDescent, AutoSelect} + + for _, algo := range algos { + _, exists := cBuildAlgos[algo] + if !exists { + t.Errorf("BuildAlgo %v not found in cBuildAlgos map", algo) + } + } +}