MB-66396: IndexSnapshot API to retrieve highest frequency terms

Author: abhinavdangeti

index/scorch/scorch_test.go

@@ -2810,3 +2810,117 @@ func TestPersistorMergerOptions(t *testing.T) {
 		}
 	}
 }
+
+func TestIndexSnapshotHighestFrequencyTerms(t *testing.T) {
+	cfg := CreateConfig("TestIndexSnapshotHighestFrequencyTerms")
+	err := InitTest(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := DestroyTest(cfg)
+		if err != nil {
+			t.Log(err)
+		}
+	}()
+
+	testConfig := cfg
+	mp := mapping.NewIndexMapping()
+
+	textMapping := mapping.NewTextFieldMapping()
+	textMapping.Analyzer = "en"
+
+	docMapping := mapping.NewDocumentStaticMapping()
+	docMapping.AddFieldMappingsAt("text", textMapping)
+
+	mp.DefaultMapping = docMapping
+
+	analysisQueue := index.NewAnalysisQueue(1)
+	idx, err := NewScorch("storeName", testConfig, analysisQueue)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	err = idx.Open()
+	if err != nil {
+		t.Errorf("error opening index: %v", err)
+	}
+	defer func() {
+		err := idx.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	data := []map[string]string{
+		{
+			"id":   "one",
+			"text": "She sells sea shells by the sea shore",
+		},
+		{
+			"id":   "two",
+			"text": "The quick brown fox jumps over the lazy dog",
+		},
+		{
+			"id":   "three",
+			"text": "She sold sea shells to the person with the dog",
+		},
+		{
+			"id":   "four",
+			"text": "But there are a lot of dogs on the beach",
+		},
+		{
+			"id":   "five",
+			"text": "To hell with the foxes",
+		},
+		{
+			"id":   "six",
+			"text": "What about the dogs",
+		},
+		{
+			"id":   "seven",
+			"text": "Dogs are OK, foxes are not",
+		},
+	}
+
+	for _, d := range data {
+		doc := document.NewDocument(d["id"])
+		err = mp.MapDocument(doc, d)
+		if err != nil {
+			t.Errorf("error mapping doc: %v", err)
+		}
+
+		err = idx.Update(doc)
+		if err != nil {
+			t.Errorf("Error updating index: %v", err)
+		}
+	}
+
+	reader, err := idx.Reader()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := reader.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	limit := 5
+
+	// FIXME: Type assert to IndexSnapshot for now
+	if snap, ok := reader.(*IndexSnapshot); ok {
+		termFreqs, err := snap.HighestFrequencyTerms("text", limit)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if len(termFreqs) > limit {
+			t.Fatal("terms freqs have more than limit")
+		}
+
+		for i := range termFreqs {
+			fmt.Println(termFreqs[i].Term, termFreqs[i].Frequency)
+		}
+	}
+}

index/scorch/snapshot_index.go

@@ -23,6 +23,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"sort"
+	"strings"
 	"sync"
 	"sync/atomic"
 
@@ -1161,3 +1162,88 @@ func (is *IndexSnapshot) ThesaurusKeysRegexp(name string,
 func (is *IndexSnapshot) UpdateSynonymSearchCount(delta uint64) {
 	atomic.AddUint64(&is.parent.stats.TotSynonymSearches, delta)
 }
+
+type TermFreq struct {
+	Term      string `json:"term"`
+	Frequency uint64 `json:"frequency"`
+}
+
+// HighestFrequencyTerms returns the top N terms with the highest frequencies
+// for a given field across all segments in the index snapshot.
+// Returns a slice of term-frequency pairs sorted by frequency (descending).
+func (is *IndexSnapshot) HighestFrequencyTerms(field string, limit int) (
+	termsFreqs []TermFreq, err error) {
+	if len(is.segment) == 0 {
+		return nil, fmt.Errorf("no segments available")
+	}
+
+	if limit <= 0 {
+		return nil, fmt.Errorf("limit must be positive")
+	}
+
+	// Use a map to aggregate frequencies across segments
+	termFreqs := make(map[string]uint64)
+
+	// Process each segment to collect term frequencies
+	for _, segment := range is.segment {
+		dict, err := segment.segment.Dictionary(field)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get dictionary for field %s in segment: %v", field, err)
+		}
+
+		// Get iterator for all terms in this segment
+		itr := dict.AutomatonIterator(nil, nil, nil)
+		if itr == nil {
+			continue
+		}
+
+		// Iterate through all terms in this segment
+		for {
+			next, err := itr.Next()
+			if err != nil {
+				return nil, fmt.Errorf("error iterating dictionary: %v", err)
+			}
+			if next == nil {
+				break // End of terms
+			}
+
+			// Get postings list for this term
+			pl, err := dict.PostingsList([]byte(next.Term), segment.deleted, nil)
+			if err != nil {
+				continue // Skip this term if we can't get postings
+			}
+
+			// Aggregate frequency across segments
+			termStr := string(next.Term)
+			termFreqs[termStr] += pl.Count()
+		}
+	}
+
+	if len(termFreqs) == 0 {
+		return nil, fmt.Errorf("no terms found for field %s", field)
+	}
+
+	var termFreqList []TermFreq
+	for termStr, freq := range termFreqs {
+		termFreqList = append(termFreqList, TermFreq{
+			Term:      termStr,
+			Frequency: freq,
+		})
+	}
+
+	// Sort by frequency (descending)
+	sort.Slice(termFreqList, func(i, j int) bool {
+		if termFreqList[i].Frequency == termFreqList[j].Frequency {
+			// If frequencies are equal, sort by term lexicographically
+			return strings.Compare(termFreqList[i].Term, termFreqList[j].Term) < 0
+		}
+		return termFreqList[i].Frequency > termFreqList[j].Frequency
+	})
+
+	// Limit results
+	if limit > len(termFreqList) {
+		limit = len(termFreqList)
+	}
+
+	return termFreqList[:limit], nil
+}