MaigoLabs/needLe
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: pre-filter and custom comparer for documents

Browse Source
This commit is contained in:
Menci
2026-01-06 22:34:25 +08:00
parent 1a08454351
commit 63dcd3fb75
4 changed files with 134 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dotnet/MaigoLabs.NeedLe.Searcher/InvertedIndexSearcher.cs
+27 -5
View File
@@ -24,6 +24,21 @@ public class SearchResult
public required int MatchRatioLevel { get; set; }
}
public class InvertedIndexSearcherOptions
{
/// <summary>
/// Called when all other comparisons are equal.
/// <code>(int documentIdA, int documentIdB) => int</code>
/// </summary>
public Func<int, int, int>? NextComparer { get; set; }
/// <summary>
/// If return false for a document, it will be excluded from the final results.
/// <code>(int documentId) => bool</code>
/// </summary>
public Func<int, bool>? FilterDocument { get; set; }
}
public static class InvertedIndexSearcher
{
public abstract class ComparableStateBase<T> : IComparable<T>
@@ -36,7 +51,6 @@ public static class InvertedIndexSearcher
protected virtual SearchResultToken? GetLastToken() => null; // Not on intermediate results
protected virtual int? GetMatchRatioLevel() => null; // Not on intermediate/candidate results
protected abstract double GetMatchRatio();
protected virtual int FallbackCompareTo(T other) => 0; // Called when all other comparisons are equal
public int CompareTo(T other)
{
@@ -73,7 +87,7 @@ public static class InvertedIndexSearcher
double aMatchRatio = GetMatchRatio(), bMatchRatio = other.GetMatchRatio();
if (aMatchRatio != bMatchRatio) return bMatchRatio < aMatchRatio ? -1 : bMatchRatio > aMatchRatio ? 1 : 0;
return FallbackCompareTo(other);
return 0;
}
}
@@ -123,7 +137,7 @@ public static class InvertedIndexSearcher
protected override SearchResultToken? GetLastToken() => Result.Tokens[^1];
protected override double GetMatchRatio() => Result.MatchRatio;
protected override int? GetMatchRatioLevel() => Result.MatchRatioLevel;
protected override int FallbackCompareTo(FinalResult other) => string.Compare(Result.DocumentText, other.Result.DocumentText, StringComparison.InvariantCulture);
// protected override int FallbackCompareTo(FinalResult other) => string.Compare(Result.DocumentText, other.Result.DocumentText, StringComparison.InvariantCulture);
}
private static bool IsIgnorableCodePoint(int codePoint) => CommonUtils.IsWhitespace(codePoint) || codePoint == 0x3099 || codePoint == 0x309A;
@@ -151,7 +165,7 @@ public static class InvertedIndexSearcher
private static bool HasNonEmptyCharacters(int[] documentCodePoints, int start, int end) =>
start != end && !documentCodePoints.Skip(start).Take(end - start).All(CommonUtils.IsWhitespace);
public static SearchResult[] Search(LoadedInvertedIndex invertedIndex, string text)
public static SearchResult[] Search(LoadedInvertedIndex invertedIndex, string text, InvertedIndexSearcherOptions? options = null)
{
var documents = invertedIndex.Documents;
var documentCodePoints = invertedIndex.DocumentCodePoints;
@@ -183,6 +197,7 @@ public static class InvertedIndexSearcher
];
foreach (var tokenId in matchingTokenIds) foreach (var reference in tokenDefinitions[tokenId].References)
{
if (options?.FilterDocument != null && !options.FilterDocument(reference.DocumentId)) continue;
var isTokenPrefixMatching = !romajiNode.IsTokenExactMatch(tokenId) && !kanaNode.IsTokenExactMatch(tokenId) && !otherNode.IsTokenExactMatch(tokenId);
var previousMatchesOfDocument = l != 0 && dp[l - 1].TryGetValue(reference.DocumentId, out var previousMatches) ? previousMatches : null;
if (l != 0 && previousMatchesOfDocument == null) continue;
@@ -265,6 +280,13 @@ public static class InvertedIndexSearcher
MatchRatioLevel = matchRatioLevel,
}
};
}).OrderBy(result => result).Select(result => result.Result).ToArray();
}).OrderBy(result => result, Comparer<FinalResult>.Create((a, b) =>
{
var compareResult = a.CompareTo(b);
if (compareResult != 0) return compareResult;
return options?.NextComparer == null
? string.Compare(a.Result.DocumentText, b.Result.DocumentText, StringComparison.InvariantCulture)
: options.NextComparer(a.Result.DocumentId, b.Result.DocumentId);
})).Select(result => result.Result).ToArray();
}
}
dotnet/MaigoLabs.NeedLe.Tests/E2E/SearchTests.cs
+51
View File
@@ -139,3 +139,54 @@ public sealed class Search_BundleDocumentsOption_ThrowsWhenNoneProvidedTest : Ne
Assert.Throws<ArgumentException>(() => InvertedIndexLoader.Load(compressed));
}
}
public sealed class Search_FilterDocumentOption_ExcludesFilteredDocumentsTest : NeedleTestBase
{
private static readonly string[] TestDocuments =
[
"ミーティア",
"エンドマークに希望と涙を添えて",
"宵の鳥",
"僕の和風本当上手",
];
[Fact]
public void Execute()
{
var compressed = InvertedIndexBuilder.BuildInvertedIndex(TestDocuments, TokenizerOptions);
var invertedIndex = InvertedIndexLoader.Load(compressed);
// Search without filter - should find "宵の鳥" (documentId 2)
var resultsWithoutFilter = InvertedIndexSearcher.Search(invertedIndex, "yoi");
Assert.Contains("宵の鳥", resultsWithoutFilter.Select(r => r.DocumentText));
// Search with filter excluding documentId 2
var resultsWithFilter = InvertedIndexSearcher.Search(invertedIndex, "yoi", new InvertedIndexSearcherOptions
{
FilterDocument = id => id != 2
});
Assert.DoesNotContain("宵の鳥", resultsWithFilter.Select(r => r.DocumentText));
}
}
public sealed class Search_NextComparerOption_UsesCustomComparerTest : NeedleTestBase
{
[Fact]
public void Execute()
{
// Create documents that would have similar match scores
var similarDocs = new[] { "テストA", "テストB", "テストC" };
var compressed = InvertedIndexBuilder.BuildInvertedIndex(similarDocs, TokenizerOptions);
var invertedIndex = InvertedIndexLoader.Load(compressed);
// Search with reverse order comparer
var results = InvertedIndexSearcher.Search(invertedIndex, "テスト", new InvertedIndexSearcherOptions
{
NextComparer = (a, b) => b - a // Reverse by documentId
});
// Should be in reverse documentId order (2, 1, 0) when other criteria equal
Assert.Equal([2, 1, 0], results.Select(r => r.DocumentId).ToArray());
}
}
packages/needle/src/e2e/search.test.ts
+34
View File
@@ -100,4 +100,38 @@ describe('search options', () => {
expect(() => loadInvertedIndex(compressed)).toThrow();
});
});
describe('filterDocument option', () => {
it('should exclude filtered documents from results', () => {
const compressed = buildInvertedIndex(testDocuments, { kuromoji });
const invertedIndex = loadInvertedIndex(compressed);
// Search without filter - should find "宵の鳥" (documentId 2)
const resultsWithoutFilter = searchInvertedIndex(invertedIndex, 'yoi');
expect(resultsWithoutFilter.map(r => r.documentText)).toContain('宵の鳥');
// Search with filter excluding documentId 2
const resultsWithFilter = searchInvertedIndex(invertedIndex, 'yoi', {
filterDocument: id => id !== 2,
});
expect(resultsWithFilter.map(r => r.documentText)).not.toContain('宵の鳥');
});
});
describe('nextComparer option', () => {
it('should use custom comparer for final sorting when other criteria are equal', () => {
// Create documents that would have similar match scores
const similarDocs = ['テストA', 'テストB', 'テストC'];
const compressed = buildInvertedIndex(similarDocs, { kuromoji });
const invertedIndex = loadInvertedIndex(compressed);
// Search with reverse order comparer
const results = searchInvertedIndex(invertedIndex, 'テスト', {
nextComparer: (a, b) => b - a, // Reverse by documentId
});
// Should be in reverse documentId order (2, 1, 0) when other criteria equal
expect(results.map(r => r.documentId)).toEqual([2, 1, 0]);
});
});
});
packages/needle/src/searcher/search.ts
+22 -3
View File
@@ -133,12 +133,24 @@ const compareFinalResult = getComparerForTraits<SearchResult>({
getLastToken: state => state.tokens[state.tokens.length - 1]!,
getMatchRatio: state => state.matchRatio,
getMatchRatioLevel: state => Math.round(state.matchRatio * 5),
nextComparer: (a, b) => a.documentText === b.documentText ? 0 : a.documentText < b.documentText ? -1 : 1,
});
const hasNonEmptyCharacters = (documentCodePoints: string[], start: number, end: number) => start !== end && !documentCodePoints.slice(start, end).every(char => /\s/.test(char));
export const searchInvertedIndex = (invertedIndex: LoadedInvertedIndex, text: string): SearchResult[] => {
export const searchInvertedIndex = (
invertedIndex: LoadedInvertedIndex,
text: string,
options?: {
/**
* Called when all other comparisons are equal.
*/
nextComparer?: (documentIdA: number, documentIdB: number) => number;
/**
* If return falsy value for a document, it will be excluded from the final results.
*/
filterDocument?: (documentId: number) => unknown;
},
): SearchResult[] => {
const { documents, documentCodePoints, tokenDefinitions, tries } = invertedIndex;
const codePoints = [...toKatakana(normalizeByCodePoint(text))];
@@ -162,6 +174,7 @@ export const searchInvertedIndex = (invertedIndex: LoadedInvertedIndex, text: st
...getTrieNodeTokenIds(otherNode, reachingInputEnd),
]);
for (const tokenId of matchingTokenIds) for (const { documentId, offsets } of tokenDefinitions[tokenId]!.references) {
if (options?.filterDocument && !options.filterDocument(documentId)) continue;
const isTokenPrefixMatching = !romajiNode?.tokenIds.includes(tokenId) && !kanaNode?.tokenIds.includes(tokenId) && !otherNode?.tokenIds.includes(tokenId);
const previousMatchesOfDocument = dp[l - 1]?.get(documentId);
if (l !== 0 && !previousMatchesOfDocument) continue;
@@ -231,7 +244,13 @@ export const searchInvertedIndex = (invertedIndex: LoadedInvertedIndex, text: st
matchRatio,
matchRatioLevel,
};
}).sort(compareFinalResult);
}).sort((a, b) => {
const compareResult = compareFinalResult(a, b);
if (compareResult !== 0) return compareResult;
return options?.nextComparer
? options.nextComparer(a.documentId, b.documentId)
: a.documentText === b.documentText ? 0 : a.documentText < b.documentText ? -1 : 1;
});
};
// For debugging