Use the computed TS from the map phase

Allows no zero in the reduce phase

Author: matthewmcneely

dgraph/cmd/bulk/loader.go

@@ -145,7 +145,7 @@ type loader struct {
 	dg      *dgo.Dgraph
 }
 
-func newLoader(opt *BulkOptions) *loader {
+func newLoader(opt *BulkOptions, precomputedWriteTs uint64) *loader {
 	if opt == nil {
 		log.Fatalf("Cannot create loader with nil options.")
 	}
@@ -186,17 +186,27 @@ func newLoader(opt *BulkOptions) *loader {
 		fmt.Printf("Error logging enabled, writing to: %s\n", opt.ErrorLogPath)
 	}
 
+	writeTs := precomputedWriteTs
+	if writeTs == 0 {
+		writeTs = getWriteTimestamp(zero, dg)
+	}
 	st := &state{
 		opt:    opt,
 		prog:   newProgress(),
 		shards: newShardMap(opt.MapShards),
 		// Lots of gz readers, so not much channel buffer needed.
 		readerChunkCh: make(chan *chunkWithMeta, opt.NumGoroutines),
-		writeTs:       getWriteTimestamp(zero, dg),
+		writeTs:       writeTs,
 		namespaces:    &sync.Map{},
 		errorLog:      errLog,
 	}
-	st.schema = newSchemaStore(readSchema(opt), opt, st)
+	var parsedSchema *schema.ParsedSchema
+	if !opt.SkipMapPhase {
+		parsedSchema = readSchema(opt)
+	} else {
+		parsedSchema = &schema.ParsedSchema{}
+	}
+	st.schema = newSchemaStore(parsedSchema, opt, st)
 	ld := &loader{
 		state:   st,
 		mappers: make([]*mapper, opt.NumGoroutines),

dgraph/cmd/bulk/run.go

@@ -208,27 +208,29 @@ func RunBulkLoader(opt BulkOptions) {
 	}
 	fmt.Printf("Encrypted input: %v; Encrypted output: %v\n", opt.Encrypted, opt.EncryptedOut)
 
-	if opt.SchemaFile == "" {
-		// if only graphql schema is provided, we can generate DQL schema from it.
-		if opt.GqlSchemaFile == "" {
-			fmt.Fprint(os.Stderr, "Schema file must be specified.\n")
-			os.Exit(1)
+	if !opt.SkipMapPhase {
+		if opt.SchemaFile == "" {
+			// if only graphql schema is provided, we can generate DQL schema from it.
+			if opt.GqlSchemaFile == "" {
+				fmt.Fprint(os.Stderr, "Schema file must be specified.\n")
+				os.Exit(1)
+			}
+		} else {
+			if !filestore.Exists(opt.SchemaFile) {
+				fmt.Fprintf(os.Stderr, "Schema path(%v) does not exist.\n", opt.SchemaFile)
+				os.Exit(1)
+			}
 		}
-	} else {
-		if !filestore.Exists(opt.SchemaFile) {
-			fmt.Fprintf(os.Stderr, "Schema path(%v) does not exist.\n", opt.SchemaFile)
+		if opt.DataFiles == "" {
+			fmt.Fprint(os.Stderr, "RDF or JSON file(s) location must be specified.\n")
 			os.Exit(1)
-		}
-	}
-	if opt.DataFiles == "" {
-		fmt.Fprint(os.Stderr, "RDF or JSON file(s) location must be specified.\n")
-		os.Exit(1)
-	} else {
-		fileList := strings.SplitSeq(opt.DataFiles, ",")
-		for file := range fileList {
-			if !filestore.Exists(file) {
-				fmt.Fprintf(os.Stderr, "Data path(%v) does not exist.\n", file)
-				os.Exit(1)
+		} else {
+			fileList := strings.SplitSeq(opt.DataFiles, ",")
+			for file := range fileList {
+				if !filestore.Exists(file) {
+					fmt.Fprintf(os.Stderr, "Data path(%v) does not exist.\n", file)
+					os.Exit(1)
+				}
 			}
 		}
 	}
@@ -323,10 +325,26 @@ func RunBulkLoader(opt BulkOptions) {
 	x.Check(os.MkdirAll(bufDir, 0700))
 	defer os.RemoveAll(bufDir)
 
-	loader := newLoader(&opt)
-
 	const bulkMetaFilename = "bulk.meta"
 	bulkMetaPath := filepath.Join(opt.TmpDir, bulkMetaFilename)
+	const writeTsFilename = "write.ts"
+	writeTsPath := filepath.Join(opt.TmpDir, writeTsFilename)
+
+	var precomputedWriteTs uint64
+	if opt.SkipMapPhase {
+		writeTsData, err := os.ReadFile(writeTsPath)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "Error reading write timestamp file; was --skip_reduce_phase used in the map run?")
+			os.Exit(1)
+		}
+		precomputedWriteTs, err = strconv.ParseUint(strings.TrimSpace(string(writeTsData)), 10, 64)
+		if err != nil {
+			fmt.Fprintln(os.Stderr, "Error parsing write timestamp file")
+			os.Exit(1)
+		}
+	}
+
+	loader := newLoader(&opt, precomputedWriteTs)
 
 	if opt.SkipMapPhase {
 		bulkMetaData, err := os.ReadFile(bulkMetaPath)
@@ -363,6 +381,10 @@ func RunBulkLoader(opt BulkOptions) {
 			fmt.Fprintln(os.Stderr, "Error writing to bulk meta file")
 			os.Exit(1)
 		}
+		if err = os.WriteFile(writeTsPath, []byte(strconv.FormatUint(loader.writeTs, 10)), 0600); err != nil {
+			fmt.Fprintln(os.Stderr, "Error writing write timestamp file")
+			os.Exit(1)
+		}
 	}
 
 	if opt.SkipReducePhase {

systest/integration2/bulk_loader_test.go

@@ -89,12 +89,11 @@ func TestBulkLoaderSkipReducePhase(t *testing.T) {
 	}
 	require.NoError(t, c.BulkLoad(mapOpts))
 
-	// Second run: reduce phase only, using the same tmp dir
+	// Second run: reduce phase only, using the same tmp dir.
+	// Data and schema files are not needed; all input was processed in the map phase.
 	reduceOpts := dgraphtest.BulkOpts{
-		DataFiles:      []string{dataFile},
-		GQLSchemaFiles: []string{gqlSchemaFile},
-		TmpDir:         tmpDir,
-		SkipMapPhase:   true,
+		TmpDir:       tmpDir,
+		SkipMapPhase: true,
 	}
 	require.NoError(t, c.BulkLoad(reduceOpts))