Improve snowball stemming by adding an insert-only cache

lucene/CHANGES.txt

@@ -340,6 +340,9 @@ Optimizations
 * GITHUB#16370: Sort two-phase iterators by matchCost in DenseConjunctionBulkScorer.
   (Alan Woodward)
 
+* GITHUB#16356: Cache stem results in SnowballFilter to avoid redundant stemmer invocations on
+  repeated tokens. 2.3-2.7x throughput improvement. Disableable via new constructor. (Costin Leau)
+
 Bug Fixes
 ---------------------
 * GITHUB#16350: Disable bulk-scoring in monitor queries. (Alan Woodward)

lucene/analysis/common/src/java/org/apache/lucene/analysis/snowball/SnowballFilter.java

@@ -18,6 +18,7 @@
 
 import java.io.IOException;
 import java.util.Objects;
+import org.apache.lucene.analysis.CharArrayMap;
 import org.apache.lucene.analysis.LowerCaseFilter;
 import org.apache.lucene.analysis.TokenFilter;
 import org.apache.lucene.analysis.TokenStream;
@@ -47,14 +48,33 @@
  */
 public final class SnowballFilter extends TokenFilter {
 
+  private static final int DEFAULT_MAX_CACHE_SIZE = 1024;
+  private static final int MAX_CACHEABLE_LENGTH = 10;
+
   private final SnowballStemmer stemmer;
+  private final int maxCacheSize;
+  private final CharArrayMap<char[]> cache;
 
   private final CharTermAttribute termAtt = addAttribute(CharTermAttribute.class);
   private final KeywordAttribute keywordAttr = addAttribute(KeywordAttribute.class);
 
   public SnowballFilter(TokenStream input, SnowballStemmer stemmer) {
+    this(input, stemmer, DEFAULT_MAX_CACHE_SIZE);
+  }
+
+  /**
+   * Creates a SnowballFilter with the specified cache size. The cache stores stem results keyed by
+   * the original token text. Only tokens with length &le; {@value #MAX_CACHEABLE_LENGTH} are
+   * cached. Set {@code maxCacheSize} to 0 to disable caching.
+   */
+  public SnowballFilter(TokenStream input, SnowballStemmer stemmer, int maxCacheSize) {
     super(input);
     this.stemmer = Objects.requireNonNull(stemmer, "stemmer");
+    if (maxCacheSize < 0) {
+      throw new IllegalArgumentException("maxCacheSize must be >= 0, got: " + maxCacheSize);
+    }
+    this.maxCacheSize = maxCacheSize;
+    this.cache = maxCacheSize > 0 ? new CharArrayMap<>(Math.min(maxCacheSize, 256), false) : null;
   }
 
   /**
@@ -68,8 +88,16 @@ public SnowballFilter(TokenStream input, SnowballStemmer stemmer) {
    * @param name the name of a stemmer
    */
   public SnowballFilter(TokenStream in, String name) {
+    this(in, name, DEFAULT_MAX_CACHE_SIZE);
+  }
+
+  /** Same as {@link #SnowballFilter(TokenStream, String)} with a configurable cache size. */
+  public SnowballFilter(TokenStream in, String name, int maxCacheSize) {
     super(in);
     Objects.requireNonNull(name, "name");
+    if (maxCacheSize < 0) {
+      throw new IllegalArgumentException("maxCacheSize must be >= 0, got: " + maxCacheSize);
+    }
     // it was called "German2" for eons, but snowball folded it into "German" and deleted "German2"
     // for now, don't annoy our users...
     if (name.equals("German2")) {
@@ -85,6 +113,8 @@ public SnowballFilter(TokenStream in, String name) {
     } catch (ReflectiveOperationException e) {
       throw new IllegalArgumentException("Invalid stemmer class specified: " + name, e);
     }
+    this.maxCacheSize = maxCacheSize;
+    this.cache = maxCacheSize > 0 ? new CharArrayMap<>(Math.min(maxCacheSize, 256), false) : null;
   }
 
   /** Returns the next input Token, after being stemmed */
@@ -94,12 +124,34 @@ public boolean incrementToken() throws IOException {
       if (!keywordAttr.isKeyword()) {
         char[] termBuffer = termAtt.buffer();
         final int length = termAtt.length();
+        final boolean lookupCache = cache != null && length <= MAX_CACHEABLE_LENGTH;
+
+        if (lookupCache) {
+          char[] cached = cache.get(termBuffer, 0, length);
+          if (cached != null) {
+            termAtt.copyBuffer(cached, 0, cached.length);
+            return true;
+          }
+        }
+
+        char[] key = null;
+        if (lookupCache && cache.size() < maxCacheSize) {
+          key = new char[length];
+          System.arraycopy(termBuffer, 0, key, 0, length);
+        }
+
         stemmer.setCurrent(termBuffer, length);
         stemmer.stem();
         final char[] finalTerm = stemmer.getCurrentBuffer();
         final int newLength = stemmer.getCurrentBufferLength();
         if (finalTerm != termBuffer) termAtt.copyBuffer(finalTerm, 0, newLength);
         else termAtt.setLength(newLength);
+
+        if (key != null) {
+          char[] value = new char[newLength];
+          System.arraycopy(finalTerm, 0, value, 0, newLength);
+          cache.put(key, value);
+        }
       }
       return true;
     } else {

lucene/analysis/common/src/test/org/apache/lucene/analysis/snowball/TestSnowball.java

@@ -18,8 +18,10 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.StringReader;
 import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.List;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.analysis.TokenStream;
@@ -142,4 +144,123 @@ protected TokenStreamComponents createComponents(String fieldName) {
     checkRandomData(random(), a, 20 * RANDOM_MULTIPLIER);
     a.close();
   }
+
+  public void testCacheProducesSameOutput() throws Exception {
+    String input = "he abhorred accents running acknowledging internationalization";
+
+    List<String> uncached = stemAll(input, 0);
+    List<String> cached = stemAll(input, 1024);
+
+    assertEquals(uncached, cached);
+  }
+
+  public void testCacheDisabled() throws Exception {
+    Analyzer a =
+        new Analyzer() {
+          @Override
+          protected TokenStreamComponents createComponents(String fieldName) {
+            Tokenizer tokenizer = new MockTokenizer();
+            return new TokenStreamComponents(
+                tokenizer, new SnowballFilter(tokenizer, "English", 0));
+          }
+        };
+
+    assertAnalyzesTo(a, "he abhorred accents", new String[] {"he", "abhor", "accent"});
+    a.close();
+  }
+
+  public void testLongTokensBypassCache() throws Exception {
+    String input = "acknowledging internationalization";
+
+    List<String> uncached = stemAll(input, 0);
+    List<String> cached = stemAll(input, 1024);
+
+    assertEquals(uncached, cached);
+    assertEquals(2, cached.size());
+  }
+
+  public void testCacheAccumulatesAcrossReset() throws Exception {
+    String input = "he abhorred accents";
+    String[] expected = new String[] {"he", "abhor", "accent"};
+
+    Analyzer a =
+        new Analyzer() {
+          @Override
+          protected TokenStreamComponents createComponents(String fieldName) {
+            Tokenizer tokenizer = new MockTokenizer();
+            return new TokenStreamComponents(
+                tokenizer, new SnowballFilter(tokenizer, "English", 1024));
+          }
+        };
+
+    assertAnalyzesTo(a, input, expected);
+    assertAnalyzesTo(a, input, expected);
+    a.close();
+  }
+
+  public void testCacheActuallyHits() throws Exception {
+    int[] stemCallCount = {0};
+    org.tartarus.snowball.ext.EnglishStemmer countingStemmer =
+        new org.tartarus.snowball.ext.EnglishStemmer() {
+          @Override
+          public boolean stem() {
+            stemCallCount[0]++;
+            return super.stem();
+          }
+        };
+
+    Analyzer a =
+        new Analyzer() {
+          @Override
+          protected TokenStreamComponents createComponents(String fieldName) {
+            Tokenizer tokenizer = new MockTokenizer();
+            return new TokenStreamComponents(
+                tokenizer, new SnowballFilter(tokenizer, countingStemmer, 1024));
+          }
+        };
+
+    try (TokenStream ts = a.tokenStream("field", new StringReader("running jumping"))) {
+      ts.reset();
+      while (ts.incrementToken()) {}
+      ts.end();
+    }
+    assertEquals(2, stemCallCount[0]);
+
+    try (TokenStream ts = a.tokenStream("field", new StringReader("running jumping"))) {
+      CharTermAttribute termAtt = ts.addAttribute(CharTermAttribute.class);
+      ts.reset();
+      List<String> stems = new ArrayList<>();
+      while (ts.incrementToken()) {
+        stems.add(termAtt.toString());
+      }
+      ts.end();
+      assertEquals(List.of("run", "jump"), stems);
+    }
+    assertEquals("Cache should prevent additional stem() calls", 2, stemCallCount[0]);
+
+    a.close();
+  }
+
+  private List<String> stemAll(String input, int cacheSize) throws IOException {
+    List<String> result = new ArrayList<>();
+    Analyzer a =
+        new Analyzer() {
+          @Override
+          protected TokenStreamComponents createComponents(String fieldName) {
+            Tokenizer tokenizer = new MockTokenizer();
+            return new TokenStreamComponents(
+                tokenizer, new SnowballFilter(tokenizer, "English", cacheSize));
+          }
+        };
+    try (TokenStream ts = a.tokenStream("field", new StringReader(input))) {
+      CharTermAttribute termAtt = ts.addAttribute(CharTermAttribute.class);
+      ts.reset();
+      while (ts.incrementToken()) {
+        result.add(termAtt.toString());
+      }
+      ts.end();
+    }
+    a.close();
+    return result;
+  }
 }

lucene/benchmark-jmh/build.gradle

@@ -19,6 +19,7 @@ description = 'Lucene JMH micro-benchmarking module'
 
 dependencies {
   moduleImplementation project(':lucene:core')
+  moduleImplementation project(':lucene:analysis:common')
   moduleImplementation project(':lucene:expressions')
   moduleImplementation project(':lucene:join')
   moduleImplementation project(':lucene:sandbox')

lucene/benchmark-jmh/src/java/module-info.java

@@ -23,6 +23,7 @@
   requires jmh.core;
   requires jdk.unsupported;
   requires org.apache.lucene.core;
+  requires org.apache.lucene.analysis.common;
   requires org.apache.lucene.expressions;
   requires org.apache.lucene.join;
   requires org.apache.lucene.sandbox;