Search: inverted indexes, ranking and sync

// Analyze lowercases text, splits it on anything that is not a letter or a digit, drops stop
// words and stems each word. Index time and query time use the same function, so a query term
// and a document term meet in the same form.
func Analyze(text string) []Token {
  words := strings.FieldsFunc(strings.ToLower(text), func(r rune) bool {
    return !unicode.IsLetter(r) && !unicode.IsDigit(r)
  })
  out := make([]Token, 0, len(words))
  for pos, w := range words {
    if stopWords[w] {
      continue
    }
    out = append(out, Token{Term: Stem(w), Pos: pos})
  }
  return out
}

// idf is the Lucene form of inverse document frequency. A term in every document scores near 0;
// a term in one document out of N scores near ln(N).
func idf(n, df int) float64 {
  return math.Log(1 + (float64(n)-float64(df)+0.5)/(float64(df)+0.5))
}

// bm25 scores one term in one field. Term frequency saturates: the fifth occurrence adds far less
// than the first. A field longer than average gets a lower score for the same count.
func bm25(tf int, idf float64, length int, avgLen, k1, b float64) float64 {
  norm := k1 * (1 - b + b*float64(length)/avgLen)
  return idf * float64(tf) * (k1 + 1) / (float64(tf) + norm)
}

// Phrase ranks documents where the terms appear next to each other, in order, in one field. It
// walks the position list of the first term and checks that term i sits at position p + i.
func (ix *Index) Phrase(text string, p Params, k int) []Hit {
  toks := Analyze(text)
  if len(toks) == 0 {
    return nil
  }
  match := map[int32]bool{}
  for f := range numFields {
    for _, first := range ix.postings[f][toks[0].Term] {
      for _, start := range first.pos {
        if ix.hasPhraseAt(Field(f), first.doc, toks, start) {
          match[first.doc] = true
          break
        }
      }
    }
  }
  terms := make([]string, len(toks))
  for i, t := range toks {
    terms[i] = t.Term
  }
  return ix.scoreFiltered(terms, p, k, func(d int32) bool { return match[d] })
}

// PrefixTerms finds every term that starts with prefix. The dictionary is sorted, so the matches
// are one contiguous run: two binary searches find it.
func (ix *Index) PrefixTerms(prefix string, max int) []string {
  d := ix.Dictionary()
  lo := sort.SearchStrings(d, prefix)
  var out []string
  for i := lo; i < len(d) && strings.HasPrefix(d[i], prefix); i++ {
    if len(out) == max {
      break
    }
    out = append(out, d[i])
  }
  return out
}

// AutoFuzziness is the edit budget by term length: 0 edits up to 2 characters, 1 edit for 3 to 5,
// 2 edits above 5. Elasticsearch calls this AUTO.
func AutoFuzziness(term string) int {
  switch n := len(term); {
  case n <= 2:
    return 0
  case n <= 5:
    return 1
  }
  return 2
}

// FuzzyTerms finds dictionary terms within the edit budget of term. An edit is an insert, a
// delete, a substitution, or a swap of two neighbours.
func (ix *Index) FuzzyTerms(term string) []string {
  budget := AutoFuzziness(term)
  var out []string
  for _, t := range ix.Dictionary() {
    if Distance(term, t) <= budget {
      out = append(out, t)
    }
  }
  return out
}

// checkResultClaims asserts every difference the sheet states between the engine and Postgres.
func checkResultClaims(t *testing.T, res map[string]QueryResult) {
  t.Helper()
  bm := func(q, boost string) []string { return ids(res[q].BM25[boost]) }
  pg := func(q string) []string { return ids(res[q].PGRank) }
  trgm := func(q string) []string { return ids(res[q].PGTrigram) }

  // IDF: "tent" is rare and "waterproof" is common. BM25 puts three tents above the dry bag,
  // which says waterproof 4 times but never tent. ts_rank has no IDF and puts the bag second.
  if b := bm("waterproof-tent", "1"); pos(b, "p08") < 3 || !slices.Equal(b[:3], []string{"p06", "p21", "p07"}) {
    t.Errorf("waterproof tent, BM25: %v; want p06 p21 p07 first, then the dry bag p08", b)
  }
  if p := pg("waterproof-tent"); pos(p, "p08") != 1 {
    t.Errorf("waterproof tent, ts_rank: %v; want the dry bag p08 second", p)
  }
  // Both stem running to run and shoes to shoe. Neither stems runner, so the vest is missed.
  if r := res["running-shoes"]; !slices.Equal(r.Terms, []string{"run", "shoe"}) || r.TSQuery != "'run' | 'shoe'" {
    t.Errorf("running shoes: terms %v, tsquery %s", r.Terms, r.TSQuery)
  }
  if slices.Contains(bm("running-shoes", "1"), "p18") || slices.Contains(pg("running-shoes"), "p18") {
    t.Error("running shoes: the runner vest p18 should match neither engine")
  }
  // Words match "jacket ... rain"; the phrase matches only "rain jacket".
  if !slices.Contains(bm("rain-jacket", "1"), "p05") || !slices.Contains(pg("rain-jacket"), "p05") {
    t.Error("rain jacket as words: both engines should match p05")
  }
  if !slices.Equal(bm("rain-jacket-phrase", "1"), []string{"p04"}) || !slices.Equal(pg("rain-jacket-phrase"), []string{"p04"}) {
    t.Errorf("rain jacket as a phrase: %v and %v, want only p04", bm("rain-jacket-phrase", "1"), pg("rain-jacket-phrase"))
  }
  // A prefix expands to every dictionary term that starts with it.
  if e := res["wat"].Expansions["wat"]; !slices.Equal(e, []string{"watch", "water", "waterproof"}) || top(bm("wat", "1")) != "p09" || top(pg("wat")) != "p09" {
    t.Errorf("wat: expansions %v, tops %s and %s", e, top(bm("wat", "1")), top(pg("wat")))
  }
  // Typos: full text search finds nothing. Edit distance finds both. Trigrams find the
  // missing letter (similarity 0.44) and miss the swapped pair (0.27, under 0.3).
  for _, q := range []string{"jaket", "jakcet"} {
    if len(pg(q)) != 0 || !slices.Equal(res[q].Expansions[q], []string{"jacket"}) || len(bm(q, "1")) != 3 {
      t.Errorf("%s: ts_rank %v, fuzzy %v, BM25 %v", q, pg(q), res[q].Expansions[q], bm(q, "1"))
    }
  }
  if res["jaket"].TrgmSim["jacket"] < 0.3 || res["jakcet"].TrgmSim["jacket"] >= 0.3 {
    t.Errorf("similarity to jacket: jaket %v, jakcet %v; want one above 0.3, one below", res["jaket"].TrgmSim, res["jakcet"].TrgmSim)
  }
  if len(trgm("jaket")) != 3 || len(trgm("jakcet")) != 0 {
    t.Errorf("trigram: jaket %v, jakcet %v; want 3 hits, then none", trgm("jaket"), trgm("jakcet"))
  }
  // Stop words leave the query; a title boost changes the order.
  if r := res["stop-words"]; !slices.Equal(r.Terms, []string{"tent", "rain"}) || r.TSQuery != "'tent' | 'rain'" {
    t.Errorf("stop words: terms %v, tsquery %s", r.Terms, r.TSQuery)
  }
  if top(bm("stop-words", "1")) != "p21" || top(bm("stop-words", "3")) != "p04" {
    t.Errorf("title boost: top %s at x1, %s at x3; want p21, then p04", top(bm("stop-words", "1")), top(bm("stop-words", "3")))
  }
}

// BuildTypeahead precomputes, for every prefix of every logged query, the k most searched
// completions. A keystroke then costs one lookup, with no ranking at query time.
func BuildTypeahead(log []Suggestion, k, maxPrefix int) map[string][]Suggestion {
  table := map[string][]Suggestion{}
  for _, s := range log {
    q := strings.ToLower(s.Query)
    for n := 1; n <= len(q) && n <= maxPrefix; n++ {
      table[q[:n]] = append(table[q[:n]], s)
    }
  }
  for p, list := range table {
    slices.SortFunc(list, func(a, b Suggestion) int {
      if a.Count != b.Count {
        return b.Count - a.Count
      }
      return strings.Compare(a.Query, b.Query)
    })
    table[p] = list[:min(k, len(list))]
  }
  return table
}

// EncodeGaps stores a sorted posting list as the gaps between document numbers, each gap as a
// varint: 7 bits a byte. A common term has small gaps, so most gaps fit in one byte.
func EncodeGaps(docs []int32) []byte {
  out := make([]byte, 0, len(docs))
  prev := int32(0)
  for _, d := range docs {
    out = binary.AppendUvarint(out, uint64(d-prev))
    prev = d
  }
  return out
}