Add support for autocompletion suggestions (#9)

* feat: Add support for autocompletion suggestions

* lint: Fix ktlint violations

Author: haroldadmin

core/src/main/kotlin/FtsIndex.kt

@@ -1,5 +1,6 @@
 package com.haroldadmin.lucilla.core
 
+import com.haroldadmin.lucilla.core.rank.ld
 import com.haroldadmin.lucilla.core.rank.tfIdf
 import com.haroldadmin.lucilla.ir.Posting
 import com.haroldadmin.lucilla.ir.extractDocumentId
@@ -32,6 +33,11 @@ public data class SearchResult(
     val matchTerm: String,
 )
 
+public data class AutocompleteSuggestion(
+    val score: Double,
+    val suggestion: String,
+)
+
 /**
  * Alias for a token in a document
  */
@@ -239,10 +245,48 @@ public class FtsIndex<DocType : Any>(
             }
         }
 
-        results.sortByDescending { result -> result.score }
+        results.sortBy { result -> result.score }
         return results
     }
 
+    /**
+     * Fetches autocompletion suggestions for the given query.
+     *
+     * An autocompletion suggestion is a term present in the index that
+     * has the same prefix as the given search query. The results are sorted
+     * in order of their relevance score.
+     * e.g. "foo" -> "fool", "foot", "football"
+     *
+     * **Autocompletion suggestions can be unexpected if stemming is a part
+     * of your text processing pipeline.**
+     *
+     * For example, the Porter stemmer stems "football" to "footbal". Therefore,
+     * even if your input text contains the word "football", you will see "footbal"
+     * as an autocompletion suggestion instead.
+     *
+     * The simplest way around this is to use a [Pipeline] that does not contain
+     * a stemming step. Alternatively you can use a custom stemmer that emits
+     * both the original word and its stemmed variant to ensure the original
+     * word appears in the suggestions.
+     *
+     * *Expect the autocompletion ranking algorithm to change in future releases*
+     *
+     * @param query The search query to fetch autocompletion suggestions for
+     * @return List of autocompletion suggestions, sorted by their scores
+     */
+    public fun autocomplete(query: String): List<AutocompleteSuggestion> {
+        val suggestions = _index.prefixMap(query).keys
+            .fold(mutableListOf<AutocompleteSuggestion>()) { suggestions, prefixKey ->
+                val score = ld(query, prefixKey).toDouble() / prefixKey.length
+                val suggestion = AutocompleteSuggestion(score, prefixKey)
+                suggestions.apply { add(suggestion) }
+            }
+
+        suggestions.sortByDescending { it.score }
+
+        return suggestions
+    }
+
     /**
      * Clears all documents added to the FTS index
      */

core/src/test/kotlin/App.kt

@@ -12,17 +12,47 @@ fun main() {
     }
     println("index:complete ($timeToBuildIndex ms)")
 
-    do {
-        println("Enter search query (EXIT to stop)")
+    println("s: search, a: autocompletion suggestions")
+    when (readln().trim().lowercase()) {
+        "s" -> search(index, books)
+        "a" -> autocomplete(index, books)
+        else -> println("Invalid input")
+    }
+}
+
+fun search(index: FtsIndex<Book>, data: Map<Int, Book>) {
+    println("Search")
+    println("Enter search query (EXIT to stop)")
+    while (true) {
         val query = readln()
+        if (query == "EXIT") {
+            break
+        }
+
         println("Searching for '$query'")
-        var results: List<SearchResult>
+        val results: List<SearchResult>
         val searchTime = measureTimeMillis { results = index.search(query) }
         println("${results.size} results, $searchTime ms")
 
         results.forEachIndexed { i, result ->
-            val book = books[result.documentId]!!
+            val book = data[result.documentId]!!
             println("$i\t(${result.score}, ${result.matchTerm})\t${book.title}, ${book.author}")
         }
-    } while (query != "EXIT")
+    }
+}
+
+fun autocomplete(index: FtsIndex<Book>, data: Map<Int, Book>) {
+    println("Autocomplete Suggestions")
+    println("Enter search query (EXIT TO STOP)")
+    while (true) {
+        val query = readln()
+        if (query == "EXIT") {
+            break
+        }
+
+        println("Suggestions for '$query'")
+        val suggestions: List<AutocompleteSuggestion>
+        val searchTime = measureTimeMillis { suggestions = index.autocomplete(query) }
+        println("($searchTime ms) ${suggestions.joinToString { it.suggestion }}")
+    }
 }

core/src/test/kotlin/FtsIndexTest.kt

@@ -125,4 +125,21 @@ class FtsIndexTest : DescribeSpec({
             fts.search("test") shouldHaveSize 0
         }
     }
+
+    context("Autocomplete suggestions") {
+        it("should return zero results if there are no matches") {
+            val index = useFts<Book>()
+            val results = index.autocomplete("foo")
+            results shouldHaveSize 0
+        }
+
+        it("should return matching results") {
+            val index = useFts<Sentence>()
+            index.add(Sentence(0, "football"))
+            index.add(Sentence(1, "foil"))
+
+            val results = index.autocomplete("fo")
+            results shouldHaveSize 2
+        }
+    }
 })