Allow FST builder to use different writer (#12543)

Author: dungba88

lucene/core/src/java/org/apache/lucene/store/FileChannelDataOutput.java

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.store;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.channels.FileChannel;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * A {@link DataOutput} which writes to a {@link FileChannel}
+ *
+ * @lucene.experimental
+ */
+public class FileChannelDataOutput extends DataOutput implements Accountable {
+
+  private static final long BASE_RAM_BYTES_USED =
+      RamUsageEstimator.shallowSizeOfInstance(FileChannelDataOutput.class);
+
+  // buffer to write to the channel
+  private ByteBuffer buf = ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN);
+
+  private final FileChannel channel;
+
+  /**
+   * ctor
+   *
+   * @param channel the channel to write to
+   */
+  public FileChannelDataOutput(FileChannel channel) {
+    this.channel = channel;
+  }
+
+  @Override
+  public void writeByte(byte b) throws IOException {
+    buf.clear().put(b).flip();
+    this.channel.write(buf);
+  }
+
+  @Override
+  public void writeBytes(byte[] b, int offset, int length) throws IOException {
+    this.channel.write(ByteBuffer.wrap(b, offset, length).order(ByteOrder.LITTLE_ENDIAN));
+  }
+
+  @Override
+  public void writeInt(int i) throws IOException {
+    buf.clear().putInt(i).flip();
+    this.channel.write(buf);
+  }
+
+  @Override
+  public void writeShort(short i) throws IOException {
+    buf.clear().putShort(i).flip();
+    this.channel.write(buf);
+  }
+
+  @Override
+  public void writeLong(long i) throws IOException {
+    buf.clear().putLong(i).flip();
+    this.channel.write(buf);
+  }
+
+  public FileChannel getChannel() {
+    return channel;
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    return BASE_RAM_BYTES_USED + buf.capacity();
+  }
+}

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

@@ -26,7 +26,7 @@
 // TODO: merge with PagedBytes, except PagedBytes doesn't
 // let you read while writing which FST needs
 
-class BytesStore extends DataOutput implements FSTReader {
+class BytesStore extends DataOutput implements FSTWriter {
 
   private static final long BASE_RAM_BYTES_USED =
       RamUsageEstimator.shallowSizeOfInstance(BytesStore.class)
@@ -359,6 +359,7 @@ public void truncate(long newLen) {
     assert newLen == getPosition();
   }
 
+  @Override
   public void finish() {
     if (current != null) {
       byte[] lastBuffer = new byte[nextWrite];
@@ -368,6 +369,16 @@ public void finish() {
     }
   }
 
+  /** Writes all of our bytes to the target {@link FSTWriter}. */
+  public void writeTo(FSTWriter out) throws IOException {
+    // TODO: if the FSTWriter is also BytesStore we are doing double write
+    // once to reverse the bytes and once to write to the BytesStore
+    // maybe we should combine it into reverseAndWriteTo()?
+    for (byte[] block : blocks) {
+      out.writeBytes(block, 0, block.length);
+    }
+  }
+
   /** Writes all of our bytes to the target {@link DataOutput}. */
   @Override
   public void writeTo(DataOutput out) throws IOException {
@@ -447,6 +458,11 @@ public FST.BytesReader getReverseBytesReader() {
     return getReverseReader(true);
   }
 
+  @Override
+  public FST.BytesReader getReverseReaderForSuffixSharing() {
+    return getReverseReader(false);
+  }
+
   FST.BytesReader getReverseReader(boolean allowSingle) {
     if (allowSingle && blocks.size() == 1) {
       return new ReverseBytesReader(blocks.get(0));

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

@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.util.fst;
+
+import java.io.Closeable;
+import java.io.IOException;
+import org.apache.lucene.store.DataOutput;
+import org.apache.lucene.store.RandomAccessInput;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
+
+/**
+ * An {@link FSTWriter} which write to a DataOutput
+ *
+ * @lucene.experimental
+ */
+public class DataOutputFSTWriter implements FSTWriter {
+
+  private static final long BASE_RAM_BYTES_USED =
+      RamUsageEstimator.shallowSizeOfInstance(DataOutputFSTWriter.class);
+
+  private final DataOutput dataOutput;
+
+  private final RandomAccessInput dataInput;
+
+  protected boolean finish = false;
+
+  private long size = 0L;
+
+  /**
+   * ctor
+   *
+   * @param dataOutput the data output to write to
+   * @param dataInput the data input containing the written bytes to read from
+   */
+  public DataOutputFSTWriter(DataOutput dataOutput, RandomAccessInput dataInput) {
+    this.dataOutput = dataOutput;
+    this.dataInput = dataInput;
+  }
+
+  @Override
+  public long ramBytesUsed() {
+    long size = BASE_RAM_BYTES_USED;
+    if (dataOutput instanceof Accountable) {
+      size += ((Accountable) dataOutput).ramBytesUsed();
+    }
+    return size;
+  }
+
+  @Override
+  public long size() {
+    return size;
+  }
+
+  @Override
+  public void writeByte(byte b) throws IOException {
+    size++;
+    dataOutput.writeByte(b);
+  }
+
+  @Override
+  public void writeBytes(byte[] b, int offset, int length) throws IOException {
+    size += length;
+    dataOutput.writeBytes(b, offset, length);
+  }
+
+  @Override
+  public void finish() throws IOException {
+    finish = true;
+    if (dataOutput instanceof Closeable) {
+      ((Closeable) dataOutput).close();
+    }
+  }
+
+  @Override
+  public void writeTo(DataOutput out) throws IOException {
+    // Technically we can support this method, as the DataOutput by this time has already been
+    // closed.
+    // But allow the FST which is already written to a DataOutput to be saved to another DataOutput
+    // would be rather a strange use case
+    throw new UnsupportedOperationException("writeTo(DataOutput) is not supported by this class");
+  }
+
+  @Override
+  public FST.BytesReader getReverseBytesReader() {
+    // Technically we can support this method, as the DataOutput by this time has already been
+    // closed.
+    // However, I think ideally we would want the FSTWriter/FSTCompiler to only write the FST to the
+    // DataOutput, and
+    // some process later on can construct the FST using the FSTStore method.
+    throw new UnsupportedOperationException(
+        "getReverseBytesReader() is not supported by this class");
+  }
+
+  @Override
+  public FST.BytesReader getReverseReaderForSuffixSharing() {
+    return new ReverseRandomAccessReader(dataInput);
+  }
+}

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

@@ -500,6 +500,16 @@ void setEmptyOutput(T v) {
   }
 
   public void save(DataOutput metaOut, DataOutput out) throws IOException {
+    saveMetadata(metaOut);
+    fstReader.writeTo(out);
+  }
+
+  /**
+   * Save the metadata to a DataOutput
+   *
+   * @param metaOut the DataOutput to save
+   */
+  public void saveMetadata(DataOutput metaOut) throws IOException {
     if (startNode == -1) {
       throw new IllegalStateException("call finish first");
     }
@@ -541,7 +551,6 @@ public void save(DataOutput metaOut, DataOutput out) throws IOException {
     metaOut.writeByte(t);
     metaOut.writeVLong(startNode);
     metaOut.writeVLong(numBytes());
-    fstReader.writeTo(out);
   }
 
   /** Writes an automaton to a file. */