Algorithms to handle UidPack

algo/packed.go

@@ -0,0 +1,296 @@
+/*
+ * Copyright 2019 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package algo
+
+import (
+	"container/heap"
+	"sort"
+
+	"github.com/dgraph-io/dgraph/codec"
+	"github.com/dgraph-io/dgraph/protos/pb"
+)
+
+// ApplyFilterPacked applies the filter to a list of packed uids.
+func ApplyFilterPacked(u *pb.UidPack, f func(uint64, int) bool) *pb.UidPack {
+	it := codec.NewUidPackIterator(u)
+	index := 0
+	encoder := codec.Encoder{BlockSize: int(u.BlockSize)}
+
+	for ; it.Valid(); it.Next() {
+		uid := it.Get()
+		if f(uid, index) {
+			encoder.Add(uid)
+		}
+		index++
+	}
+
+	return encoder.Done()
+}
+
+// IntersectWithLinPacked performs the intersection of two UidPack objects.
+func IntersectWithLinPacked(u, v *pb.UidPack) *pb.UidPack {
+	if u == nil || v == nil {
+		return nil
+	}
+
+	n := codec.ExactLen(u)
+	m := codec.ExactLen(v)
+	i, k := 0, 0
+	uIt := codec.NewUidPackIterator(u)
+	vIt := codec.NewUidPackIterator(v)
+	encoder := codec.Encoder{BlockSize: int(u.BlockSize)}
+
+	for i < n && k < m {
+		uid := uIt.Get()
+		vid := vIt.Get()
+
+		switch {
+		case uid > vid:
+			k++
+			vIt.Next()
+
+			for {
+				if !(k < m && vIt.Get() < uid) {
+					break
+				}
+				k++
+				vIt.Next()
+			}
+		case uid == vid:
+			encoder.Add(uid)
+			i++
+			uIt.Next()
+			k++
+			vIt.Next()
+		default:
+			i++
+			uIt.Next()
+
+			for {
+				if !(i < n && uIt.Get() < vid) {
+					break
+				}
+
+				i++
+				uIt.Next()
+			}
+		}
+	}
+
+	return encoder.Done()
+}
+
+type listInfoPacked struct {
+	l      *pb.UidPack
+	length int
+}
+
+// IntersectSortedPacked calculates the intersection of multiple lists and performs
+// the intersections from the smallest to the largest list.
+func IntersectSortedPacked(lists []*pb.UidPack) *pb.UidPack {
+	if len(lists) == 0 {
+		encoder := codec.Encoder{BlockSize: 10}
+		return encoder.Done()
+	}
+	ls := make([]listInfoPacked, 0, len(lists))
+	for _, list := range lists {
+		ls = append(ls, listInfoPacked{
+			l:      list,
+			length: codec.ExactLen(list),
+		})
+	}
+	// Sort the lists based on length.
+	sort.Slice(ls, func(i, j int) bool {
+		return ls[i].length < ls[j].length
+	})
+
+	if len(ls) == 1 {
+		// Return a copy of the UidPack.
+		return codec.CopyUidPack(ls[0].l)
+	}
+
+	out := IntersectWithLinPacked(ls[0].l, ls[1].l)
+	// Intersect from smallest to largest.
+	for i := 2; i < len(ls); i++ {
+		out := IntersectWithLinPacked(out, ls[i].l)
+		// Break if we reach size 0 as we can no longer
+		// add any element.
+		if codec.ExactLen(out) == 0 {
+			break
+		}
+	}
+	return out
+}
+
+// DifferencePacked performs the difference operation between two UidPack objects.
+func DifferencePacked(u, v *pb.UidPack) *pb.UidPack {
+	if u == nil || v == nil {
+		// If v == nil, then it's empty so the value of u - v is just u.
+		// Return a copy of u.
+		if v == nil {
+			return codec.CopyUidPack(u)
+		}
+
+		return nil
+	}
+
+	n := codec.ExactLen(u)
+	m := codec.ExactLen(v)
+	encoder := codec.Encoder{BlockSize: int(u.BlockSize)}
+	uIt := codec.NewUidPackIterator(u)
+	vIt := codec.NewUidPackIterator(v)
+	i, k := 0, 0
+
+	for i < n && k < m {
+		uid := uIt.Get()
+		vid := vIt.Get()
+
+		switch {
+		case uid < vid:
+			for i < n && uIt.Get() < vid {
+				encoder.Add(uIt.Get())
+				i++
+				uIt.Next()
+			}
+		case uid == vid:
+			i++
+			uIt.Next()
+			k++
+			vIt.Next()
+		default:
+			k++
+			vIt.Next()
+			for {
+				if !(k < m && vIt.Get() < uid) {
+					break
+				}
+				k++
+				vIt.Next()
+			}
+		}
+	}
+
+	for i < n && k >= m {
+		encoder.Add(uIt.Get())
+		i++
+		uIt.Next()
+	}
+
+	return encoder.Done()
+}
+
+// MergeSortedPacked merges already sorted UidPack objects into a single UidPack.
+func MergeSortedPacked(lists []*pb.UidPack) *pb.UidPack {
+	if len(lists) == 0 {
+		return nil
+	}
+
+	h := &uint64Heap{}
+	heap.Init(h)
+	maxSz := 0
+	// iters stores the iterator for each corresponding, list.
+	iters := make([]*codec.UidPackIterator, len(lists))
+	// lenghts stores the length of each list so they are only computed once.
+	lenghts := make([]int, len(lists))
+	blockSize := 0
+
+	for i, l := range lists {
+		if l == nil {
+			continue
+		}
+
+		if blockSize == 0 {
+			blockSize = int(l.BlockSize)
+		}
+
+		iters[i] = codec.NewUidPackIterator(lists[i])
+
+		lenList := codec.ExactLen(lists[i])
+		lenghts[i] = lenList
+		if lenList > 0 {
+			heap.Push(h, elem{
+				val:     iters[i].Get(),
+				listIdx: i,
+			})
+			if lenList > maxSz {
+				maxSz = lenList
+			}
+		}
+	}
+
+	// Our final output.
+	output := codec.Encoder{BlockSize: blockSize}
+	// empty is used to keep track of whether the encoder contains data since the
+	// encoder does not have an equivalent of len.
+	empty := true
+	// idx[i] is the element we are looking at for lists[i].
+	idx := make([]int, len(lists))
+	var last uint64   // Last element added to sorted / final output.
+	for h.Len() > 0 { // While heap is not empty.
+		me := (*h)[0] // Peek at the top element in heap.
+		if empty || me.val != last {
+			output.Add(me.val) // Add if unique.
+			last = me.val
+			empty = false
+		}
+		if idx[me.listIdx] >= lenghts[me.listIdx]-1 {
+			heap.Pop(h)
+		} else {
+			idx[me.listIdx]++
+			iters[me.listIdx].Next()
+			val := iters[me.listIdx].Get()
+			(*h)[0].val = val
+			heap.Fix(h, 0) // Faster than Pop() followed by Push().
+		}
+	}
+	return output.Done()
+}
+
+// IndexOfPacked finds the index of the given uid in the UidPack. If it doesn't find it,
+// it returns -1.
+func IndexOfPacked(u *pb.UidPack, uid uint64) int {
+	if u == nil {
+		return -1
+	}
+	decoder := codec.Decoder{Pack: u}
+	decoder.Seek(0, codec.SeekStart)
+
+	for {
+		if !decoder.Valid() {
+			return -1
+		}
+
+		if decoder.PeekNextBase() < uid {
+			decoder.Next()
+			continue
+		}
+
+		uids := decoder.Uids()
+		if len(uids) == 0 {
+			return -1
+		}
+
+		i := sort.Search(len(uids), func(i int) bool { return uids[i] >= uid })
+		if i < len(uids) && uids[i] == uid {
+			return i + int(u.BlockSize)*decoder.BlockIdx()
+		}
+
+		decoder.Next()
+	}
+
+	return -1
+}