Allow append to existing lists with write iterator

algo/uidlist.go

@@ -27,13 +27,27 @@ type WriteIterator struct {
 
 }
 
-func NewWriteIterator(l *task.List) WriteIterator {
+func NewWriteIterator(l *task.List, whence int) WriteIterator {
+	blen := len(l.Blocks)
+	var cur *task.Block
+	if whence == 1 && blen != 0 {
+		// Set the iterator to the end of the list.
+		llen := len(l.Blocks[blen-1].List)
+		cur = l.Blocks[blen-1]
+		// Set the list to its actual size as we want to append.
+		cur.List = cur.List[:blockSize]
+		return WriteIterator{
+			list:     l,
+			curBlock: cur,
+			bidx:     blen - 1,
+			lidx:     llen,
+		}
+	}
 	// Initialise and allocate some memory.
 	if len(l.Blocks) == 0 {
 		l.Blocks = make([]*task.Block, 0, 2)
 	}
 
-	var cur *task.Block
 	if len(l.Blocks) > 0 {
 		cur = l.Blocks[0]
 	}
@@ -170,8 +184,7 @@ func (l *ListIterator) Next() {
 // Slice returns a new task.List with the elements between start index and end index
 // of  the list passed to it.
 func Slice(ul *task.List, start, end int) {
-	out := NewWriteIterator(ul)
-
+	out := NewWriteIterator(ul, 0)
 	it := NewListIterator(ul)
 	it.SeekToIndex(start)
 
@@ -188,8 +201,7 @@ func SortedListToBlock(l []uint64) *task.List {
 	if len(l) == 0 {
 		return b
 	}
-	wit := NewWriteIterator(b)
-
+	wit := NewWriteIterator(b, 0)
 	for _, it := range l {
 		wit.Append(it)
 	}
@@ -211,7 +223,7 @@ func ListLen(l *task.List) int {
 func IntersectWith(u, v *task.List) {
 	itu := NewListIterator(u)
 	itv := NewListIterator(v)
-	out := NewWriteIterator(u)
+	out := NewWriteIterator(u, 0)
 	for itu.Valid() && itv.Valid() {
 		uid := itu.Val()
 		vid := itv.Val()
@@ -230,7 +242,7 @@ func IntersectWith(u, v *task.List) {
 
 // ApplyFilter applies a filter to our UIDList.
 func ApplyFilter(u *task.List, f func(uint64, int) bool) {
-	out := NewWriteIterator(u)
+	out := NewWriteIterator(u, 0)
 	for i, block := range u.Blocks {
 		for j, uid := range block.List {
 			if f(uid, Idx(u, i, j)) {
@@ -255,8 +267,7 @@ func IntersectSorted(lists []*task.List) *task.List {
 	if len(lists) == 0 {
 		return o
 	}
-	out := NewWriteIterator(o)
-
+	out := NewWriteIterator(o, 0)
 	// Scan through the smallest list. Denote as A.
 	// For each x in A,
 	//   For each other list B,
@@ -312,7 +323,7 @@ func IntersectSorted(lists []*task.List) *task.List {
 func Difference(u, v *task.List) {
 	itu := NewListIterator(u)
 	itv := NewListIterator(v)
-	out := NewWriteIterator(u)
+	out := NewWriteIterator(u, 0)
 	for itu.Valid() && itv.Valid() {
 		uid := itu.Val()
 		vid := itv.Val()
@@ -335,8 +346,7 @@ func MergeSorted(lists []*task.List) *task.List {
 	if len(lists) == 0 {
 		return o
 	}
-	out := NewWriteIterator(o)
-
+	out := NewWriteIterator(o, 0)
 	h := &uint64Heap{}
 	heap.Init(h)
 	var lIt []ListIterator

algo/uidlist_test.go

@@ -259,6 +259,28 @@ func TestApplyFilterUint(t *testing.T) {
 	require.Equal(t, []uint64{1, 3, 5}, BlockToList(u))
 }
 
+func TestWriteAppend1(t *testing.T) {
+	l := []uint64{1, 2, 3}
+	u := newList(l)
+	w := NewWriteIterator(u, 1)
+	w.Append(4)
+	w.Append(5)
+	w.Append(6)
+	w.End()
+	require.Equal(t, []uint64{1, 2, 3, 4, 5, 6}, BlockToList(u))
+}
+
+func TestWriteAppend2(t *testing.T) {
+	l := []uint64{1, 2, 3}
+	u := newList(l)
+	w := NewWriteIterator(u, 0)
+	w.Append(4)
+	w.Append(5)
+	w.Append(6)
+	w.End()
+	require.Equal(t, []uint64{4, 5, 6}, BlockToList(u))
+}
+
 // sort interface for []uint64
 type uint64Slice []uint64
 

posting/list.go

@@ -564,7 +564,7 @@ func (l *List) Uids(opt ListOptions) *task.List {
 	defer l.RUnlock()
 
 	res := new(task.List)
-	wit := algo.NewWriteIterator(res)
+	wit := algo.NewWriteIterator(res, 0)
 	it := algo.NewListIterator(opt.Intersect)
 	l.iterate(opt.AfterUID, func(p *types.Posting) bool {
 		if postingType(p) != valueUid {

query/query.go

@@ -558,7 +558,7 @@ func newGraph(ctx context.Context, gq *gql.GraphQuery) (*SubGraph, error) {
 	}
 	if len(gq.UID) > 0 {
 		o := new(task.List)
-		it := algo.NewWriteIterator(o)
+		it := algo.NewWriteIterator(o, 0)
 		for _, uid := range gq.UID {
 			it.Append(uid)
 		}
@@ -746,7 +746,7 @@ func populateVarMap(sg *SubGraph, doneVars map[string]*task.List, isCascade bool
 	}
 
 	o := new(task.List)
-	out := algo.NewWriteIterator(o)
+	out := algo.NewWriteIterator(o, 0)
 	it := algo.NewListIterator(sg.DestUIDs)
 	i := -1
 	if !isCascade {

types/geofilter.go

@@ -326,7 +326,7 @@ func (q GeoQueryData) intersects(g geom.T) bool {
 func FilterGeoUids(uids *task.List, values []*task.Value, q *GeoQueryData) *task.List {
 	x.AssertTruef(len(values) == algo.ListLen(uids), "lengths not matching")
 	o := new(task.List)
-	out := algo.NewWriteIterator(o)
+	out := algo.NewWriteIterator(o, 0)
 	it := algo.NewListIterator(uids)
 	for i := -1; it.Valid(); it.Next() {
 		i++

worker/assign.go

@@ -75,7 +75,7 @@ func assignUids(ctx context.Context, num *task.Num) (*task.List, error) {
 	// Mutations successfully applied.
 
 	o := new(task.List)
-	out := algo.NewWriteIterator(o)
+	out := algo.NewWriteIterator(o, 0)
 	// Only the First N entities are newly assigned UIDs, so we collect them.
 	for i := 0; i < val; i++ {
 		out.Append(mutations.Edges[i].Entity)

worker/sort.go

@@ -221,13 +221,7 @@ func intersectBucket(ts *task.Sort, attr, token string, out []intersectedList) e
 				n = slack
 			}
 		}
-
-		in1 := algo.NewListIterator(il.ulist)
-		o := new(task.List)
-		out := algo.NewWriteIterator(o)
-		for ; in1.Valid(); in1.Next() {
-			out.Append(in1.Val())
-		}
+		out := algo.NewWriteIterator(il.ulist, 1)
 		i := 0
 		in2 := algo.NewListIterator(result)
 		for ; in2.Valid() && i < n; in2.Next() {
@@ -235,7 +229,6 @@ func intersectBucket(ts *task.Sort, attr, token string, out []intersectedList) e
 			i++
 		}
 		out.End()
-		il.ulist.Blocks = o.Blocks
 	} // end for loop over UID lists in UID matrix.
 
 	// Check out[i] sizes for all i.