feat(lyra): adds stemming capabilities during indexing and search

packages/lyra/package.json

@@ -24,10 +24,12 @@
   "types": "./dist/esm/lyra.d.ts",
   "dependencies": {
     "fastq": "^1.13.0",
-    "nanoid": "3.3.4"
+    "nanoid": "3.3.4",
+    "natural": "^5.2.2"
   },
   "devDependencies": {
     "@types/jest": "^27.5.0",
+    "@types/natural": "^5.1.1",
     "jest": "^28.1.0",
     "ts-jest": "^28.0.2",
     "typescript": "^4.6.4"

packages/lyra/src/lyra.ts

@@ -5,6 +5,7 @@ import { Trie } from "./prefix-tree/trie";
 import * as ERRORS from "./errors";
 import { tokenize } from "./tokenizer";
 import { formatNanoseconds, getNanosecondsTime } from "./utils";
+import { Language } from "./stemmer";
 
 export type PropertyType = "string" | "number" | "boolean";
 
@@ -30,6 +31,7 @@ type LyraIndex = Map<string, Trie>;
 type QueueDocParams = {
   id: string;
   doc: object;
+  language: Language;
 };
 
 type SearchResult = Promise<{
@@ -183,7 +185,10 @@ export class Lyra {
     return ids;
   }
 
-  public async insert(doc: object): Promise<{ id: string }> {
+  public async insert(
+    doc: object,
+    language: Language = "english"
+  ): Promise<{ id: string }> {
     const id = nanoid();
 
     if (!(await this.checkInsertDocSchema(doc))) {
@@ -193,12 +198,13 @@ export class Lyra {
     await this.queue.push({
       id,
       doc,
+      language,
     });
 
     return { id };
   }
 
-  private async _insert({ doc, id }: QueueDocParams): Promise<void> {
+  private async _insert({ doc, id, language }: QueueDocParams): Promise<void> {
     const index = this.index;
     this.docs.set(id, doc);
 
@@ -210,7 +216,7 @@ export class Lyra {
           // recursiveTrieInsertion((doc as any)[key]);
         } else if (typeof (doc as any)[key] === "string") {
           const requestedTrie = index.get(key);
-          const tokens = tokenize((doc as any)[key]);
+          const tokens = tokenize((doc as any)[key], language);
 
           for (const token of tokens) {
             requestedTrie?.insert(token, id);

packages/lyra/src/stemmer.ts

@@ -0,0 +1,57 @@
+import natural from "natural";
+
+export type Language = typeof SUPPORTED_LANGUAGES[number];
+
+export const SUPPORTED_LANGUAGES = [
+  "dutch",
+  "english",
+  "french",
+  "italian",
+  "norwegian",
+  "portugese",
+  "russian",
+  "spanish",
+  "swedish",
+] as const;
+
+export function stemArray(input: string[], language: Language): string[] {
+  let stemmer;
+
+  switch (language) {
+    case "dutch":
+      stemmer = natural.PorterStemmerNl;
+      break;
+    case "english":
+      stemmer = natural.PorterStemmer;
+      break;
+    case "french":
+      stemmer = natural.PorterStemmerFr;
+      break;
+    case "italian":
+      stemmer = natural.PorterStemmerIt;
+      break;
+    case "norwegian":
+      stemmer = natural.PorterStemmerNo;
+      break;
+    case "portugese":
+      stemmer = natural.PorterStemmerPt;
+      break;
+    case "russian":
+      stemmer = natural.PorterStemmerRu;
+      break;
+    case "spanish":
+      stemmer = natural.PorterStemmerEs;
+      break;
+    case "swedish":
+      stemmer = natural.PorterStemmerSv;
+      break;
+    default:
+      break;
+  }
+
+  if (stemmer) {
+    return input.map(stemmer.stem);
+  }
+
+  return input;
+}

packages/lyra/src/tokenizer.ts

@@ -1,9 +1,12 @@
-export function tokenize(input: string): Set<string> {
-  return new Set(
-    input
-      .toLowerCase()
-      .replace(/[^a-z0-9 -]/g, "")
-      .split(" ")
-      .filter(Boolean)
-  );
+import { WordTokenizer } from "natural";
+import { Language, stemArray } from "./stemmer";
+
+const tokenizer = new WordTokenizer();
+
+export function tokenize(
+  input: string,
+  language: Language = "english"
+): Set<string> {
+  const tokens = tokenizer.tokenize(input);
+  return new Set(stemArray(tokens, language));
 }