Function TestBulkLoadVectorIndexMultipleGroups

systest/vector/load_test.go:206–350

(t *testing.T)

}

func TestBulkLoadVectorIndexMultipleGroups(t *testing.T) {
	// Define 3 different vector predicates - each will potentially go to different shards
	predicates := []string{"vec_pred_alpha", "vec_pred_beta", "vec_pred_gamma"}
	numVectorsPerPred := 1000
	vectorDim := 10
	numShards := 3

	// Schema with 3 vector predicates
	multiPredSchema := `
		vec_pred_alpha: float32vector @index(hnsw(exponent: "5", metric: "euclidean")) .
		vec_pred_beta: float32vector @index(hnsw(exponent: "5", metric: "euclidean")) .
		vec_pred_gamma: float32vector @index(hnsw(exponent: "5", metric: "euclidean")) .
	`

	// Step 1: Create a source cluster and load vectors into it
	sourceConf := dgraphtest.NewClusterConfig().
		WithNumAlphas(1).
		WithNumZeros(1).
		WithReplicas(1).
		WithACL(time.Hour)
	sourceCluster, err := dgraphtest.NewLocalCluster(sourceConf)
	require.NoError(t, err)
	defer func() { sourceCluster.Cleanup(t.Failed()) }()
	require.NoError(t, sourceCluster.Start())

	gc, cleanup, err := sourceCluster.Client()
	require.NoError(t, err)
	defer cleanup()
	require.NoError(t, gc.LoginIntoNamespace(context.Background(),
		dgraphapi.DefaultUser, dgraphapi.DefaultPassword, x.RootNamespace))

	hc, err := sourceCluster.HTTPClient()
	require.NoError(t, err)
	require.NoError(t, hc.LoginIntoNamespace(dgraphapi.DefaultUser,
		dgraphapi.DefaultPassword, x.RootNamespace))

	// Set up schema with multiple vector predicates
	require.NoError(t, gc.SetupSchema(multiPredSchema))

	// Generate and load vectors for each predicate
	allVectors := make(map[string][][]float32)
	for _, pred := range predicates {
		rdfs, vectors := dgraphapi.GenerateRandomVectors(0, numVectorsPerPred, vectorDim, pred)
		allVectors[pred] = vectors

		mu := &api.Mutation{SetNquads: []byte(rdfs), CommitNow: true}
		_, err = gc.Mutate(mu)
		require.NoError(t, err)
	}

	// Verify vectors are loaded and queryable in source cluster
	for _, pred := range predicates {
		vectors := allVectors[pred]
		similarVectors, err := gc.QueryMultipleVectorsUsingSimilarTo(vectors[0], pred, 5)
		require.NoError(t, err)
		require.GreaterOrEqual(t, len(similarVectors), 3, "similar_to query should return results for %s", pred)
	}