perf: optimize TID entry lookup

frac/active_index.go

@@ -240,8 +240,8 @@ type activeTokenIndex struct {
 	inverser  *inverser
 }
 
-func (si *activeTokenIndex) GetValByTID(tid uint32) []byte {
-	return si.tokenList.GetValByTID(tid)
+func (si *activeTokenIndex) GetValByTID(tid uint32, field string) []byte {
+	return si.tokenList.GetValByTID(tid, field)
 }
 
 func (si *activeTokenIndex) GetTIDsByTokenExpr(t parser.Token) ([]uint32, error) {

frac/active_token_list.go

@@ -177,7 +177,7 @@ func (tl *TokenList) initSystemTokens() {
 	tl.allTokenLIDs = tlids[0]
 }
 
-func (tl *TokenList) GetValByTID(tid uint32) []byte {
+func (tl *TokenList) GetValByTID(tid uint32, _ string) []byte {
 	tl.tidMu.RLock()
 	defer tl.tidMu.RUnlock()
 

frac/processor/aggregator.go

@@ -375,6 +375,7 @@ type SourcedNodeIterator struct {
 	sourcedNode node.Sourced
 	ti          tokenIndex
 	tids        []uint32
+	field       string
 
 	tokensCache map[uint32]string
 
@@ -385,12 +386,13 @@ type SourcedNodeIterator struct {
 	lastSource uint32
 }
 
-func NewSourcedNodeIterator(sourced node.Sourced, ti tokenIndex, tids []uint32, limit iteratorLimit) *SourcedNodeIterator {
+func NewSourcedNodeIterator(sourced node.Sourced, ti tokenIndex, tids []uint32, field string, limit iteratorLimit) *SourcedNodeIterator {
 	lastID, lastSource := sourced.NextSourced()
 	return &SourcedNodeIterator{
 		sourcedNode:      sourced,
 		ti:               ti,
 		tids:             tids,
+		field:            field,
 		tokensCache:      make(map[uint32]string),
 		uniqSourcesLimit: limit,
 		countBySource:    make(map[uint32]int),
@@ -425,14 +427,14 @@ func (s *SourcedNodeIterator) ConsumeTokenSource(lid node.LID) (uint32, bool, er
 func (s *SourcedNodeIterator) ValueBySource(source uint32) string {
 	const useCacheThreshold = 2
 	if s.countBySource[source] < useCacheThreshold {
-		return string(s.ti.GetValByTID(s.tids[source]))
+		return string(s.ti.GetValByTID(s.tids[source], s.field))
 	}
 
 	val, ok := s.tokensCache[source]
 	if ok {
 		return val
 	}
-	val = string(s.ti.GetValByTID(s.tids[source]))
+	val = string(s.ti.GetValByTID(s.tids[source], s.field))
 	s.tokensCache[source] = val
 	return val
 }

frac/processor/aggregator_test.go

@@ -29,7 +29,7 @@ func TestSingleSourceCountAggregator(t *testing.T) {
 	}
 
 	source := node.BuildORTreeAgg(node.MakeStaticNodes(sources))
-	iter := NewSourcedNodeIterator(source, nil, nil, iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+	iter := NewSourcedNodeIterator(source, nil, nil, "", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
 	agg := NewSingleSourceCountAggregator(iter, provideExtractTimeFunc(nil, nil, 0))
 	for _, id := range searchDocs {
 		if err := agg.Next(node.NewDescLID(id)); err != nil {
@@ -57,7 +57,7 @@ func TestSingleSourceCountAggregatorWithInterval(t *testing.T) {
 	}
 
 	source := node.BuildORTreeAgg(node.MakeStaticNodes(sources))
-	iter := NewSourcedNodeIterator(source, nil, nil, iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+	iter := NewSourcedNodeIterator(source, nil, nil, "", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
 
 	agg := NewSingleSourceCountAggregator(iter, func(l seq.LID) seq.MID {
 		return seq.MID(l) % 3
@@ -99,7 +99,7 @@ func BenchmarkAggDeep(b *testing.B) {
 			r := rand.New(rand.NewSource(benchRandSeed))
 			v, _ := Generate(r, s)
 			src := node.NewSourcedNodeWrapper(node.NewStatic(v, false), 0)
-			iter := NewSourcedNodeIterator(src, nil, make([]uint32, 1), iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+			iter := NewSourcedNodeIterator(src, nil, make([]uint32, 1), "", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
 			n := NewSingleSourceCountAggregator(iter, provideExtractTimeFunc(nil, nil, 0))
 			vals, _ := Generate(r, s)
 
@@ -133,7 +133,7 @@ func BenchmarkAggWide(b *testing.B) {
 
 			source := node.BuildORTreeAgg(node.MakeStaticNodes(wide))
 
-			iter := NewSourcedNodeIterator(source, nil, make([]uint32, len(wide)), iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+			iter := NewSourcedNodeIterator(source, nil, make([]uint32, len(wide)), "", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
 			n := NewSingleSourceCountAggregator(iter, provideExtractTimeFunc(nil, nil, 0))
 			vals, _ := Generate(r, s)
 
@@ -152,7 +152,7 @@ type MockTokenIndex struct {
 	tokenIndex // embed to implement TokenIndex interface and override only needed methods
 }
 
-func (m *MockTokenIndex) GetValByTID(tid uint32) []byte {
+func (m *MockTokenIndex) GetValByTID(tid uint32, _ string) []byte {
 	return []byte(strconv.Itoa(int(tid)))
 }
 
@@ -211,8 +211,8 @@ func TestTwoSourceAggregator(t *testing.T) {
 
 	fieldTIDs := []uint32{42, 73}
 	groupByTIDs := []uint32{1, 2}
-	groupIterator := NewSourcedNodeIterator(groupBy, dp, groupByTIDs, iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
-	fieldIterator := NewSourcedNodeIterator(field, dp, fieldTIDs, iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+	groupIterator := NewSourcedNodeIterator(groupBy, dp, groupByTIDs, "group_by", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+	fieldIterator := NewSourcedNodeIterator(field, dp, fieldTIDs, "field", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
 	limits := AggLimits{}
 	aggregator := NewGroupAndFieldAggregator(
 		fieldIterator, groupIterator, provideExtractTimeFunc(nil, nil, 0), true, false, limits,
@@ -264,7 +264,7 @@ func TestSingleTreeCountAggregator(t *testing.T) {
 		},
 	}
 
-	iter := NewSourcedNodeIterator(field, dp, []uint32{0}, iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
+	iter := NewSourcedNodeIterator(field, dp, []uint32{0}, "field", iteratorLimit{limit: 0, err: consts.ErrTooManyGroupTokens})
 	aggregator := NewSingleSourceCountAggregator(iter, provideExtractTimeFunc(nil, nil, 0))
 
 	r.NoError(aggregator.Next(node.NewDescLID(1)))
@@ -302,7 +302,7 @@ func TestAggregatorLimitExceeded(t *testing.T) {
 
 	for _, expectedErr := range []error{consts.ErrTooManyGroupTokens, consts.ErrTooManyFieldTokens} {
 		source := node.BuildORTreeAgg(node.MakeStaticNodes(sources))
-		iter := NewSourcedNodeIterator(source, nil, nil, iteratorLimit{limit: limit, err: expectedErr})
+		iter := NewSourcedNodeIterator(source, nil, nil, "", iteratorLimit{limit: limit, err: expectedErr})
 		agg := NewSingleSourceCountAggregator(iter, provideExtractTimeFunc(nil, nil, 0))
 
 		var limitErr error

frac/processor/eval_tree.go

@@ -223,5 +223,5 @@ func iteratorFromLiteral(
 	}
 
 	sourcedNode := node.BuildORTreeAgg(lidsTids)
-	return NewSourcedNodeIterator(sourcedNode, ti, tids, iteratorLimit), nil
+	return NewSourcedNodeIterator(sourcedNode, ti, tids, literal.Field, iteratorLimit), nil
 }

frac/processor/search.go

@@ -32,7 +32,7 @@ type idsIndex interface {
 }
 
 type tokenIndex interface {
-	GetValByTID(tid uint32) []byte
+	GetValByTID(tid uint32, field string) []byte
 	GetTIDsByTokenExpr(token parser.Token) ([]uint32, error)
 	GetLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.Node
 }

frac/sealed/token/table.go

@@ -71,26 +71,29 @@ func (t Table) SelectEntries(field, hint string) []*TableEntry {
 	return data.Entries[l:r]
 }
 
-func (t Table) GetEntryByTID(tid uint32) *TableEntry {
+func (t Table) GetEntryByTID(tid uint32, field string) *TableEntry {
 	if tid == 0 {
 		return nil
 	}
-	for _, data := range t {
-		from := data.Entries[0].StartTID
-		to := data.Entries[len(data.Entries)-1].GetLastTID()
-		if tid < from || tid > to {
-			continue
-		}
 
-		i := sort.Search(len(data.Entries), func(j int) bool {
-			return data.Entries[j].StartTID > tid
-		})
+	data, ok := t[field]
+	if !ok {
+		logger.Panic("can't find field", zap.String("field", field))
+		return nil
+	}
 
-		return data.Entries[i-1]
+	from := data.Entries[0].StartTID
+	to := data.Entries[len(data.Entries)-1].GetLastTID()
+	if tid < from || tid > to {
+		logger.Panic("tid out of range", zap.Uint32("tid", tid), zap.Uint32("from", from), zap.Uint32("to", to))
+		return nil
 	}
 
-	logger.Panic("can't find tid", zap.Uint32("tid", tid))
-	return nil
+	i := sort.Search(len(data.Entries), func(j int) bool {
+		return data.Entries[j].StartTID > tid
+	})
+
+	return data.Entries[i-1]
 }
 
 // Size calculates a very approximate amount of memory occupied

frac/sealed_index.go

@@ -228,9 +228,9 @@ type sealedTokenIndex struct {
 	tokenBlockLoader *token.BlockLoader
 }
 
-func (ti *sealedTokenIndex) GetValByTID(tid uint32) []byte {
+func (ti *sealedTokenIndex) GetValByTID(tid uint32, field string) []byte {
 	tokenTable := ti.tokenTableLoader.Load()
-	if entry := tokenTable.GetEntryByTID(tid); entry != nil {
+	if entry := tokenTable.GetEntryByTID(tid, field); entry != nil {
 		block := ti.tokenBlockLoader.Load(entry.BlockIndex)
 		return block.GetToken(entry.GetIndexInTokensBlock(tid))
 	}