perf: do not match each token in aggregation

.seqbench/continuous.env

@@ -22,6 +22,11 @@ SEQDB_LIMITS_SEARCH_REQUESTS=1024
 SEQDB_LIMITS_BULK_REQUESTS=128
 SEQDB_LIMITS_INFLIGHT_BULK=128
 
+SEQDB_LIMITS_AGGREGATION_FIELD_TOKENS=9223372036854775807
+SEQDB_LIMITS_AGGREGATION_FIELD_VALUES=9223372036854775807
+SEQDB_LIMITS_AGGREGATION_GROUP_TOKENS=9223372036854775807
+SEQDB_LIMITS_AGGREGATION_FRACTION_TOKENS=9223372036854775807
+
 # DO NOT CHANGE FOLLOWING VALUES
 # SEQBAZOOKA RELIES ON THEM
 SEQDB_STORAGE_DATA_DIR=/var/seqdb

frac/active_index.go

@@ -244,6 +244,10 @@ func (si *activeTokenIndex) GetValByTID(tid uint32, field string) []byte {
 	return si.tokenList.GetValByTID(tid, field)
 }
 
+func (si *activeTokenIndex) GetTIDsByField(field string) ([]uint32, error) {
+	return si.tokenList.GetTIDsByField(field), nil
+}
+
 func (si *activeTokenIndex) GetTIDsByTokenExpr(t parser.Token) ([]uint32, error) {
 	return si.tokenList.FindPattern(si.ctx, t)
 }

frac/processor/aggregator.go

@@ -1,8 +1,10 @@
 package processor
 
 import (
+	"cmp"
 	"fmt"
 	"math"
+	"slices"
 	"strconv"
 
 	"github.com/ozontech/seq-db/consts"
@@ -113,6 +115,9 @@ func (n *TwoSourceAggregator) Next(lid node.LID) error {
 
 // Aggregate processes and returns the final aggregation result.
 func (n *TwoSourceAggregator) Aggregate() (seq.AggregatableSamples, error) {
+	n.groupBy.prefetchTokenValues()
+	n.field.prefetchTokenValues()
+
 	aggMap := make(map[seq.AggBin]*seq.SamplesContainer, n.groupBy.UniqueSources())
 
 	var sourceValuePool []string
@@ -226,6 +231,8 @@ func (n *SingleSourceCountAggregator) Next(lid node.LID) error {
 }
 
 func (n *SingleSourceCountAggregator) Aggregate() (seq.AggregatableSamples, error) {
+	n.group.prefetchTokenValues()
+
 	aggMap := make(map[seq.AggBin]*seq.SamplesContainer, n.group.UniqueSources())
 
 	for bin, cnt := range n.countBySource {
@@ -289,6 +296,8 @@ func (n *SingleSourceUniqueAggregator) Next(lid node.LID) error {
 }
 
 func (n *SingleSourceUniqueAggregator) Aggregate() (seq.AggregatableSamples, error) {
+	n.group.prefetchTokenValues()
+
 	aggMap := make(map[seq.AggBin]*seq.SamplesContainer, n.group.UniqueSources())
 
 	for val := range n.values {
@@ -342,6 +351,9 @@ func (n *SingleSourceHistogramAggregator) Next(lid node.LID) error {
 		return nil
 	}
 
+	// TODO(dkharms): Sequence of `source` values
+	// is in a random order so we again lose benefits of kernel read-ahead.
+	// Maybe it's worth it to do something like [prefetchTokenValues].
 	value := n.field.ValueBySource(source)
 	num, err := parseNum(value)
 	if err != nil {
@@ -411,31 +423,53 @@ func (s *SourcedNodeIterator) ConsumeTokenSource(lid node.LID) (uint32, bool, er
 		return 0, false, nil
 	}
 
-	if s.uniqSourcesLimit.limit <= 0 {
-		return s.lastSource, true, nil
-	}
-
 	s.countBySource[s.lastSource]++
-
-	if len(s.countBySource) > s.uniqSourcesLimit.limit {
+	if s.uniqSourcesLimit.limit > 0 && len(s.countBySource) > s.uniqSourcesLimit.limit {
 		return lid.Unpack(), true, fmt.Errorf("%w: iterator limit is exceeded", s.uniqSourcesLimit.err)
 	}
 
 	return s.lastSource, true, nil
 }
 
+func (s *SourcedNodeIterator) prefetchTokenValues() {
+	if s.ti == nil || len(s.countBySource) == 0 {
+		return
+	}
+
+	// NOTE(dkharms): Since `countBySource` is a hashmap and
+	// its iteration order is not determined, we lose benefits
+	// of kernel read-ahead.
+	//
+	// In this method we establish the order again.
+	sources := make([]uint32, 0, len(s.countBySource))
+	for source := range s.countBySource {
+		if _, ok := s.tokensCache[source]; !ok {
+			sources = append(sources, source)
+		}
+	}
+
+	slices.SortFunc(sources, func(a, b uint32) int {
+		return cmp.Compare(s.tids[a], s.tids[b])
+	})
+
+	for _, source := range sources {
+		s.tokensCache[source] = string(s.ti.GetValByTID(s.tids[source], s.field))
+	}
+}
+
 func (s *SourcedNodeIterator) ValueBySource(source uint32) string {
+	if val, ok := s.tokensCache[source]; ok {
+		return val
+	}
+
 	const useCacheThreshold = 2
 	if s.countBySource[source] < useCacheThreshold {
 		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], s.field))
+	val := string(s.ti.GetValByTID(s.tids[source], s.field))
 	s.tokensCache[source] = val
+
 	return val
 }
 

frac/processor/eval_tree.go

@@ -200,11 +200,13 @@ func iteratorFromLiteral(
 	iteratorLimit iteratorLimit,
 	order seq.DocsOrder,
 ) (*SourcedNodeIterator, error) {
-	m := sw.Start("get_tids_by_token_expr")
-	tids, err := ti.GetTIDsByTokenExpr(literal)
+	m := sw.Start("get_tids_by_field")
+	// For aggregations we can receive the first and the last TID for field,
+	// because we build tree over *all* token values.
+	tids, err := ti.GetTIDsByField(literal.Field)
 	m.Stop()
 	if err != nil {
-		return nil, fmt.Errorf("getting TIDs by token expression: %s", err)
+		return nil, fmt.Errorf("getting TIDs for field %q: %s", literal.Field, err)
 	}
 
 	if len(tids) > maxTIDs && maxTIDs > 0 {

frac/processor/search.go

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

frac/sealed_index.go

@@ -237,6 +237,25 @@ func (ti *sealedTokenIndex) GetValByTID(tid uint32, field string) []byte {
 	return nil
 }
 
+func (ti *sealedTokenIndex) GetTIDsByField(field string) ([]uint32, error) {
+	table := ti.tokenTableLoader.Load()
+
+	entries := table.SelectEntries(field, "")
+	if len(entries) == 0 {
+		return nil, nil
+	}
+
+	first := entries[0].StartTID
+	last := entries[len(entries)-1].GetLastTID()
+
+	tids := make([]uint32, (last-first)+1)
+	for i := range tids {
+		tids[i] = first + uint32(i)
+	}
+
+	return tids, nil
+}
+
 func (ti *sealedTokenIndex) GetTIDsByTokenExpr(t parser.Token) ([]uint32, error) {
 	field := parser.GetField(t)
 	searchStr := parser.GetHint(t)