Method copyPredicateWithBatchCount

dgraph/cmd/bulk/reduce.go:552–599

copyPredicateWithBatchCount is like copyPredicateWithBatch but returns the count of entries copied.

(destDb *badger.DB, srcDb *badger.DB, pred string)

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

	stream := srcDb.NewStreamAt(math.MaxUint64)
	stream.LogPrefix = fmt.Sprintf("copying predicate %s", pred)
	stream.Prefix = x.PredicatePrefix(pred)
	stream.NumGo = 1

	var count int
	stream.Send = func(buf *z.Buffer) error {
		kvs, err := badger.BufferToKVList(buf)
		if err != nil {
			return err
		}

		for _, kv := range kvs.Kv {
			userMeta := byte(0)
			if len(kv.UserMeta) > 0 {
				userMeta = kv.UserMeta[0]
			}

			entry := &badger.Entry{
				Key:      kv.Key,
				Value:    kv.Value,
				UserMeta: userMeta,
			}

			if err := wb.SetEntryAt(entry, kv.Version); err != nil {
				return fmt.Errorf("error writing to batch: %w", err)
			}
			count++
		}
		return nil
	}

	if err := stream.Orchestrate(context.Background()); err != nil {
		glog.Warningf("Error streaming predicate %s: %v", pred, err)
		return 0
	}

	if err := wb.Flush(); err != nil {
		glog.Errorf("Error flushing batch for predicate %s: %v", pred, err)
		return 0
	}

	return count
}

const limit = 2 << 30