From 03380ca9dec8796ca778bc813d585e912e04c7f3 Mon Sep 17 00:00:00 2001
From: Abdullah Atta <abdullahatta@streetwriters.co>
Date: Mon, 19 Feb 2024 16:34:39 +0500
Subject: [PATCH] core: improve notes lookup performance

---
 packages/core/__benches__/lookup.bench.ts | 65 +++++++++++++++++++++++
 packages/core/src/api/lookup.ts           | 59 +++++++++++---------
 2 files changed, 98 insertions(+), 26 deletions(-)
 create mode 100644 packages/core/__benches__/lookup.bench.ts

diff --git a/packages/core/__benches__/lookup.bench.ts b/packages/core/__benches__/lookup.bench.ts
new file mode 100644
index 000000000..76765a797
--- /dev/null
+++ b/packages/core/__benches__/lookup.bench.ts
@@ -0,0 +1,65 @@
+/*
+This file is part of the Notesnook project (https://notesnook.com/)
+
+Copyright (C) 2023 Streetwriters (Private) Limited
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <http://www.gnu.org/licenses/>.
+*/
+
+import { bench, describe } from "vitest";
+import { databaseTest } from "../__tests__/utils";
+import Database from "../src/api";
+
+async function addNotes(db: Database) {
+  const words = [
+    "title",
+    "who",
+    "what",
+    "never",
+    "works",
+    "here",
+    "say",
+    "there"
+  ];
+  for (let i = 0; i < 10000; ++i) {
+    await db.notes.add({
+      title: new Array(10)
+        .fill(0)
+        .map(() => words[getRandom(0, words.length)])
+        .join(" "),
+      content: {
+        type: "tiptap",
+        data: new Array(50)
+          .fill(0)
+          .map(() => words[getRandom(0, words.length)])
+          .join(" ")
+      }
+    });
+    if (i % 100 === 0) console.log(i);
+  }
+  console.log("DONE");
+}
+
+describe("lookup", async () => {
+  const db = await databaseTest();
+  await addNotes(db);
+
+  bench("using multiple queries", async () => {
+    await db.lookup.notes("never works here say").ids();
+  });
+});
+
+function getRandom(min: number, max: number) {
+  return Math.round(Math.random() * (max - min) + min);
+}
diff --git a/packages/core/src/api/lookup.ts b/packages/core/src/api/lookup.ts
index cc2ee075f..75e97e169 100644
--- a/packages/core/src/api/lookup.ts
+++ b/packages/core/src/api/lookup.ts
@@ -20,7 +20,7 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 import { match } from "fuzzyjs";
 import Database from ".";
 import { Item, Note, TrashItem } from "../types";
-import { DatabaseSchema, RawDatabaseSchema, isFalse } from "../database";
+import { DatabaseSchema, RawDatabaseSchema } from "../database";
 import { AnyColumnWithTable, Kysely, sql } from "kysely";
 import { FilteredSelector } from "../database/sql-collection";
 import { VirtualizedGrouping } from "../utils/virtualized-grouping";
@@ -45,34 +45,41 @@ export default class Lookup {
 
       const db = this.db.sql() as unknown as Kysely<RawDatabaseSchema>;
       query = query.replace(/"/, '""');
-      const result = await db
-        .with("matching", (eb) =>
+
+      const excludedIds = this.db.trash.cache.notes;
+      const results = await db
+        .selectFrom((eb) =>
           eb
-            .selectFrom("content_fts")
-            .where("data", "match", `"${query}"`)
-            .select(["noteId as id", "rank"])
-            .unionAll(
-              eb
-                .selectFrom("notes_fts")
-                .where("title", "match", `"${query}"`)
-                // add 10 weight to title
-                .select(["id", sql.raw<number>(`rank * 10`).as("rank")])
+            .selectFrom("notes_fts")
+            .$if(!!notes, (eb) =>
+              eb.where("id", "in", notes!.filter.select("id"))
             )
+            .$if(excludedIds.length > 0, (eb) =>
+              eb.where("id", "not in", excludedIds)
+            )
+            .where("title", "match", `"${query}"`)
+            .select(["id", sql<number>`rank * 10`.as("rank")])
+            .unionAll((eb) =>
+              eb
+                .selectFrom("content_fts")
+                .$if(!!notes, (eb) =>
+                  eb.where("id", "in", notes!.filter.select("id"))
+                )
+                .$if(excludedIds.length > 0, (eb) =>
+                  eb.where("id", "not in", excludedIds)
+                )
+                .where("data", "match", `"${query}"`)
+                .select(["noteId as id", "rank"])
+                .$castTo<{ id: string; rank: number }>()
+            )
+            .as("results")
         )
-        .selectFrom("notes")
-        .$if(!!notes, (eb) =>
-          eb.where("notes.id", "in", notes!.filter.select("id"))
-        )
+        .select(["results.id"])
+        .groupBy("results.id")
+        .orderBy(sql`SUM(results.rank)`, "asc")
         .$if(!!limit, (eb) => eb.limit(limit!))
-        .where(isFalse("notes.deleted"))
-        .where(isFalse("notes.dateDeleted"))
-        .innerJoin("matching", (eb) =>
-          eb.onRef("notes.id", "==", "matching.id")
-        )
-        .orderBy("matching.rank desc")
-        .select(["notes.id"])
         .execute();
-      return result.map((id) => id.id);
+      return results.map((r) => r.id);
     }, notes || this.db.notes.all);
   }
 
@@ -212,10 +219,10 @@ export default class Lookup {
       this.db.options.batchSize,
       () => Promise.resolve(ids),
       async (start, end) => {
-        const items = await selector.items(ids);
+        const items = await selector.records(ids);
         return {
           ids: ids.slice(start, end),
-          items: items.slice(start, end)
+          items: Object.values(items).slice(start, end)
         };
       }
     );