From 1a084543513b548a7d7778bfa415a5e75c92ce8f Mon Sep 17 00:00:00 2001 From: Menci Date: Tue, 6 Jan 2026 22:33:54 +0800 Subject: [PATCH] feat: bundle documents in index optionally --- .../Types/CompressedInvertedIndex.cs | 2 +- .../InvertedIndexBuilder.cs | 16 +++++- .../InvertedIndexLoader.cs | 5 +- .../MaigoLabs.NeedLe.Tests/E2E/SearchTests.cs | 50 +++++++++++++++++++ packages/needle/src/common/types.ts | 2 +- packages/needle/src/e2e/search.test.ts | 30 +++++++++++ packages/needle/src/indexer/inverted-index.ts | 13 +++-- .../needle/src/searcher/inverted-index.ts | 5 +- 8 files changed, 112 insertions(+), 11 deletions(-) diff --git a/dotnet/MaigoLabs.NeedLe.Common/Types/CompressedInvertedIndex.cs b/dotnet/MaigoLabs.NeedLe.Common/Types/CompressedInvertedIndex.cs index 43f76ad..f3546b4 100644 --- a/dotnet/MaigoLabs.NeedLe.Common/Types/CompressedInvertedIndex.cs +++ b/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[] tokenReferences { get; set; } // tokenId -> [documentId, start1, end1, start2, end2, ...] public required CompressedInvertedIndexTries tries { get; set; } diff --git a/dotnet/MaigoLabs.NeedLe.Indexer/InvertedIndexBuilder.cs b/dotnet/MaigoLabs.NeedLe.Indexer/InvertedIndexBuilder.cs index 51bdd5a..2726888 100644 --- a/dotnet/MaigoLabs.NeedLe.Indexer/InvertedIndexBuilder.cs +++ b/dotnet/MaigoLabs.NeedLe.Indexer/InvertedIndexBuilder.cs @@ -6,6 +6,15 @@ using MaigoLabs.NeedLe.Indexer.Trie; namespace MaigoLabs.NeedLe.Indexer; +public class InvertedIndexBuilderOptions +{ + /// + /// If false, the documents will not be bundled with the inverted index. You must pass documents explicitly when loading the index. + /// Defaults to true. + /// + public bool BundleDocuments { get; set; } = true; +} + public static class InvertedIndexBuilder { private static TrieNode BuildTypedTrie(IEnumerable tokenDefinitions, Func 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())], tries = new CompressedInvertedIndexTries diff --git a/dotnet/MaigoLabs.NeedLe.Searcher/InvertedIndexLoader.cs b/dotnet/MaigoLabs.NeedLe.Searcher/InvertedIndexLoader.cs index d08f5e8..425e317 100644 --- a/dotnet/MaigoLabs.NeedLe.Searcher/InvertedIndexLoader.cs +++ b/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); diff --git a/dotnet/MaigoLabs.NeedLe.Tests/E2E/SearchTests.cs b/dotnet/MaigoLabs.NeedLe.Tests/E2E/SearchTests.cs index cc7dedd..66d8e51 100644 --- a/dotnet/MaigoLabs.NeedLe.Tests/E2E/SearchTests.cs +++ b/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(() => InvertedIndexLoader.Load(compressed)); + } +} diff --git a/packages/needle/src/common/types.ts b/packages/needle/src/common/types.ts index f88482f..58151e4 100644 --- a/packages/needle/src/common/types.ts +++ b/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: { diff --git a/packages/needle/src/e2e/search.test.ts b/packages/needle/src/e2e/search.test.ts index 26787c3..b089d31 100644 --- a/packages/needle/src/e2e/search.test.ts +++ b/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(); + }); + }); +}); diff --git a/packages/needle/src/indexer/inverted-index.ts b/packages/needle/src/indexer/inverted-index.ts index 83ac851..264ac0b 100644 --- a/packages/needle/src/indexer/inverted-index.ts +++ b/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: { diff --git a/packages/needle/src/searcher/inverted-index.ts b/packages/needle/src/searcher/inverted-index.ts index 11731f0..dd0da1d 100644 --- a/packages/needle/src/searcher/inverted-index.ts +++ b/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);