feat: bundle documents in index optionally

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) => {
-  const tokenizer = createTokenizer(tokenizerOptions);
+export const buildInvertedIndex = (documents: string[], options: 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 => {
-  const documents = compressed.documents;
+export const loadInvertedIndex = (compressed: CompressedInvertedIndex, documents?: string[]): LoadedInvertedIndex => {
+  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);