feat: bundle documents in index optionally

dotnet/MaigoLabs.NeedLe.Common/Types/CompressedInvertedIndex.cs

@@ -6,7 +6,7 @@ namespace MaigoLabs.NeedLe.Common.Types;
 
 public class CompressedInvertedIndex
 {
-    public required string[] documents { get; set; }
+    public required string[]? documents { get; set; }
     public required int[] tokenTypes { get; set; } // Use int values here instead of TokenType enum to avoid JSON serialization issues.
     public required List<int[]>[] tokenReferences { get; set; } // tokenId -> [documentId, start1, end1, start2, end2, ...]
     public required CompressedInvertedIndexTries tries { get; set; }

dotnet/MaigoLabs.NeedLe.Indexer/InvertedIndexBuilder.cs

@@ -6,6 +6,15 @@ using MaigoLabs.NeedLe.Indexer.Trie;
 
 namespace MaigoLabs.NeedLe.Indexer;
 
+public class InvertedIndexBuilderOptions
+{
+    /// <summary>
+    /// If false, the documents will not be bundled with the inverted index. You must pass documents explicitly when loading the index.
+    /// Defaults to true.
+    /// </summary>
+    public bool BundleDocuments { get; set; } = true;
+}
+
 public static class InvertedIndexBuilder
 {
     private static TrieNode BuildTypedTrie(IEnumerable<TokenDefinition> tokenDefinitions, Func<TokenType, bool> typePredicate) =>
@@ -13,7 +22,10 @@ public static class InvertedIndexBuilder
             .Where(token => typePredicate(token.Type))
             .Select(token => (token.Id, CodePoints: token.Text.ToCodePoints())));
 
-    public static CompressedInvertedIndex BuildInvertedIndex(string[] documents, TokenizerOptions? tokenizerOptions = null)
+    public static CompressedInvertedIndex BuildInvertedIndex(
+        string[] documents,
+        TokenizerOptions? tokenizerOptions = null,
+        InvertedIndexBuilderOptions? invertedIndexBuilderOptions = null)
     {
         var tokenizer = new Tokenizer(tokenizerOptions);
         var documentTokens = documents.Select(tokenizer.Tokenize).ToArray();
@@ -27,7 +39,7 @@ public static class InvertedIndexBuilder
 
         var invertedIndex = new CompressedInvertedIndex
         {
-            documents = documents,
+            documents = (invertedIndexBuilderOptions?.BundleDocuments ?? true) ? documents : null,
             tokenTypes = [.. tokenDefinitions.Select(token => (int)token.Type)],
             tokenReferences = [.. tokenDefinitions.Select(_ => new List<int[]>())],
             tries = new CompressedInvertedIndexTries

dotnet/MaigoLabs.NeedLe.Searcher/InvertedIndexLoader.cs

@@ -33,9 +33,10 @@ public class LoadedInvertedIndex
 
 public class InvertedIndexLoader
 {
-    public static LoadedInvertedIndex Load(CompressedInvertedIndex compressed)
+    public static LoadedInvertedIndex Load(CompressedInvertedIndex compressed, string[]? documents = null)
     {
-        var documents = compressed.documents;
+        documents ??= compressed.documents;
+        if (documents == null) throw new ArgumentException("Loading an inverted index without documents bundled requires documents to be provided explicitly.");
         var documentCodePoints = documents.Select(document => document.ToCodePoints().ToArray()).ToArray();
 
         var romajiTrie = TrieDeserializer.Deserialize(compressed.tries.romaji);

dotnet/MaigoLabs.NeedLe.Tests/E2E/SearchTests.cs

@@ -89,3 +89,53 @@ public sealed class Search_MatchesRomajiInputToKanaDocumentsTest : NeedleTestBas
     }
 }
 
+public sealed class Search_BundleDocumentsOption_WorksWhenNotBundledTest : NeedleTestBase
+{
+    private static readonly string[] TestDocuments =
+    [
+        "ミーティア",
+        "エンドマークに希望と涙を添えて",
+        "宵の鳥",
+        "僕の和風本当上手",
+    ];
+
+    [Fact]
+    public void Execute()
+    {
+        var compressed = InvertedIndexBuilder.BuildInvertedIndex(
+            TestDocuments,
+            TokenizerOptions,
+            new InvertedIndexBuilderOptions { BundleDocuments = false });
+
+        // Documents should not be in the compressed index
+        Assert.Null(compressed.documents);
+
+        // Load with documents provided explicitly
+        var invertedIndex = InvertedIndexLoader.Load(compressed, TestDocuments);
+
+        var results = InvertedIndexSearcher.Search(invertedIndex, "yoi");
+        Assert.Contains("宵の鳥", results.Select(r => r.DocumentText));
+    }
+}
+
+public sealed class Search_BundleDocumentsOption_ThrowsWhenNoneProvidedTest : NeedleTestBase
+{
+    private static readonly string[] TestDocuments =
+    [
+        "ミーティア",
+        "エンドマークに希望と涙を添えて",
+        "宵の鳥",
+        "僕の和風本当上手",
+    ];
+
+    [Fact]
+    public void Execute()
+    {
+        var compressed = InvertedIndexBuilder.BuildInvertedIndex(
+            TestDocuments,
+            TokenizerOptions,
+            new InvertedIndexBuilderOptions { BundleDocuments = false });
+
+        Assert.Throws<ArgumentException>(() => InvertedIndexLoader.Load(compressed));
+    }
+}

packages/needle/src/common/types.ts

@@ -20,7 +20,7 @@ export interface OffsetSpan {
 }
 
 export type CompressedInvertedIndex = {
-  documents: string[];
+  documents?: string[];
   tokenTypes: TokenType[];
   tokenReferences: number[][][]; // tokenId -> [documentId, start1, end1, start2, end2, ...][]
   tries: {

packages/needle/src/e2e/search.test.ts

@@ -71,3 +71,33 @@ describe('search', () => {
     expect(matchedTexts).toContain('宵の鳥');
   });
 });
+
+describe('search options', () => {
+  const testDocuments = [
+    'ミーティア',
+    'エンドマークに希望と涙を添えて',
+    '宵の鳥',
+    '僕の和風本当上手',
+  ];
+
+  describe('bundleDocuments option', () => {
+    it('should work when documents are not bundled and provided at load time', () => {
+      const compressed = buildInvertedIndex(testDocuments, { kuromoji, bundleDocuments: false });
+
+      // Documents should not be in the compressed index
+      expect(compressed.documents).toBeUndefined();
+
+      // Load with documents provided explicitly
+      const invertedIndex = loadInvertedIndex(compressed, testDocuments);
+
+      const results = searchInvertedIndex(invertedIndex, 'yoi');
+      expect(results.map(r => r.documentText)).toContain('宵の鳥');
+    });
+
+    it('should throw when loading without documents and none provided', () => {
+      const compressed = buildInvertedIndex(testDocuments, { kuromoji, bundleDocuments: false });
+
+      expect(() => loadInvertedIndex(compressed)).toThrow();
+    });
+  });
+});

packages/needle/src/indexer/inverted-index.ts

@@ -7,8 +7,15 @@ import { TokenType } from '../common/types';
 const buildTypedTrie = (tokens: TokenDefinition[], typePredicate: (tokenType: TokenType) => boolean) =>
   buildTrie(tokens.filter(token => typePredicate(token.type)).map(token => [token.id, token.text]));
 
-export const buildInvertedIndex = (documents: string[], tokenizerOptions: TokenizerOptions) => {
+export const buildInvertedIndex = (documents: string[], options: TokenizerOptions & {
-  const tokenizer = createTokenizer(tokenizerOptions);
+  /**
+   * If false, the documents will not be bundled with the inverted index. You must pass documents explicitly when loading the index.
+   *
+   * @default true
+   */
+  bundleDocuments?: boolean;
+}) => {
+  const tokenizer = createTokenizer(options);
   const documentTokens = documents.map(document => tokenizer.tokenize(document));
 
   const tokenDefinitions = [...tokenizer.tokens.values()];
@@ -19,7 +26,7 @@ export const buildInvertedIndex = (documents: string[], tokenizerOptions: Tokeni
   graftTriePaths(kanaRoot, NORMALIZE_RULES_KANA_DAKUTEN);
 
   const invertedIndex: CompressedInvertedIndex = {
-    documents,
+    documents: (options.bundleDocuments ?? true) ? documents : undefined,
     tokenTypes: tokenDefinitions.map(token => token.type),
     tokenReferences: Array.from({ length: tokenDefinitions.length }, () => []),
     tries: {

packages/needle/src/searcher/inverted-index.ts

@@ -28,8 +28,9 @@ export interface LoadedInvertedIndex {
   };
 }
 
-export const loadInvertedIndex = (compressed: CompressedInvertedIndex): LoadedInvertedIndex => {
+export const loadInvertedIndex = (compressed: CompressedInvertedIndex, documents?: string[]): LoadedInvertedIndex => {
-  const documents = compressed.documents;
+  documents ??= compressed.documents;
+  if (!documents) throw new Error('Loading an inverted index without documents bundled requires documents to be provided explicitly.');
   const documentCodePoints = documents.map(document => [...document]);
 
   const romajiTrie = deserializeTrie(compressed.tries.romaji);