improve recall on short documents

lofcz · lofcz · commit 37753aa91f7b · 2025-11-28T11:48:43.000+01:00
diff --git a/src/Infidex.Tests/FuzzyRegressionTests.cs b/src/Infidex.Tests/FuzzyRegressionTests.cs
@@ -0,0 +1,60 @@
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Infidex.Core;
+using Infidex.Api;
+using System.Linq;
+using System;
+
+namespace Infidex.Tests;
+
+[TestClass]
+public class FuzzyRegressionTests
+{
+    private SearchEngine _engine = null!;
+    
+    [TestInitialize]
+    public void Setup()
+    {
+        _engine = SearchEngine.CreateDefault();
+        
+        var docs = new[]
+        {
+            new Document(1L, "The Mat"),
+            new Document(2L, "The Matrix"),
+            new Document(3L, "The Matriarx"),
+            new Document(4L, "The Match"),
+            new Document(5L, "The Meatrix")
+        };
+        
+        _engine.IndexDocuments(docs);
+    }
+    
+    [TestMethod]
+    public void Search_TheMatrx_RanksMatrixAboveMat()
+    {
+        // Query: "the matrx"
+        // Target: "The Matrix" (Doc 2) should be higher than "The Mat" (Doc 1)
+        // "matrx" is a typo for "matrix".
+        // "The Matriarx" (Doc 3) contains "matrx" exactly? If tokenizer splits it? 
+        // "Matriarx" -> "matriarx". "matrx" is not "matriarx".
+        // Unless "matrx" is in "Matriarx" as n-gram? Yes.
+        // But "The Matrix" should beat "The Mat" because "matrx" -> "matrix" (fuzzy) is a whole-word match
+        // whereas "The Mat" is only a partial n-gram match (which should be suppressed!).
+        
+        var result = _engine.Search(new Query("the matrx", 10));
+        
+        Console.WriteLine("Results for 'the matrx':");
+        foreach (var r in result.Records)
+        {
+            Console.WriteLine($"[{r.Score:F1}] Doc {r.DocumentId}");
+        }
+
+        var doc1 = result.Records.FirstOrDefault(r => r.DocumentId == 1L); // The Mat
+        var doc2 = result.Records.FirstOrDefault(r => r.DocumentId == 2L); // The Matrix
+        
+        Assert.IsNotNull(doc2, "The Matrix should be found");
+     
+        // The Matrix should score higher than The Mat
+        Assert.IsTrue(doc2.Score > doc1.Score, 
+            $"The Matrix ({doc2.Score}) should rank higher than The Mat ({doc1.Score})");
+    }
+}
diff --git a/src/Infidex/Coverage/CoverageEngine.cs b/src/Infidex/Coverage/CoverageEngine.cs
@@ -289,7 +289,7 @@ private CoverageResult CalculateCoverageInternal(string query, string documentTe
                 PrefixSuffixMatcher.Match(ref state);
 
             if (_setup.CoverFuzzyWords && qCount > 0 && !FuzzyWordMatcher.AllTermsFullyMatched(ref state))
-                FuzzyWordMatcher.Match(ref state, _setup.MinWordSize, _setup.LevenshteinMaxWordSize);
+                FuzzyWordMatcher.Match(ref state, _setup);
         }
         finally
         {
diff --git a/src/Infidex/Coverage/CoverageSetup.cs b/src/Infidex/Coverage/CoverageSetup.cs
@@ -14,6 +14,24 @@ public class CoverageSetup
     /// Maximum word size for Levenshtein fuzzy matching
     /// </summary>
     public int LevenshteinMaxWordSize { get; set; } = 20;
+
+    /// <summary>
+    /// Maximum number of typographical errors (0, 1 or 2) that would be tolerated.
+    /// Default: 2.
+    /// </summary>
+    public int NumTypos { get; set; } = 2;
+
+    /// <summary>
+    /// Minimum word length for 1-typo correction to be applied.
+    /// Default: 3.
+    /// </summary>
+    public int MinLengthOneTypo { get; set; } = 3;
+
+    /// <summary>
+    /// Minimum word length for 2-typo correction to be applied.
+    /// Default: 7.
+    /// </summary>
+    public int MinLengthTwoTypos { get; set; } = 7;
     
     /// <summary>
     /// Minimum absolute number of word matches required
@@ -98,6 +116,9 @@ internal CoverageSetup(CoverageSetup source)
     {
         MinWordSize = source.MinWordSize;
         LevenshteinMaxWordSize = source.LevenshteinMaxWordSize;
+        NumTypos = source.NumTypos;
+        MinLengthOneTypo = source.MinLengthOneTypo;
+        MinLengthTwoTypos = source.MinLengthTwoTypos;
         CoverageMinWordHitsAbs = source.CoverageMinWordHitsAbs;
         CoverageMinWordHitsRelative = source.CoverageMinWordHitsRelative;
         CoverageQLimitForErrorTolerance = source.CoverageQLimitForErrorTolerance;
diff --git a/src/Infidex/Coverage/FuzzyWordMatcher.cs b/src/Infidex/Coverage/FuzzyWordMatcher.cs
@@ -11,7 +11,7 @@ internal static class FuzzyWordMatcher
     private const double DefaultErrorProbability = 0.04;
     private const double DefaultTailProbability = 0.01;
 
-    public static void Match(ref MatchState state, int minWordSize, int levenshteinMaxWordSize)
+    public static void Match(ref MatchState state, CoverageSetup setup)
     {
         int qCount = state.QCount;
         int dCount = state.DCount;
@@ -23,23 +23,30 @@ public static void Match(ref MatchState state, int minWordSize, int levenshteinM
         
         if (maxQueryLength == 0) return;
 
-        // Use a principled maximum edit distance based on a simple
-        // Binomial(L, p) error model instead of an uncalibrated
-        // relative distance knob.
-        //
-        // We choose the smallest d such that:
-        //   Pr[D ≤ d] ≥ 1 - alpha
-        //
-        // where D ~ Binomial(L, p), p ≈ DefaultErrorProbability and
-        // alpha ≈ DefaultTailProbability.
-        int maxEditDist = EditDistanceModel.GetMaxEditsForLength(
-            maxQueryLength,
-            DefaultErrorProbability,
-            DefaultTailProbability);
+        int maxEditDist;
+        
+        if (maxQueryLength >= setup.MinLengthTwoTypos) 
+            maxEditDist = 2;
+        else if (maxQueryLength >= setup.MinLengthOneTypo) 
+            maxEditDist = 1;
+        else 
+            maxEditDist = 0; 
 
-        // Ensure we allow at least one edit for non-empty queries.
-        if (maxEditDist < 1)
+        // Special handling for len=2 queries if default logic forbids typos (maxEditDist=0)
+        // If we have very short words (len 2) that are disallowed normal typos, 
+        // we conditionally allow 1 typo ONLY if it matches a target of length 3 (Insertion).
+        // This supports common cases like "te" -> "the" while avoiding high-noise substitutions like "at" -> "it".
+        bool hasSpecialShortWord = (maxQueryLength == 2 && maxEditDist == 0 && setup.NumTypos >= 1);
+        if (hasSpecialShortWord)
+        {
             maxEditDist = 1;
+        }
+
+        // Respect the global cap
+        if (maxEditDist > setup.NumTypos)
+            maxEditDist = setup.NumTypos;
+        
+        if (maxEditDist == 0) return;
         
         for (int editDist = 1; editDist <= maxEditDist; editDist++)
         {
@@ -55,12 +62,32 @@ public static void Match(ref MatchState state, int minWordSize, int levenshteinM
                 int qLen = qSlice.Length;
                 
                 // Skip query tokens that are too short for meaningful fuzzy matching
-                // but allow 2-char tokens since they can fuzzy-match 3-char words (e.g., "te" → "the")
-                if (qLen < minWordSize) continue;
+                if (qLen < setup.MinWordSize) continue;
+                
+                // Calculate max edits for THIS token
+                int tokenMaxEdits = 0;
+                if (qLen >= setup.MinLengthTwoTypos) tokenMaxEdits = 2;
+                else if (qLen >= setup.MinLengthOneTypo) tokenMaxEdits = 1;
+                else tokenMaxEdits = 0;
                 
+                // Apply special short word logic for individual token
+                bool isSpecialShortCase = false;
+                if (qLen == 2 && tokenMaxEdits == 0 && setup.NumTypos >= 1)
+                {
+                    tokenMaxEdits = 1;
+                    isSpecialShortCase = true;
+                }
+
+                if (tokenMaxEdits > setup.NumTypos) tokenMaxEdits = setup.NumTypos;
+                
+                if (editDist > tokenMaxEdits) continue;
+                
+                // For special short case, we only process editDist=1
+                if (isSpecialShortCase && editDist != 1) continue;
+
                 // Calculate the valid document word length range for this query token and edit distance
-                int minLen = Math.Max(minWordSize, qLen - editDist);
-                int maxLen = Math.Min(levenshteinMaxWordSize, qLen + editDist);
+                int minLen = Math.Max(setup.MinWordSize, qLen - editDist);
+                int maxLen = Math.Min(setup.LevenshteinMaxWordSize, qLen + editDist);
                 if (maxLen > 63) maxLen = 63;
                 
                 ReadOnlySpan<char> qText = state.QuerySpan.Slice(qSlice.Offset, qSlice.Length);
@@ -73,6 +100,16 @@ public static void Match(ref MatchState state, int minWordSize, int levenshteinM
                     if (dLen > maxLen || dLen < minLen) continue;
                     
                     ReadOnlySpan<char> dText = state.DocSpan.Slice(dSlice.Offset, dSlice.Length);
+
+                    // Enforce special short word constraints
+                    if (isSpecialShortCase)
+                    {
+                        // Special handling for short words: First character MUST match.
+                        // This allows "te" -> "the" (ins) or "te" -> "to" (sub), 
+                        // but prevents high-noise matches like "at" -> "cat" (prefix ins) or "at" -> "it" (start sub).
+                        if (dText.Length == 0 || char.ToLowerInvariant(dText[0]) != char.ToLowerInvariant(qText[0]))
+                            continue;
+                    }
                     
                     int dist = LevenshteinDistance.CalculateDamerau(qText, dText, editDist, ignoreCase: true);
                     
@@ -107,4 +144,3 @@ public static bool AllTermsFullyMatched(ref MatchState state)
         return true;
     }
 }
-
diff --git a/src/Infidex/Indexing/VectorModel.cs b/src/Infidex/Indexing/VectorModel.cs
@@ -400,6 +400,10 @@ internal TopKHeap SearchWithMaxScore(string queryText, int topK, Dictionary<int,
                 term.QueryOccurrences = (byte)shingle.Occurrences;
                 queryTerms.Add(term);
             }
+            else if (term == null && _fstIndex != null && shingle.Text.Length >= 4)
+            {
+                ExpandMissingTerm(shingle, queryTerms);
+            }
         }
 
         if (queryTerms.Count == 0 || _documents.Count == 0)
@@ -412,6 +416,22 @@ internal TopKHeap SearchWithMaxScore(string queryText, int topK, Dictionary<int,
         return Bm25Scorer.Search(queryTerms, topK, totalDocs, _docLengths!, _avgDocLength, _stopTermLimit, _documents, bestSegmentsMap, queryIndex, _shortQueryIndex, queryText);
     }
 
+    private void ExpandMissingTerm(Shingle shingle, List<Term> queryTerms)
+    {
+        // Fuzzy fallback for missing terms (Edit Distance 1 only)
+        List<int> fuzzyOutputs = new List<int>();
+        _fstIndex!.GetWithinEditDistance1(shingle.Text.AsSpan(), fuzzyOutputs);
+        
+        foreach (int termId in fuzzyOutputs)
+        {
+            Term? fuzzyTerm = _termCollection.GetTermByIndex(termId);
+            if (fuzzyTerm != null && fuzzyTerm.DocumentFrequency <= _stopTermLimit)
+            {
+                queryTerms.Add(fuzzyTerm);
+            }
+        }
+    }
+
     public void Save(string filePath)
     {
         using FileStream stream = File.Create(filePath);
diff --git a/src/Infidex/Tokenization/Tokenizer.cs b/src/Infidex/Tokenization/Tokenizer.cs
@@ -78,7 +78,48 @@ public List<Shingle> TokenizeForIndexing(string text, bool isSegmentContinuation
         // Add padding
         string paddedText = startPad + text + _stopPadding;
         
-        return GenerateShingles(paddedText);
+        List<Shingle> shingles = GenerateShingles(paddedText);
+
+        // index full words to support fuzzy correction and exact word matching
+        // this ensures that words longer than max n-gram size are present in the index/FST
+        if (TokenizerSetup != null)
+        {
+            ReadOnlySpan<char> span = text.AsSpan();
+            char[] delimiters = TokenizerSetup.Delimiters;
+            int baseOffset = isSegmentContinuation ? 0 : StartPadSize;
+            
+            int i = 0;
+            while (i < span.Length)
+            {
+                // Skip delimiters
+                int start = i;
+                while (start < span.Length && delimiters.Contains(span[start]))
+                {
+                    start++;
+                }
+                
+                if (start >= span.Length) break;
+                
+                // Find end of word
+                int end = start;
+                while (end < span.Length && !delimiters.Contains(span[end]))
+                {
+                    end++;
+                }
+                
+                int len = end - start;
+                if (len >= IndexSizes[0])
+                {
+                    string word = new string(span.Slice(start, len));
+                    // Exact position in padded text
+                    shingles.Add(new Shingle(word, 1, baseOffset + start));
+                }
+                
+                i = end;
+            }
+        }
+        
+        return shingles;
     }
     
     /// <summary>
@@ -351,4 +392,3 @@ public HashSet<string> GetWordTokensForCoverage(string text, int minWordSize)
         return result;
     }
 }
-

Original file line number	Diff line number	Diff line change
`@@ -289,7 +289,7 @@ private CoverageResult CalculateCoverageInternal(string query, string documentTe`
`289`	`289`	`PrefixSuffixMatcher.Match(ref state);`
`290`	`290`
`291`	`291`	`if (_setup.CoverFuzzyWords && qCount > 0 && !FuzzyWordMatcher.AllTermsFullyMatched(ref state))`
`292`		`- FuzzyWordMatcher.Match(ref state, _setup.MinWordSize, _setup.LevenshteinMaxWordSize);`
	`292`	`+ FuzzyWordMatcher.Match(ref state, _setup);`
`293`	`293`	`}`
`294`	`294`	`finally`
`295`	`295`	`{`