Make FSTCompiler.compile() to only return the FSTMetadata (apache#12831)

* Make FSTCompiler.compile() to only return the FSTMetadata

* tidy code

lucene/CHANGES.txt

@@ -80,6 +80,10 @@ API Changes
 * GITHUB#12875: Ensure token position is always increased in PathHierarchyTokenizer and ReversePathHierarchyTokenizer
   and resulting tokens do not overlap. (Michael Froh, Lukáš Vlček)
 
+* GITHUB#12624, GITHUB#12831: Allow FSTCompiler to stream to any DataOutput while building, and
+  make compile() only return the FSTMetadata. For on-heap (default) use case, please use
+  FST.fromFSTReader(fstMetadata, fstCompiler.getFSTReader()) to create the FST. (Anh Dung Bui)
+
 New Features
 ---------------------
 

lucene/analysis/common/src/java/org/apache/lucene/analysis/charfilter/NormalizeCharMap.java

@@ -111,7 +111,7 @@ public NormalizeCharMap build() {
         for (Map.Entry<String, String> ent : pendingPairs.entrySet()) {
           fstCompiler.add(Util.toUTF16(ent.getKey(), scratch), new CharsRef(ent.getValue()));
         }
-        map = fstCompiler.compile();
+        map = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
         pendingPairs.clear();
       } catch (IOException ioe) {
         // Bogus FST IOExceptions!!  (will never happen)

lucene/analysis/common/src/java/org/apache/lucene/analysis/hunspell/ConvTable.java

@@ -51,7 +51,7 @@ class ConvTable {
         fstCompiler.add(scratchInts.get(), new CharsRef(entry.getValue()));
       }
 
-      fst = fstCompiler.compile();
+      fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
     } catch (IOException bogus) {
       throw new RuntimeException(bogus);
     }

lucene/analysis/common/src/java/org/apache/lucene/analysis/hunspell/Dictionary.java

@@ -657,7 +657,7 @@ private FST<IntsRef> affixFST(TreeMap<String, List<Integer>> affixes) throws IOE
       }
       fstCompiler.add(scratch.get(), output);
     }
-    return fstCompiler.compile();
+    return FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
   }
 
   /**

lucene/analysis/common/src/java/org/apache/lucene/analysis/miscellaneous/StemmerOverrideFilter.java

@@ -222,7 +222,8 @@ public StemmerOverrideMap build() throws IOException {
         intsSpare.copyUTF8Bytes(bytesRef);
         fstCompiler.add(intsSpare.get(), new BytesRef(outputValues.get(id)));
       }
-      return new StemmerOverrideMap(fstCompiler.compile(), ignoreCase);
+      return new StemmerOverrideMap(
+          FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader()), ignoreCase);
     }
   }
 }

lucene/analysis/common/src/java/org/apache/lucene/analysis/synonym/SynonymMap.java

@@ -291,7 +291,7 @@ public SynonymMap build() throws IOException {
         fstCompiler.add(Util.toUTF32(input, scratchIntsRef), scratch.toBytesRef());
       }
 
-      FST<BytesRef> fst = fstCompiler.compile();
+      FST<BytesRef> fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
       return new SynonymMap(fst, words, maxHorizontalContext);
     }
   }

lucene/analysis/kuromoji/src/java/org/apache/lucene/analysis/ja/dict/TokenInfoDictionaryBuilder.java

@@ -126,7 +126,7 @@ private TokenInfoDictionaryWriter buildDictionary(List<Path> csvFiles) throws IO
       dictionary.addMapping((int) ord, offset);
       offset = next;
     }
-    dictionary.setFST(fstCompiler.compile());
+    dictionary.setFST(FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader()));
     return dictionary;
   }
 

lucene/analysis/kuromoji/src/java/org/apache/lucene/analysis/ja/dict/UserDictionary.java

@@ -147,7 +147,9 @@ public int compare(String[] left, String[] right) {
       segmentations.add(wordIdAndLength);
       ord++;
     }
-    this.fst = new TokenInfoFST(fstCompiler.compile(), false);
+    this.fst =
+        new TokenInfoFST(
+            FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader()), false);
     this.morphAtts = new UserMorphData(data.toArray(new String[0]));
     this.segmentations = segmentations.toArray(new int[segmentations.size()][]);
   }

lucene/analysis/nori/src/java/org/apache/lucene/analysis/ko/dict/TokenInfoDictionaryBuilder.java

@@ -122,7 +122,7 @@ private TokenInfoDictionaryWriter buildDictionary(List<Path> csvFiles) throws IO
       dictionary.addMapping((int) ord, offset);
       offset = next;
     }
-    dictionary.setFST(fstCompiler.compile());
+    dictionary.setFST(FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader()));
     return dictionary;
   }
 }

lucene/analysis/nori/src/java/org/apache/lucene/analysis/ko/dict/UserDictionary.java

@@ -130,7 +130,8 @@ private UserDictionary(List<String> entries) throws IOException {
       lastToken = token;
       ord++;
     }
-    this.fst = new TokenInfoFST(fstCompiler.compile());
+    this.fst =
+        new TokenInfoFST(FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader()));
     int[][] segmentations = _segmentations.toArray(new int[_segmentations.size()][]);
     short[] rightIds = new short[_rightIds.size()];
     for (int i = 0; i < _rightIds.size(); i++) {

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene40/blocktree/Lucene40BlockTreeTermsWriter.java

@@ -498,7 +498,7 @@ public void compileIndex(
         }
       }
 
-      index = fstCompiler.compile();
+      index = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
 
       assert subIndices == null;
 

lucene/classification/src/java/org/apache/lucene/classification/BooleanPerceptronClassifier.java

@@ -216,7 +216,7 @@ private void updateFST(SortedMap<String, Double> weights) throws IOException {
       fstCompiler.add(
           Util.toIntsRef(scratchBytes.get(), scratchInts), entry.getValue().longValue());
     }
-    fst = fstCompiler.compile();
+    fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
   }
 
   @Override

lucene/codecs/src/java/org/apache/lucene/codecs/blockterms/VariableGapTermsIndexWriter.java

@@ -283,7 +283,7 @@ public void add(BytesRef text, TermStats stats, long termsFilePointer) throws IO
 
     @Override
     public void finish(long termsFilePointer) throws IOException {
-      fst = fstCompiler.compile();
+      fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
       if (fst != null) {
         fst.save(out, out);
       }

lucene/codecs/src/java/org/apache/lucene/codecs/blocktreeords/OrdsBlockTreeTermsWriter.java

@@ -425,7 +425,7 @@ public void compileIndex(
 
       assert sumTotalTermCount == totFloorTermCount;
 
-      index = fstCompiler.compile();
+      index = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
       assert subIndices == null;
 
       /*

lucene/codecs/src/java/org/apache/lucene/codecs/memory/FSTTermsWriter.java

@@ -277,7 +277,8 @@ public void finishTerm(BytesRef text, BlockTermState state) throws IOException {
     public void finish(long sumTotalTermFreq, long sumDocFreq, int docCount) throws IOException {
       // save FST dict
       if (numTerms > 0) {
-        final FST<FSTTermOutputs.TermData> fst = fstCompiler.compile();
+        final FST<FSTTermOutputs.TermData> fst =
+            FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
         fields.add(
             new FieldMetaData(fieldInfo, numTerms, sumTotalTermFreq, sumDocFreq, docCount, fst));
       }

lucene/codecs/src/java/org/apache/lucene/codecs/simpletext/SimpleTextFieldsReader.java

@@ -738,7 +738,7 @@ private void loadTerms() throws IOException {
         }
       }
       docCount = visitedDocs.cardinality();
-      fst = fstCompiler.compile();
+      fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
       /*
       PrintStream ps = new PrintStream("out.dot");
       fst.toDot(ps);

lucene/codecs/src/java/org/apache/lucene/codecs/uniformsplit/FSTDictionary.java

@@ -185,7 +185,8 @@ public void add(BytesRef blockKey, long blockFilePointer) throws IOException {
 
     @Override
     public FSTDictionary build() throws IOException {
-      return new FSTDictionary(fstCompiler.compile());
+      return new FSTDictionary(
+          FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader()));
     }
   }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene90/blocktree/Lucene90BlockTreeTermsWriter.java

@@ -587,7 +587,7 @@ public void compileIndex(
         }
       }
 
-      index = fstCompiler.compile();
+      index = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
 
       assert subIndices == null;
 

lucene/core/src/java/org/apache/lucene/util/fst/FST.java

@@ -437,6 +437,21 @@ public FST(FSTMetadata<T> metadata, DataInput in, FSTStore fstStore) throws IOEx
     this.fstReader = fstReader;
   }
 
+  /**
+   * Create a FST from a {@link FSTReader}. Return null if the metadata is null.
+   *
+   * @param fstMetadata the metadata
+   * @param fstReader the FSTReader
+   * @return the FST
+   */
+  public static <T> FST<T> fromFSTReader(FSTMetadata<T> fstMetadata, FSTReader fstReader) {
+    // FSTMetadata could be null if there is no node accepted by the FST
+    if (fstMetadata == null) {
+      return null;
+    }
+    return new FST<>(fstMetadata, Objects.requireNonNull(fstReader, "FSTReader cannot be null"));
+  }
+
   /**
    * Read the FST metadata from DataInput
    *
@@ -516,9 +531,7 @@ public FSTMetadata<T> getMetadata() {
   }
 
   /**
-   * Save the FST to DataOutput. If you use an {@link org.apache.lucene.store.IndexOutput} to build
-   * the FST, then you should not and do not need to call this method, as the FST is already saved.
-   * Doing so will throw an {@link UnsupportedOperationException}.
+   * Save the FST to DataOutput.
    *
    * @param metaOut the DataOutput to write the metadata to
    * @param out the DataOutput to write the FST bytes to

lucene/core/src/java/org/apache/lucene/util/fst/FSTCompiler.java

@@ -99,6 +99,7 @@ public class FSTCompiler<T> {
   private static final FSTReader NULL_FST_READER = new NullFSTReader();
 
   private final NodeHash<T> dedupHash;
+  // a temporary FST used during building for NodeHash cache
   final FST<T> fst;
   private final T NO_OUTPUT;
 
@@ -173,9 +174,7 @@ private FSTCompiler(
     paddingBytePending = true;
     this.dataOutput = dataOutput;
     fst =
-        new FST<>(
-            new FST.FSTMetadata<>(inputType, outputs, null, -1, version, 0),
-            toFSTReader(dataOutput));
+        new FST<>(new FST.FSTMetadata<>(inputType, outputs, null, -1, version, 0), NULL_FST_READER);
     if (suffixRAMLimitMB < 0) {
       throw new IllegalArgumentException("ramLimitMB must be >= 0; got: " + suffixRAMLimitMB);
     } else if (suffixRAMLimitMB > 0) {
@@ -193,16 +192,6 @@ private FSTCompiler(
     }
   }
 
-  // Get the respective FSTReader of the DataOutput. If the DataOutput is also a FSTReader then we
-  // will use it, otherwise we will return a NullFSTReader. Attempting to read from a FST with
-  // NullFSTReader will throw UnsupportedOperationException
-  private FSTReader toFSTReader(DataOutput dataOutput) {
-    if (dataOutput instanceof FSTReader) {
-      return (FSTReader) dataOutput;
-    }
-    return NULL_FST_READER;
-  }
-
   /**
    * This class is used for FST backed by non-FSTReader DataOutput. It does not allow getting the
    * reverse BytesReader nor writing to a DataOutput.
@@ -227,6 +216,22 @@ public void writeTo(DataOutput out) {
     }
   }
 
+  /**
+   * Get the respective {@link FSTReader} of the {@link DataOutput}. To call this method, you need
+   * to use the default DataOutput or {@link #getOnHeapReaderWriter(int)}, otherwise we will throw
+   * an exception.
+   *
+   * @return the DataOutput as FSTReader
+   * @throws IllegalStateException if the DataOutput does not implement FSTReader
+   */
+  public FSTReader getFSTReader() {
+    if (dataOutput instanceof FSTReader) {
+      return (FSTReader) dataOutput;
+    }
+    throw new IllegalStateException(
+        "The DataOutput must implement FSTReader, but got " + dataOutput);
+  }
+
   /**
    * Fluent-style constructor for FST {@link FSTCompiler}.
    *
@@ -967,10 +972,31 @@ private boolean validOutput(T output) {
     return output == NO_OUTPUT || !output.equals(NO_OUTPUT);
   }
 
-  /** Returns final FST. NOTE: this will return null if nothing is accepted by the FST. */
-  // TODO: make this method to only return the FSTMetadata and user needs to construct the FST
-  // themselves
-  public FST<T> compile() throws IOException {
+  /**
+   * Returns the metadata of the final FST. NOTE: this will return null if nothing is accepted by
+   * the FST themselves.
+   *
+   * <p>To create the FST, you need to:
+   *
+   * <p>- If a FSTReader DataOutput was used, such as the one returned by {@link
+   * #getOnHeapReaderWriter(int)}
+   *
+   * <pre class="prettyprint">
+   *     fstMetadata = fstCompiler.compile();
+   *     fst = FST.fromFSTReader(fstMetadata, fstCompiler.getFSTReader());
+   * </pre>
+   *
+   * <p>- If a non-FSTReader DataOutput was used, such as {@link
+   * org.apache.lucene.store.IndexOutput}, you need to first create the corresponding {@link
+   * org.apache.lucene.store.DataInput}, such as {@link org.apache.lucene.store.IndexInput} then
+   * pass it to the FST construct
+   *
+   * <pre class="prettyprint">
+   *     fstMetadata = fstCompiler.compile();
+   *     fst = new FST&lt;&gt;(fstMetadata, dataInput, new OffHeapFSTStore());
+   * </pre>
+   */
+  public FST.FSTMetadata<T> compile() throws IOException {
 
     final UnCompiledNode<T> root = frontier[0];
 
@@ -990,7 +1016,7 @@ public FST<T> compile() throws IOException {
     // root.output=" + root.output);
     finish(compileNode(root).node);
 
-    return fst;
+    return fst.metadata;
   }
 
   /** Expert: holds a pending (seen but not yet serialized) arc. */

lucene/core/src/java/org/apache/lucene/util/fst/package-info.java

@@ -49,7 +49,7 @@
  *       scratchBytes.copyChars(inputValues[i]);
  *       fstCompiler.add(Util.toIntsRef(scratchBytes.toBytesRef(), scratchInts), outputValues[i]);
  *     }
- *     FST<Long> fst = fstCompiler.compile();
+ *     FST<Long> fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
  * </pre>
  *
  * Retrieval by key:

lucene/core/src/test/org/apache/lucene/util/fst/Test2BFST.java

@@ -90,7 +90,7 @@ public void test() throws Exception {
           nextInput(r, ints2);
         }
 
-        FST<Object> fst = fstCompiler.compile();
+        FST<Object> fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
 
         for (int verify = 0; verify < 2; verify++) {
           System.out.println(
@@ -183,7 +183,7 @@ public void test() throws Exception {
           nextInput(r, ints);
         }
 
-        FST<BytesRef> fst = fstCompiler.compile();
+        FST<BytesRef> fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
         for (int verify = 0; verify < 2; verify++) {
 
           System.out.println(
@@ -273,7 +273,7 @@ public void test() throws Exception {
           nextInput(r, ints);
         }
 
-        FST<Long> fst = fstCompiler.compile();
+        FST<Long> fst = FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
 
         for (int verify = 0; verify < 2; verify++) {
 

lucene/core/src/test/org/apache/lucene/util/fst/Test2BFSTOffHeap.java

@@ -92,10 +92,10 @@ public void test() throws Exception {
         nextInput(r, ints2);
       }
 
-      FST<Object> fst = fstCompiler.compile();
+      FST.FSTMetadata<Object> fstMetadata = fstCompiler.compile();
       indexOutput.close();
       try (IndexInput indexInput = dir.openInput("fst", IOContext.DEFAULT)) {
-        fst = new FST<>(fst.getMetadata(), indexInput, new OffHeapFSTStore());
+        FST<Object> fst = new FST<>(fstMetadata, indexInput, new OffHeapFSTStore());
 
         for (int verify = 0; verify < 2; verify++) {
           System.out.println(
@@ -180,10 +180,10 @@ public void test() throws Exception {
         nextInput(r, ints);
       }
 
-      FST<BytesRef> fst = fstCompiler.compile();
+      FST.FSTMetadata<BytesRef> fstMetadata = fstCompiler.compile();
       indexOutput.close();
       try (IndexInput indexInput = dir.openInput("fst", IOContext.DEFAULT)) {
-        fst = new FST<>(fst.getMetadata(), indexInput, new OffHeapFSTStore());
+        FST<BytesRef> fst = new FST<>(fstMetadata, indexInput, new OffHeapFSTStore());
         for (int verify = 0; verify < 2; verify++) {
 
           System.out.println(
@@ -265,10 +265,10 @@ public void test() throws Exception {
         nextInput(r, ints);
       }
 
-      FST<Long> fst = fstCompiler.compile();
+      FST.FSTMetadata<Long> fstMetadata = fstCompiler.compile();
       indexOutput.close();
       try (IndexInput indexInput = dir.openInput("fst", IOContext.DEFAULT)) {
-        fst = new FST<>(fst.getMetadata(), indexInput, new OffHeapFSTStore());
+        FST<Long> fst = new FST<>(fstMetadata, indexInput, new OffHeapFSTStore());
 
         for (int verify = 0; verify < 2; verify++) {
 

lucene/core/src/test/org/apache/lucene/util/fst/TestFSTDirectAddressing.java

@@ -196,7 +196,7 @@ private static FST<Object> buildFST(List<BytesRef> entries, FSTCompiler<Object>
       }
       last = entry;
     }
-    return fstCompiler.compile();
+    return FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
   }
 
   public static void main(String... args) throws Exception {
@@ -333,7 +333,7 @@ private static FST<CharsRef> recompile(FST<CharsRef> fst, float oversizingFactor
     while ((inputOutput = fstEnum.next()) != null) {
       fstCompiler.add(inputOutput.input, CharsRef.deepCopyOf(inputOutput.output));
     }
-    return fstCompiler.compile();
+    return FST.fromFSTReader(fstCompiler.compile(), fstCompiler.getFSTReader());
   }
 
   private static int walk(FST<CharsRef> read) throws IOException {