diff --git a/posting/list.go b/posting/list.go index b257568ef69..df0109ffd9d 100644 --- a/posting/list.go +++ b/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 { diff --git a/posting/list_test.go b/posting/list_test.go index a27e9416402..bd12d147dbf 100644 --- a/posting/list_test.go +++ b/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) +} diff --git a/posting/mvcc.go b/posting/mvcc.go index 77491fb8a8f..e316b11c01d 100644 --- a/posting/mvcc.go +++ b/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 } diff --git a/posting/mvcc_test.go b/posting/mvcc_test.go index 3a0553c886c..7487a884753 100644 --- a/posting/mvcc_test.go +++ b/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) diff --git a/worker/task.go b/worker/task.go index a59bff64fcc..8ecd5acfbf8 100644 --- a/worker/task.go +++ b/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 }