Bulk loader allocates reserved predicates in first reduce shard.

dgraph/cmd/bulk/merge_shards.go

@@ -32,7 +32,13 @@ const (
 )
 
 func mergeMapShardsIntoReduceShards(opt options) {
-	mapShards := shardDirs(filepath.Join(opt.TmpDir, mapShardDir))
+	shardDirs := readShardDirs(filepath.Join(opt.TmpDir, mapShardDir))
+	// First shard is handled differently because it contains reserved predicates.
+	x.AssertTrue(len(shardDirs) > 0)
+	firstShard := shardDirs[0]
+	// Sort the rest of the shards by size to allow the largest shards to be shuffled first.
+	shardDirs = shardDirs[1:]
+	sortBySize(shardDirs)
 
 	var reduceShards []string
 	for i := 0; i < opt.ReduceShards; i++ {
@@ -41,9 +47,15 @@ func mergeMapShardsIntoReduceShards(opt options) {
 		reduceShards = append(reduceShards, shardDir)
 	}
 
+	// Put the first map shard in the first reduce shard since it contains all the reserved
+	// predicates.
+	reduceShard := filepath.Join(reduceShards[0], filepath.Base(firstShard))
+	fmt.Printf("Shard %s -> Reduce %s\n", firstShard, reduceShard)
+	x.Check(os.Rename(firstShard, reduceShard))
+
 	// Heuristic: put the largest map shard into the smallest reduce shard
 	// until there are no more map shards left. Should be a good approximation.
-	for _, shard := range mapShards {
+	for _, shard := range shardDirs {
 		sortBySize(reduceShards)
 		reduceShard := filepath.Join(
 			reduceShards[len(reduceShards)-1], filepath.Base(shard))
@@ -52,7 +64,7 @@ func mergeMapShardsIntoReduceShards(opt options) {
 	}
 }
 
-func shardDirs(d string) []string {
+func readShardDirs(d string) []string {
 	dir, err := os.Open(d)
 	x.Check(err)
 	shards, err := dir.Readdirnames(0)
@@ -61,9 +73,7 @@ func shardDirs(d string) []string {
 	for i, shard := range shards {
 		shards[i] = filepath.Join(d, shard)
 	}
-
-	// Allow largest shards to be shuffled first.
-	sortBySize(shards)
+	sort.Strings(shards)
 	return shards
 }
 

dgraph/cmd/bulk/reduce.go

@@ -46,7 +46,7 @@ type reducer struct {
 }
 
 func (r *reducer) run() error {
-	dirs := shardDirs(filepath.Join(r.opt.TmpDir, reduceShardDir))
+	dirs := readShardDirs(filepath.Join(r.opt.TmpDir, reduceShardDir))
 	x.AssertTrue(len(dirs) == r.opt.ReduceShards)
 	x.AssertTrue(len(r.opt.shardOutputDirs) == r.opt.ReduceShards)
 

dgraph/cmd/bulk/run.go

@@ -38,6 +38,7 @@ import (
 var Bulk x.SubCommand
 
 var defaultOutDir = "./out"
+var groupFile = "group_id"
 
 func init() {
 	Bulk.Cmd = &cobra.Command{
@@ -192,6 +193,13 @@ func run() {
 		dir := filepath.Join(opt.OutDir, strconv.Itoa(i), "p")
 		x.Check(os.MkdirAll(dir, 0700))
 		opt.shardOutputDirs = append(opt.shardOutputDirs, dir)
+
+		groupFile := filepath.Join(dir, groupFile)
+		f, err := os.OpenFile(groupFile, os.O_CREATE|os.O_WRONLY, 0600)
+		x.Check(err)
+		x.Check2(f.WriteString(strconv.Itoa(i + 1)))
+		x.Check2(f.WriteString("\n"))
+		x.Check(f.Close())
 	}
 
 	// Create a directory just for bulk loader's usage.

dgraph/cmd/bulk/shard_map.go

@@ -16,7 +16,11 @@
 
 package bulk
 
-import "sync"
+import (
+	"sync"
+
+	"github.com/dgraph-io/dgraph/x"
+)
 
 type shardMap struct {
 	sync.RWMutex
@@ -33,6 +37,11 @@ func newShardMap(numShards int) *shardMap {
 }
 
 func (m *shardMap) shardFor(pred string) int {
+	// Always assign NQuads with reserved predicates to the first map shard.
+	if x.IsReservedPredicate(pred) {
+		return 0
+	}
+
 	m.RLock()
 	shard, ok := m.predToShard[pred]
 	m.RUnlock()

dgraph/cmd/zero/zero.go

@@ -512,6 +512,12 @@ func (s *Server) Connect(ctx context.Context,
 				// We need more servers here, so let's add it.
 				proposal.Member = m
 				return proposal
+			} else if m.ForceGroupId {
+				// If the group ID was taken from the group_id file, force the member
+				// to be in this group even if the group is at capacity. This should
+				// not happen if users properly initialize a cluster after a bulk load.
+				proposal.Member = m
+				return proposal
 			}
 			// Already have plenty of servers serving this group.
 		}

edgraph/server.go

@@ -19,9 +19,11 @@ package edgraph
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"math"
 	"math/rand"
 	"os"
+	"path/filepath"
 	"sort"
 	"strconv"
 	"strings"
@@ -60,6 +62,7 @@ import (
 const (
 	methodMutate = "Server.Mutate"
 	methodQuery  = "Server.Query"
+	groupFile    = "group_id"
 )
 
 // ServerState holds the state of the Dgraph server.
@@ -96,8 +99,21 @@ func InitServerState() {
 	State.needTs = make(chan tsReq, 100)
 
 	State.initStorage()
-
 	go State.fillTimestampRequests()
+
+	contents, err := ioutil.ReadFile(filepath.Join(Config.PostingDir, groupFile))
+	if err != nil {
+		return
+	}
+
+	glog.Infof("Found group_id file inside posting directory %s. Will attempt to read.",
+		Config.PostingDir)
+	groupId, err := strconv.ParseUint(strings.TrimSpace(string(contents)), 0, 32)
+	if err != nil {
+		glog.Warningf("Could not read %s file inside posting directory %s.",
+			groupFile, Config.PostingDir)
+	}
+	x.WorkerConfig.ProposedGroupId = uint32(groupId)
 }
 
 func (s *ServerState) runVlogGC(store *badger.DB) {

protos/pb.proto

@@ -130,6 +130,7 @@ message Member {
 	uint64 last_update = 6;
 
 	bool cluster_info_only = 13;
+	bool force_group_id = 14;
 }
 
 message Group {