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

posting/list.go

@@ -1765,12 +1765,13 @@ 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
 	if opt.First == 0 {
 		opt.First = math.MaxInt32
 	}
 
 	getUidList := func() (*pb.List, error, bool) {
-		if l.canUseCalculatedUids(opt.ReadTs) {
+		if opt.Intersect == nil && requestedFirst <= 0 && l.canUseCalculatedUids(opt.ReadTs) {
 			l.RLock()
 
 			afterIdx := 0
@@ -1797,8 +1798,6 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 		// 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,9 +1808,19 @@ func (l *List) Uids(opt ListOptions) (*pb.List, error) {
 			return out, nil, false
 		}
 
+		approxLen := l.ApproxLen()
+		resCap := approxLen
+		if opt.Intersect != nil && len(opt.Intersect.Uids) < resCap {
+			resCap = len(opt.Intersect.Uids)
+		}
+		if requestedFirst > 0 && requestedFirst < resCap {
+			resCap = requestedFirst
+		}
+		res := make([]uint64, 0, resCap)
+
 		// 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.
 			var pitr pIterator
 			for _, uid := range opt.Intersect.Uids {

posting/list_test.go

@@ -1874,3 +1874,33 @@ 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()
+
+	first, err := l.Uids(ListOptions{ReadTs: 10, First: 1})
+	require.NoError(t, err)
+	require.Equal(t, []uint64{2}, first.Uids)
+	require.LessOrEqual(t, cap(first.Uids), 2)
+
+	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
@@ -792,7 +796,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
 		}

posting/mvcc_test.go

@@ -273,6 +273,32 @@ 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 TestPostingListRead(t *testing.T) {
 	attr := x.AttrInRootNamespace("emptypl")
 	key := x.DataKey(attr, 1)

worker/task.go

@@ -778,6 +778,22 @@ func retrieveUidsAndFacets(args funcArgs, pl *posting.List, facetsTree *facetsTr
 	return uidList, fcsList, nil
 }
 
+func shouldPrecalculateUids(q *pb.Query, srcFn *functionContext, facetsTree *facetsTree,
+	opts posting.ListOptions) bool {
+	if q.DoCount || q.FacetParam != nil || facetsTree != nil {
+		return false
+	}
+	if opts.Intersect != nil || opts.First > 0 {
+		return false
+	}
+	switch srcFn.fnType {
+	case compareScalarFn, hasFn, uidInFn:
+		return false
+	default:
+		return true
+	}
+}
+
 // This function handles operations on uid posting lists. Index keys, reverse keys and some data
 // keys store uid posting lists.
 func (qs *queryState) handleUidPostings(
@@ -855,7 +871,13 @@ func (qs *queryState) handleUidPostings(
 			}
 
 			// Get or create the posting list for an entity, attribute combination.
-			pl, err := qs.cache.GetUids(key)
+			var pl *posting.List
+			var err error
+			if shouldPrecalculateUids(q, srcFn, facetsTree, opts) {
+				pl, err = qs.cache.GetUids(key)
+			} else {
+				pl, err = qs.cache.Get(key)
+			}
 			if err != nil {
 				return err
 			}