Completion FSTs to be loaded off-heap at all times (#14364)

All the existing completion postings format load their FSTs on-heap. It is possible to customize that behaviour by mainintaing a custom postings format that override the fst load mode.

TestSuggestField attempted to test all modes, but in practice, it can only test the default load mode because the read path relies on the default constructor called via SPI, that does not allow providing the fst load mode.

This commit switches loading of FST to off-heap, and removes the fst load mode argument from all completion postings format classes, effectively making the loading always off-heap.

Author: javanna

lucene/CHANGES.txt

@@ -100,6 +100,7 @@ Improvements
 
 * GITHUB#14213: Allowing indexing stored-only StoredField directly from DataInput. (Tim Brooks)
 
+* GITHUB#14364: Completion FSTs to be loaded off-heap at all times. (Luca Cavanna)
 
 Optimizations
 ---------------------

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/Completion101PostingsFormat.java

@@ -25,17 +25,9 @@
  * @lucene.experimental
  */
 public class Completion101PostingsFormat extends CompletionPostingsFormat {
-  /** Creates a {@link Completion101PostingsFormat} that will load the completion FST on-heap. */
+  /** Creates a {@link Completion101PostingsFormat}. */
   public Completion101PostingsFormat() {
-    this(FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link Completion101PostingsFormat} that will use the provided <code>fstLoadMode
-   * </code> to determine if the completion FST should be loaded on or off heap.
-   */
-  public Completion101PostingsFormat(FSTLoadMode fstLoadMode) {
-    super("Completion101", fstLoadMode);
+    super("Completion101");
   }
 
   @Override

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/Completion50PostingsFormat.java

@@ -27,17 +27,9 @@
  * @lucene.experimental
  */
 public class Completion50PostingsFormat extends CompletionPostingsFormat {
-  /** Creates a {@link Completion50PostingsFormat} that will load the completion FST on-heap. */
+  /** Creates a {@link Completion50PostingsFormat}. */
   public Completion50PostingsFormat() {
-    this(FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link Completion50PostingsFormat} that will use the provided <code>fstLoadMode
-   * </code> to determine if the completion FST should be loaded on or off heap.
-   */
-  public Completion50PostingsFormat(FSTLoadMode fstLoadMode) {
-    super("completion", fstLoadMode);
+    super("completion");
   }
 
   @Override

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/Completion84PostingsFormat.java

@@ -27,17 +27,9 @@
  * @lucene.experimental
  */
 public class Completion84PostingsFormat extends CompletionPostingsFormat {
-  /** Creates a {@link Completion84PostingsFormat} that will load the completion FST on-heap. */
+  /** Creates a {@link Completion84PostingsFormat}. */
   public Completion84PostingsFormat() {
-    this(FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link Completion84PostingsFormat} that will use the provided <code>fstLoadMode
-   * </code> to determine if the completion FST should be loaded on or off heap.
-   */
-  public Completion84PostingsFormat(FSTLoadMode fstLoadMode) {
-    super("Completion84", fstLoadMode);
+    super("Completion84");
   }
 
   @Override

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/Completion90PostingsFormat.java

@@ -27,17 +27,9 @@
  * @lucene.experimental
  */
 public class Completion90PostingsFormat extends CompletionPostingsFormat {
-  /** Creates a {@link Completion90PostingsFormat} that will load the completion FST on-heap. */
+  /** Creates a {@link Completion90PostingsFormat}. */
   public Completion90PostingsFormat() {
-    this(FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link Completion90PostingsFormat} that will use the provided <code>fstLoadMode
-   * </code> to determine if the completion FST should be loaded on or off heap.
-   */
-  public Completion90PostingsFormat(FSTLoadMode fstLoadMode) {
-    super("Completion90", fstLoadMode);
+    super("Completion90");
   }
 
   @Override

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/Completion912PostingsFormat.java

@@ -27,17 +27,9 @@
  * @lucene.experimental
  */
 public class Completion912PostingsFormat extends CompletionPostingsFormat {
-  /** Creates a {@link Completion912PostingsFormat} that will load the completion FST on-heap. */
+  /** Creates a {@link Completion912PostingsFormat}. */
   public Completion912PostingsFormat() {
-    this(FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link Completion912PostingsFormat} that will use the provided <code>fstLoadMode
-   * </code> to determine if the completion FST should be loaded on or off heap.
-   */
-  public Completion912PostingsFormat(FSTLoadMode fstLoadMode) {
-    super("Completion912", fstLoadMode);
+    super("Completion912");
   }
 
   @Override

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/Completion99PostingsFormat.java

@@ -27,17 +27,9 @@
  * @lucene.experimental
  */
 public class Completion99PostingsFormat extends CompletionPostingsFormat {
-  /** Creates a {@link Completion99PostingsFormat} that will load the completion FST on-heap. */
+  /** Creates a {@link Completion99PostingsFormat}. */
   public Completion99PostingsFormat() {
-    this(FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link Completion99PostingsFormat} that will use the provided <code>fstLoadMode
-   * </code> to determine if the completion FST should be loaded on or off heap.
-   */
-  public Completion99PostingsFormat(FSTLoadMode fstLoadMode) {
-    super("Completion99", fstLoadMode);
+    super("Completion99");
   }
 
   @Override

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/CompletionFieldsProducer.java

@@ -35,7 +35,6 @@
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.Terms;
-import org.apache.lucene.search.suggest.document.CompletionPostingsFormat.FSTLoadMode;
 import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.Accountable;
@@ -56,7 +55,7 @@
 final class CompletionFieldsProducer extends FieldsProducer implements Accountable {
 
   private FieldsProducer delegateFieldsProducer;
-  private Map<String, CompletionsTermsReader> readers;
+  private final Map<String, CompletionsTermsReader> readers;
   private IndexInput dictIn;
 
   // copy ctr for merge instance
@@ -66,8 +65,7 @@ private CompletionFieldsProducer(
     this.readers = readers;
   }
 
-  CompletionFieldsProducer(String codecName, SegmentReadState state, FSTLoadMode fstLoadMode)
-      throws IOException {
+  CompletionFieldsProducer(String codecName, SegmentReadState state) throws IOException {
     String indexFile =
         IndexFileNames.segmentFileName(
             state.segmentInfo.name, state.segmentSuffix, INDEX_EXTENSION);
@@ -114,8 +112,7 @@ private CompletionFieldsProducer(
         FieldInfo fieldInfo = state.fieldInfos.fieldInfo(fieldNumber);
         // we don't load the FST yet
         readers.put(
-            fieldInfo.name,
-            new CompletionsTermsReader(dictIn, offset, minWeight, maxWeight, type, fstLoadMode));
+            fieldInfo.name, new CompletionsTermsReader(dictIn, offset, minWeight, maxWeight, type));
       }
       CodecUtil.checkFooter(index);
       success = true;

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/CompletionPostingsFormat.java

@@ -104,36 +104,9 @@ public abstract class CompletionPostingsFormat extends PostingsFormat {
   static final String INDEX_EXTENSION = "cmp";
   static final String DICT_EXTENSION = "lkp";
 
-  /** An enum that allows to control if suggester FSTs are loaded into memory or read off-heap */
-  public enum FSTLoadMode {
-    /**
-     * Always read FSTs from disk. NOTE: If this option is used the FST will be read off-heap even
-     * if buffered directory implementations are used.
-     */
-    OFF_HEAP,
-    /** Never read FSTs from disk ie. all suggest fields FSTs are loaded into memory */
-    ON_HEAP,
-    /**
-     * Automatically make the decision if FSTs are read from disk depending if the segment read from
-     * an MMAPDirectory
-     */
-    AUTO
-  }
-
-  private final FSTLoadMode fstLoadMode;
-
-  /** Used only by core Lucene at read-time via Service Provider instantiation */
+  /** Creates a {@link CompletionPostingsFormat} with the given name. */
   public CompletionPostingsFormat(String name) {
-    this(name, FSTLoadMode.ON_HEAP);
-  }
-
-  /**
-   * Creates a {@link CompletionPostingsFormat} that will use the provided <code>fstLoadMode</code>
-   * to determine if the completion FST should be loaded on or off heap.
-   */
-  public CompletionPostingsFormat(String name, FSTLoadMode fstLoadMode) {
     super(name);
-    this.fstLoadMode = fstLoadMode;
   }
 
   /** Concrete implementation should specify the delegating postings format */
@@ -153,6 +126,6 @@ public FieldsConsumer fieldsConsumer(SegmentWriteState state) throws IOException
 
   @Override
   public FieldsProducer fieldsProducer(SegmentReadState state) throws IOException {
-    return new CompletionFieldsProducer(getName(), state, fstLoadMode);
+    return new CompletionFieldsProducer(getName(), state);
   }
 }

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/CompletionsTermsReader.java

@@ -19,7 +19,6 @@
 import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
-import org.apache.lucene.search.suggest.document.CompletionPostingsFormat.FSTLoadMode;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.Accountable;
 
@@ -41,29 +40,21 @@ public final class CompletionsTermsReader implements Accountable {
   private final IndexInput dictIn;
   private final long offset;
 
-  private final FSTLoadMode fstLoadMode;
-
   private NRTSuggester suggester;
 
   /**
    * Creates a CompletionTermsReader to load a field-specific suggester from the index <code>dictIn
    * </code> with <code>offset</code>
    */
   CompletionsTermsReader(
-      IndexInput dictIn,
-      long offset,
-      long minWeight,
-      long maxWeight,
-      byte type,
-      FSTLoadMode fstLoadMode) {
+      IndexInput dictIn, long offset, long minWeight, long maxWeight, byte type) {
     assert minWeight <= maxWeight;
     assert offset >= 0l && offset < dictIn.length();
     this.dictIn = dictIn;
     this.offset = offset;
     this.minWeight = minWeight;
     this.maxWeight = maxWeight;
     this.type = type;
-    this.fstLoadMode = fstLoadMode;
   }
 
   /**
@@ -73,7 +64,7 @@ public final class CompletionsTermsReader implements Accountable {
   public synchronized NRTSuggester suggester() throws IOException {
     if (suggester == null) {
       IndexInput indexInput = dictIn.slice("NRTSuggester", offset, dictIn.length() - offset);
-      suggester = NRTSuggester.load(indexInput, fstLoadMode);
+      suggester = NRTSuggester.load(indexInput);
     }
     return suggester;
   }

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/NRTSuggester.java

@@ -24,7 +24,6 @@
 import java.util.Comparator;
 import java.util.List;
 import org.apache.lucene.search.suggest.analyzing.FSTUtil;
-import org.apache.lucene.search.suggest.document.CompletionPostingsFormat.FSTLoadMode;
 import org.apache.lucene.store.ByteArrayDataInput;
 import org.apache.lucene.store.ByteArrayDataOutput;
 import org.apache.lucene.store.IndexInput;
@@ -306,37 +305,15 @@ private static double calculateLiveDocRatio(int numDocs, int maxDocs) {
     return (numDocs > 0) ? ((double) numDocs / maxDocs) : -1;
   }
 
-  private static boolean shouldLoadFSTOffHeap(IndexInput input, FSTLoadMode fstLoadMode) {
-    switch (fstLoadMode) {
-      case ON_HEAP:
-        return false;
-      case OFF_HEAP:
-        return true;
-      case AUTO:
-        // TODO: Make this less hacky to maybe expose "off-heap" feature using a marker interface on
-        // the IndexInput
-        return input.getClass().getName().contains(".MemorySegmentIndexInput");
-      default:
-        throw new IllegalStateException("unknown enum constant: " + fstLoadMode);
-    }
-  }
-
-  /**
-   * Loads a {@link NRTSuggester} from {@link org.apache.lucene.store.IndexInput} on or off-heap
-   * depending on the provided <code>fstLoadMode</code>
-   */
-  static NRTSuggester load(IndexInput input, FSTLoadMode fstLoadMode) throws IOException {
+  /** Loads a {@link NRTSuggester} from {@link org.apache.lucene.store.IndexInput} */
+  static NRTSuggester load(IndexInput input) throws IOException {
     final FST<Pair<Long, BytesRef>> fst;
     PairOutputs<Long, BytesRef> outputs =
         new PairOutputs<>(PositiveIntOutputs.getSingleton(), ByteSequenceOutputs.getSingleton());
-    if (shouldLoadFSTOffHeap(input, fstLoadMode)) {
-      final FST.FSTMetadata<Pair<Long, BytesRef>> fstMetadata = FST.readMetadata(input, outputs);
-      OffHeapFSTStore store = new OffHeapFSTStore(input, input.getFilePointer(), fstMetadata);
-      fst = FST.fromFSTReader(fstMetadata, store);
-      input.skipBytes(store.size());
-    } else {
-      fst = new FST<>(FST.readMetadata(input, outputs), input);
-    }
+    final FST.FSTMetadata<Pair<Long, BytesRef>> fstMetadata = FST.readMetadata(input, outputs);
+    OffHeapFSTStore store = new OffHeapFSTStore(input, input.getFilePointer(), fstMetadata);
+    fst = FST.fromFSTReader(fstMetadata, store);
+    input.skipBytes(store.size());
 
     /* read some meta info */
     int maxAnalyzedPathsPerOutput = input.readVInt();

lucene/suggest/src/java/org/apache/lucene/search/suggest/document/NRTSuggesterBuilder.java

@@ -99,10 +99,7 @@ public void finishTerm() throws IOException {
     entries.clear();
   }
 
-  /**
-   * Builds and stores a FST that can be loaded with {@link NRTSuggester#load(IndexInput,
-   * CompletionPostingsFormat.FSTLoadMode)})}
-   */
+  /** Builds and stores a FST that can be loaded with {@link NRTSuggester#load(IndexInput)})} */
   public boolean store(DataOutput output) throws IOException {
     final FST<PairOutputs.Pair<Long, BytesRef>> fst =
         FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());

lucene/suggest/src/test/org/apache/lucene/search/suggest/document/TestSuggestField.java

@@ -20,7 +20,6 @@
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.startsWith;
 
-import com.carrotsearch.randomizedtesting.generators.RandomPicks;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -949,9 +948,7 @@ static IndexWriterConfig iwcWithSuggestField(Analyzer analyzer, final Set<String
     iwc.setMergePolicy(newLogMergePolicy());
     Codec filterCodec =
         new FilterCodec(TestUtil.getDefaultCodec().getName(), TestUtil.getDefaultCodec()) {
-          final CompletionPostingsFormat.FSTLoadMode fstLoadMode =
-              RandomPicks.randomFrom(random(), CompletionPostingsFormat.FSTLoadMode.values());
-          final PostingsFormat postingsFormat = new Completion101PostingsFormat(fstLoadMode);
+          final PostingsFormat postingsFormat = new Completion101PostingsFormat();
 
           @Override
           public PostingsFormat postingsFormat() {