Method copyVectorDataToShards

dgraph/cmd/bulk/reduce.go:507–549

copyVectorDataToShards copies vector data from the shared vectorTmpDb to the correct output DBs. It uses the predToOutputShard map to determine which predicate goes to which shard.

(vectorDb *badger.DB, predToShard map[string]int)

// copyVectorDataToShards copies vector data from the shared vectorTmpDb to the correct output DBs.
// It uses the predToOutputShard map to determine which predicate goes to which shard.
func (r *reducer) copyVectorDataToShards(vectorDb *badger.DB, predToShard map[string]int) {
	if len(predToShard) == 0 {
		return
	}

	// Group predicates by output shard
	shardPreds := make(map[int][]string)
	for pred, shardId := range predToShard {
		shardPreds[shardId] = append(shardPreds[shardId], pred)
	}

	// Copy each shard's predicates to its output DB
	for shardId, preds := range shardPreds {
		if shardId >= len(r.dbs) {
			glog.Errorf("Invalid shard ID %d for predicates %v (only %d DBs)", shardId, preds, len(r.dbs))
			continue
		}

		destDb := r.dbs[shardId]

		// Collect all predicate keys (base + HNSW suffixes) and sort for efficient iteration
		var allPreds []string
		for _, pred := range preds {
			allPreds = append(allPreds, pred)
			allPreds = append(allPreds, hnsw.ConcatStrings(pred, hnsw.VecKeyword)) // __vector_
			allPreds = append(allPreds, hnsw.ConcatStrings(pred, hnsw.VecEntry))   // __vector_entry
			allPreds = append(allPreds, hnsw.ConcatStrings(pred, hnsw.VecDead))    // __vector_dead
		}
		sort.Strings(allPreds)

		fmt.Printf("Copying %d vector predicates to shard %d\n", len(preds), shardId)

		// Copy each predicate
		totalCount := 0
		for _, pred := range allPreds {
			count := r.copyPredicateWithBatchCount(destDb, vectorDb, pred)
			totalCount += count
		}

		fmt.Printf("Copied %d total KV entries for %d predicates to shard %d\n",
			totalCount, len(preds), shardId)
	}
}

// copyPredicateWithBatchCount is like copyPredicateWithBatch but returns the count of entries copied.
func (r *reducer) copyPredicateWithBatchCount(destDb *badger.DB, srcDb *badger.DB, pred string) int {