summaryrefslogtreecommitdiff
path: root/core/search/search.go
blob: 721574c8aa8a37608fb51ac0544a01370d0724d7 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
package search

import (
	"cmp"
	"searchEngine/core/analyzer"
	"searchEngine/core/index"
	"searchEngine/core/scoring"
	"slices"
)

type Result struct {
	DocID int
	Score float64
}

func Search(idx *index.InvertedIndex, query string) []string {
	tokenizedQuery := analyzer.ProcessText(query)
	if len(tokenizedQuery) == 0 {
		return nil
	}

	idx.Mu.RLock()
	defer idx.Mu.RUnlock()

	totalDocs := len(idx.DocNames)
	docMatchCounts := make(map[int]int)
	scores := make(map[int]float64)

	uniqueMap := make(map[string]bool)
	var uniqueTokens []string
	for _, token := range tokenizedQuery {
		if !uniqueMap[token] {
			uniqueMap[token] = true
			uniqueTokens = append(uniqueTokens, token)
		}
	}

	for _, token := range tokenizedQuery {
		postings, exists := idx.Data[token]
		if !exists {
			return nil
		}

		docFrequency := len(postings)
		idf := scoring.CalculateIDF(totalDocs, docFrequency)

		for _, p := range postings {
			docMatchCounts[p.DocId]++
			termScore := scoring.ScoreTFIDF(p, idf)
			scores[p.DocId] += termScore
		}
	}

	var results []Result
	for docID, count := range docMatchCounts {
		if count == len(uniqueTokens) {
			results = append(results, Result{DocID: docID, Score: scores[docID]})
		}
	}

	slices.SortFunc(results, func(a, b Result) int {
		return cmp.Compare(b.Score, a.Score)
	})

	var titles []string
	for _, res := range results {
		if title, exists := idx.DocNames[res.DocID]; exists {
			titles = append(titles, title)
		}
	}

	return titles
}