Adding more tests and I find the compression settings are broken in the Golang API

go/cagra/cagra_test.go

@@ -8,116 +8,152 @@ 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
 		nQueries    = 4
 		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])
+					}
+				}
+			}
+		})
 	}
 }
 

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)
+	}
+}

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(&params)))
 	if err != nil {
 		return nil, err
 	}
 
+	if params == nil {
+		return nil, errors.New("memory allocation failed")
+	}
+
 	return &CompressionParams{params: params}, nil
 }