fix(query): avoid eager UID materialization on posting reads

posting/list.go

@@ -1712,37 +1712,57 @@ func (l *List) rollup(readTs uint64, split bool) (*rollupOutput, error) {
 func (l *List) ApproxLen() int {
 	l.RLock()
 	defer l.RUnlock()
+	return l.approxLen()
+}
+
+func (l *List) approxLen() int {
+	l.AssertRLock()
 	return l.mutationMap.len() + codec.ApproxLen(l.plist.Pack)
 }
 
 func (l *List) calculateUids() error {
-	l.RLock()
-	if l.mutationMap == nil || l.mutationMap.isUidsCalculated {
-		l.RUnlock()
-		return nil
-	}
-	res := make([]uint64, 0, l.ApproxLen())
-
-	err := l.iterate(l.mutationMap.committedUidsTime, 0, func(p *pb.Posting) error {
-		if p.PostingType == pb.Posting_REF {
-			res = append(res, p.Uid)
+	for {
+		l.RLock()
+		if l.mutationMap == nil || l.mutationMap.isUidsCalculated {
+			l.RUnlock()
+			return nil
 		}
-		return nil
-	})
+		calculatedAt := l.mutationMap.committedUidsTime
+		res := make([]uint64, 0, l.approxLen())
 
-	l.RUnlock()
+		err := l.iterate(calculatedAt, 0, func(p *pb.Posting) error {
+			if p.PostingType == pb.Posting_REF {
+				res = append(res, p.Uid)
+			}
+			return nil
+		})
 
-	if err != nil {
-		return err
-	}
+		l.RUnlock()
 
-	l.Lock()
-	defer l.Unlock()
+		if err != nil {
+			return err
+		}
+
+		l.Lock()
+		if l.mutationMap == nil || l.mutationMap.isUidsCalculated {
+			l.Unlock()
+			return nil
+		}
+		if l.mutationMap.currentEntries != nil {
+			l.Unlock()
+			return nil
+		}
+		if l.mutationMap.committedUidsTime != calculatedAt {
+			l.Unlock()
+			continue
+		}
 
-	l.mutationMap.calculatedUids = res
-	l.mutationMap.isUidsCalculated = true
+		l.mutationMap.calculatedUids = res
+		l.mutationMap.isUidsCalculated = true
+		l.Unlock()
 
-	return nil
+		return nil
+	}
 }
 
 // canUseCalculatedUids reports whether calculatedUids can serve a read at readTs. The slice is
@@ -1765,12 +1785,14 @@ func (l *List) canUseCalculatedUids(readTs uint64) bool {
 // We have to apply the filtering before applying (offset, count).
 // WARNING: Calling this function just to get UIDs is expensive
 func (l *List) Uids(opt ListOptions) (*pb.List, error) {
+	requestedFirst := opt.First
+	bounded := requestedFirst > 0 && requestedFirst < math.MaxInt32
 	if opt.First == 0 {
 		opt.First = math.MaxInt32
 	}
 
 	getUidList := func() (*pb.List, error, bool) {
-		if l.canUseCalculatedUids(opt.ReadTs) {
+		if opt.Intersect == nil && !bounded && l.canUseCalculatedUids(opt.ReadTs) {
 			l.RLock()
 
 			afterIdx := 0
@@ -1792,13 +1814,11 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			out := &pb.List{Uids: copyArr}
 			l.RUnlock()
 
-			return out, nil, opt.Intersect != nil
+			return out, nil, false
 		}
 		// Pre-assign length to make it faster.
 		l.RLock()
 		defer l.RUnlock()
-		// Use approximate length for initial capacity.
-		res := make([]uint64, 0, l.ApproxLen())
 		out := &pb.List{}
 
 		if l.mutationMap.len() == 0 && opt.Intersect != nil && len(l.plist.Splits) == 0 {
@@ -1809,10 +1829,13 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			return out, nil, false
 		}
 
+		approxLen := l.approxLen()
+
 		// If we need to intersect and the number of elements are small, in that case it's better to
 		// just check each item is present or not.
-		if opt.Intersect != nil && len(opt.Intersect.Uids) < l.ApproxLen() {
+		if opt.Intersect != nil && len(opt.Intersect.Uids) < approxLen {
 			// Cache the iterator as it makes the search space smaller each time.
+			res := make([]uint64, 0, len(opt.Intersect.Uids))
 			var pitr pIterator
 			for _, uid := range opt.Intersect.Uids {
 				ok, _, err := l.findPostingWithItr(opt.ReadTs, uid, pitr)
@@ -1836,6 +1859,12 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			uidMax = opt.Intersect.Uids[len(opt.Intersect.Uids)-1]
 		}
 
+		resCap := approxLen
+		if bounded && requestedFirst+1 < resCap {
+			resCap = requestedFirst + 1
+		}
+		res := make([]uint64, 0, resCap)
+
 		err := l.iterate(opt.ReadTs, opt.AfterUid, func(p *pb.Posting) error {
 			if p.PostingType == pb.Posting_REF {
 				if p.Uid < uidMin {

posting/list_test.go

@@ -1874,3 +1874,40 @@ func TestCalculatedUidsRespectReadTs(t *testing.T) {
 	// A read before every commit sees nothing.
 	require.Empty(t, uidsAt(5))
 }
+
+func TestCalculatedUidsSkippedForBoundedReads(t *testing.T) {
+	key := x.DataKey(x.AttrInRootNamespace("calculatedUidsBoundedReads"), 7)
+
+	txn := NewTxn(5)
+	l, err := txn.Get(key)
+	require.NoError(t, err)
+	for _, uid := range []uint64{2, 3, 4} {
+		addMutationHelper(t, l, &pb.DirectedEdge{ValueId: uid}, Set, txn)
+	}
+	require.NoError(t, l.commitMutation(5, 10))
+	require.NoError(t, l.calculateUids())
+	require.True(t, l.canUseCalculatedUids(10))
+
+	l.Lock()
+	l.mutationMap.calculatedUids = []uint64{100, 101}
+	l.Unlock()
+
+	unbounded, err := l.Uids(ListOptions{ReadTs: 10})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{100, 101}, unbounded.Uids)
+
+	workerUnbounded, err := l.Uids(ListOptions{ReadTs: 10, First: math.MaxInt32})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{100, 101}, workerUnbounded.Uids)
+
+	first, err := l.Uids(ListOptions{ReadTs: 10, First: 1})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{2}, first.Uids)
+
+	intersect, err := l.Uids(ListOptions{
+		ReadTs:    10,
+		Intersect: &pb.List{Uids: []uint64{3}},
+	})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{3}, intersect.Uids)
+}

posting/mvcc.go

@@ -409,6 +409,10 @@ func (ml *MemoryLayer) UpdateMaxCost(maxCost int64) {
 	ml.cache.data.UpdateMaxCost(maxCost)
 }
 
+func (ml *MemoryLayer) hasCache() bool {
+	return ml != nil && ml.cache != nil && ml.cache.data != nil
+}
+
 type IterateDiskArgs struct {
 	Prefix         []byte
 	Prefetch       bool
@@ -761,19 +765,24 @@ func (c *CachePL) Set(l *List, readTs uint64) {
 	}
 }
 
-func (ml *MemoryLayer) readFromCache(key []byte, readTs uint64) *List {
+func (ml *MemoryLayer) readFromCache(key []byte, readTs uint64, readUids bool) (*List, error) {
 	cacheItem, ok := ml.cache.get(key)
 
 	// Issue #9597 fix: Cache is only valid if minTs <= readTs AND maxTs >= readTs.
 	// If maxTs < readTs, the cache is missing mutations committed after maxTs.
 	if ok && cacheItem.list != nil && cacheItem.list.minTs <= readTs && cacheItem.list.maxTs >= readTs {
+		if readUids && ml.hasCache() {
+			if err := cacheItem.list.calculateUids(); err != nil {
+				return nil, err
+			}
+		}
 		cacheItem.list.RLock()
 		lCopy := copyList(cacheItem.list)
 		cacheItem.list.RUnlock()
 		checkForRollup(key, lCopy)
-		return lCopy
+		return lCopy, nil
 	}
-	return nil
+	return nil, nil
 }
 
 func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64, readUids bool) (*List, error) {
@@ -792,7 +801,7 @@ func (ml *MemoryLayer) readFromDisk(key []byte, pstore *badger.DB, readTs uint64
 	if err != nil {
 		return l, err
 	}
-	if readUids {
+	if readUids && ml.hasCache() {
 		if err := l.calculateUids(); err != nil {
 			return nil, err
 		}
@@ -814,12 +823,15 @@ func (ml *MemoryLayer) ReadData(key []byte, pstore *badger.DB, readTs uint64, re
 	// We first try to read the data from cache, if it is present. If it's not present, then we would read the
 	// latest data from the disk. This would get stored in the cache. If this read has a minTs > readTs then
 	// we would have to read the correct timestamp from the disk.
-	l := ml.readFromCache(key, readTs)
+	l, err := ml.readFromCache(key, readTs, readUids)
+	if err != nil {
+		return nil, err
+	}
 	if l != nil {
 		l.mutationMap.setTs(readTs)
 		return l, nil
 	}
-	l, err := ml.readFromDisk(key, pstore, math.MaxUint64, readUids)
+	l, err = ml.readFromDisk(key, pstore, math.MaxUint64, readUids)
 	if err != nil {
 		return nil, err
 	}

posting/mvcc_test.go

@@ -273,6 +273,64 @@ func TestCacheStaleWhenMaxTsLessThanReadTs(t *testing.T) {
 	require.True(t, hasUid2, "UID 2 missing - cache returned stale data (maxTs < readTs)")
 }
 
+func TestReadUidsHonorsPostingListCache(t *testing.T) {
+	require.NoError(t, pstore.DropAll())
+
+	origMemLayer := MemLayerInstance
+	MemLayerInstance = initMemoryLayer(0, false)
+	t.Cleanup(func() {
+		MemLayerInstance = origMemLayer
+	})
+
+	attr := x.AttrInRootNamespace("readUidsCache")
+	key := x.DataKey(attr, 1)
+	addEdgeToUID(t, attr, 1, 2, 1, 2)
+	addEdgeToUID(t, attr, 1, 3, 3, 4)
+
+	l, err := getNew(key, pstore, math.MaxUint64, true)
+	require.NoError(t, err)
+	require.False(t, l.mutationMap.isUidsCalculated,
+		"readUids should not materialize UIDs when posting-list cache is disabled")
+
+	MemLayerInstance = initMemoryLayer(10<<20, false)
+	l, err = getNew(key, pstore, math.MaxUint64, true)
+	require.NoError(t, err)
+	require.True(t, l.mutationMap.isUidsCalculated,
+		"readUids should still warm calculated UIDs when posting-list cache is enabled")
+}
+
+func TestReadUidsWarmsCachedPostingList(t *testing.T) {
+	require.NoError(t, pstore.DropAll())
+
+	origMemLayer := MemLayerInstance
+	MemLayerInstance = initMemoryLayer(10<<20, false)
+	t.Cleanup(func() {
+		MemLayerInstance = origMemLayer
+	})
+
+	attr := x.AttrInRootNamespace("readUidsCacheHit")
+	key := x.DataKey(attr, 1)
+	addEdgeToUID(t, attr, 1, 2, 1, 2)
+	addEdgeToUID(t, attr, 1, 3, 3, 4)
+
+	l, err := getNew(key, pstore, math.MaxUint64, false)
+	require.NoError(t, err)
+	require.False(t, l.mutationMap.isUidsCalculated)
+	MemLayerInstance.wait()
+
+	cacheItem, ok := MemLayerInstance.cache.get(key)
+	require.True(t, ok)
+	require.False(t, cacheItem.list.mutationMap.isUidsCalculated)
+
+	l, err = getNew(key, pstore, math.MaxUint64, true)
+	require.NoError(t, err)
+	require.True(t, l.mutationMap.isUidsCalculated)
+
+	cacheItem, ok = MemLayerInstance.cache.get(key)
+	require.True(t, ok)
+	require.True(t, cacheItem.list.mutationMap.isUidsCalculated)
+}
+
 func TestPostingListRead(t *testing.T) {
 	attr := x.AttrInRootNamespace("emptypl")
 	key := x.DataKey(attr, 1)