Use List<byte[]> instead of ByteBlockPool

Author: dungba88

lucene/core/src/java/org/apache/lucene/util/ByteBlockPool.java

@@ -256,18 +256,6 @@ public void readBytes(final long offset, final byte[] bytes, int bytesOffset, in
     }
   }
 
-  /**
-   * Read a single byte at the given offset
-   *
-   * @param offset the offset to read
-   * @return the byte
-   */
-  public byte readByte(final long offset) {
-    int bufferIndex = (int) (offset >> BYTE_BLOCK_SHIFT);
-    int pos = (int) (offset & BYTE_BLOCK_MASK);
-    return buffers[bufferIndex][pos];
-  }
-
   @Override
   public long ramBytesUsed() {
     long size = BASE_RAM_BYTES;

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

@@ -292,13 +292,13 @@ private CompiledNode compileNode(UnCompiledNode<T> nodeIn, int tailLength) throw
     long bytesPosStart = bytes.getPosition();
     if (dedupHash != null) {
       if (nodeIn.numArcs == 0) {
-        node = addNode(nodeIn).nodeAddress;
+        node = addNode(nodeIn, false).nodeAddress;
         lastFrozenNode = node;
       } else {
         node = dedupHash.add(nodeIn);
       }
     } else {
-      node = addNode(nodeIn).nodeAddress;
+      node = addNode(nodeIn, false).nodeAddress;
     }
     assert node != -2;
 
@@ -318,7 +318,7 @@ private CompiledNode compileNode(UnCompiledNode<T> nodeIn, int tailLength) throw
 
   // serializes new node by appending its bytes to the end
   // of the current byte[]
-  NodeAndBuffer addNode(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
+  NodeAndBuffer addNode(FSTCompiler.UnCompiledNode<T> nodeIn, boolean needCopy) throws IOException {
     // System.out.println("FST.addNode pos=" + bytes.getPosition() + " numArcs=" + nodeIn.numArcs);
     if (nodeIn.numArcs == 0) {
       if (nodeIn.isFinal) {
@@ -461,13 +461,24 @@ NodeAndBuffer addNode(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
     final long thisNodeAddress = bytes.getPosition() - 1;
     bytes.reverse(startAddress, thisNodeAddress);
     nodeCount++;
-    // nocommit: this is non-optimal, we should write the BytesStore to the ByteBlockPool directly
-    byte[] buf = new byte[Math.toIntExact(thisNodeAddress - startAddress + 1)];
-    bytes.copyBytes(startAddress, buf, 0, buf.length);
-    return new NodeAndBuffer(thisNodeAddress, buf);
+    if (needCopy) {
+      byte[] buf = new byte[Math.toIntExact(thisNodeAddress - startAddress + 1)];
+      bytes.copyBytes(startAddress, buf, 0, buf.length);
+      return new NodeAndBuffer(thisNodeAddress, buf);
+    }
+    return new NodeAndBuffer(thisNodeAddress, null);
   }
 
-  record NodeAndBuffer(long nodeAddress, byte[] bytes) {}
+  class NodeAndBuffer {
+
+    final long nodeAddress;
+    final byte[] bytes;
+
+    NodeAndBuffer(long nodeAddress, byte[] bytes) {
+      this.nodeAddress = nodeAddress;
+      this.bytes = bytes;
+    }
+  }
 
   private void writeLabel(DataOutput out, int v) throws IOException {
     assert v >= 0 : "v=" + v;

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

@@ -17,10 +17,10 @@
 package org.apache.lucene.util.fst;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
-import org.apache.lucene.util.ByteBlockPool;
-import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.packed.PackedInts;
 import org.apache.lucene.util.packed.PagedGrowableWriter;
 
@@ -117,14 +117,14 @@ public long add(FSTCompiler.UnCompiledNode<T> nodeIn) throws IOException {
           // not in fallback either -- freeze & add the incoming node
 
           // freeze & add
-          FSTCompiler.NodeAndBuffer nodeAndBuffer = fstCompiler.addNode(nodeIn);
-          node = nodeAndBuffer.nodeAddress();
+          FSTCompiler.NodeAndBuffer nodeAndBuffer = fstCompiler.addNode(nodeIn, true);
+          node = nodeAndBuffer.nodeAddress;
 
           // we use 0 as empty marker in hash table, so it better be impossible to get a frozen node
           // at 0:
           assert node != 0;
 
-          primaryTable.set(pos, node, nodeAndBuffer.bytes());
+          primaryTable.set(pos, node, nodeAndBuffer.bytes);
 
           // confirm frozen hash and unfrozen hash are the same
           assert hash(node) == hash : "mismatch frozenHash=" + hash(node) + " vs hash=" + hash;
@@ -192,9 +192,7 @@ private long hash(FSTCompiler.UnCompiledNode<T> node) {
   // hash code for a frozen node.  this must precisely match the hash computation of an unfrozen
   // node!
   private long hash(long node) throws IOException {
-    long offset = primaryTable.copiedOffsets.get(node).offset;
-    FST.BytesReader in =
-        new ByteBlockPoolReverseBytesReader(primaryTable.copiedNodes, node - offset);
+    FST.BytesReader in = getBytesReader(primaryTable, node);
 
     final int PRIME = 31;
 
@@ -224,8 +222,7 @@ private long hash(long node) throws IOException {
   private boolean nodesEqual(
       PagedGrowableHash table, FSTCompiler.UnCompiledNode<T> node, long address)
       throws IOException {
-    long offset = table.copiedOffsets.get(address).offset;
-    FST.BytesReader in = new ByteBlockPoolReverseBytesReader(table.copiedNodes, address - offset);
+    FST.BytesReader in = getBytesReader(table, address);
     fstCompiler.fst.readFirstRealTargetArc(address, scratchArc, in);
 
     // fail fast for a node with fixed length arcs
@@ -280,20 +277,22 @@ private boolean nodesEqual(
     return false;
   }
 
-  // nocommit: change this to just offset and somehow find the length
-  record OffsetAndLength(long offset, int length) {}
+  private static <T> FST.BytesReader getBytesReader(
+      NodeHash<T>.PagedGrowableHash table, long address) {
+    byte[] bytes = table.getBytes(address);
+    return new RelativeReverseBytesReader(bytes, address - bytes.length + 1);
+  }
 
   /** Inner class because it needs access to hash function and FST bytes. */
   private class PagedGrowableHash {
     private PagedGrowableWriter entries;
     // nocommit: use PagedGrowableWriter? there was some size overflow issue with
     // PagedGrowableWriter
     // mapping from FST real address to copiedNodes offsets & length
-    private Map<Long, OffsetAndLength> copiedOffsets;
+    private Map<Long, Integer> copiedOffsets;
     long count;
-    long currentOffsets = -1;
     private long mask;
-    private final ByteBlockPool copiedNodes;
+    private final List<byte[]> copiedNodes;
 
     // 256K blocks, but note that the final block is sized only as needed so it won't use the full
     // block size when just a few elements were written to it
@@ -303,7 +302,7 @@ public PagedGrowableHash() {
       entries = new PagedGrowableWriter(16, BLOCK_SIZE_BYTES, 8, PackedInts.COMPACT);
       copiedOffsets = new HashMap<>();
       mask = 15;
-      copiedNodes = new ByteBlockPool(new ByteBlockPool.DirectAllocator());
+      copiedNodes = new ArrayList<>();
     }
 
     public PagedGrowableHash(long lastNodeAddress, long size) {
@@ -313,19 +312,11 @@ public PagedGrowableHash(long lastNodeAddress, long size) {
       copiedOffsets = new HashMap<>();
       mask = size - 1;
       assert (mask & size) == 0 : "size must be a power-of-2; got size=" + size + " mask=" + mask;
-      copiedNodes = new ByteBlockPool(new ByteBlockPool.DirectAllocator());
+      copiedNodes = new ArrayList<>();
     }
 
-    public byte[] getBytes(long index) {
-      // nocommit: find a more efficient way to copy from fallback table to primary table
-      // here we need double copying
-      OffsetAndLength offsetAndLength = copiedOffsets.get(index);
-      byte[] bytes = new byte[offsetAndLength.length];
-      // offset is the last offset of the node, hence subtract by (length - 1) to get the first
-      // offset
-      copiedNodes.readBytes(
-          offsetAndLength.offset - offsetAndLength.length + 1, bytes, 0, offsetAndLength.length);
-      return bytes;
+    public byte[] getBytes(long node) {
+      return copiedNodes.get(copiedOffsets.get(node));
     }
 
     public long get(long index) {
@@ -334,11 +325,9 @@ public long get(long index) {
 
     public void set(long index, long pointer, byte[] bytes) {
       entries.set(index, pointer);
-      copiedNodes.append(new BytesRef(bytes));
-      count += 3;
-      currentOffsets += bytes.length;
-      // offsets is the last offset of the node, as we are reading in backward
-      copiedOffsets.put(pointer, new OffsetAndLength(currentOffsets, bytes.length));
+      copiedNodes.add(bytes);
+      copiedOffsets.put(pointer, copiedNodes.size() - 1);
+      count += 2;
     }
 
     private void rehash(long lastNodeAddress) throws IOException {

lucene/core/src/java/org/apache/lucene/util/fst/ByteBlockPoolReverseBytesReader.java renamed to lucene/core/src/java/org/apache/lucene/util/fst/RelativeReverseBytesReader.java

@@ -17,45 +17,41 @@
 package org.apache.lucene.util.fst;
 
 import java.io.IOException;
-import org.apache.lucene.util.ByteBlockPool;
 
-/** Reads in reverse from a ByteBlockPool. */
-final class ByteBlockPoolReverseBytesReader extends FST.BytesReader {
+/** Same as ReverseBytesReader but acts upon a relative position */
+final class RelativeReverseBytesReader extends FST.BytesReader {
 
-  private final ByteBlockPool buf;
+  private final FST.BytesReader reader;
   private final long relativePos;
-  private long pos;
 
-  public ByteBlockPoolReverseBytesReader(ByteBlockPool buf, long relativePos) {
-    this.buf = buf;
+  public RelativeReverseBytesReader(byte[] bytes, long relativePos) {
+    this.reader = new ReverseBytesReader(bytes);
     this.relativePos = relativePos;
   }
 
   @Override
   public byte readByte() throws IOException {
-    return buf.readByte(pos--);
+    return reader.readByte();
   }
 
   @Override
   public void readBytes(byte[] b, int offset, int len) throws IOException {
-    for (int i = 0; i < len; i++) {
-      b[offset + i] = readByte();
-    }
+    reader.readBytes(b, offset, len);
   }
 
   @Override
-  public void skipBytes(long numBytes) throws IOException {
-    pos -= numBytes;
+  public void skipBytes(long count) throws IOException {
+    reader.skipBytes(count);
   }
 
   @Override
   public long getPosition() {
-    return pos + relativePos;
+    return reader.getPosition() + relativePos;
   }
 
   @Override
   public void setPosition(long pos) {
-    this.pos = pos - relativePos;
+    reader.setPosition(pos - relativePos);
   }
 
   @Override

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

@@ -1257,7 +1257,7 @@ public void testNonFinalStopNode() throws Exception {
       node.isFinal = true;
       rootNode.addArc('a', node);
       final FSTCompiler.CompiledNode frozen = new FSTCompiler.CompiledNode();
-      frozen.node = fstCompiler.addNode(node).nodeAddress();
+      frozen.node = fstCompiler.addNode(node, false).nodeAddress;
       rootNode.arcs[0].nextFinalOutput = 17L;
       rootNode.arcs[0].isFinal = true;
       rootNode.arcs[0].output = nothing;
@@ -1270,13 +1270,13 @@ public void testNonFinalStopNode() throws Exception {
           new FSTCompiler.UnCompiledNode<>(fstCompiler, 0);
       rootNode.addArc('b', node);
       final FSTCompiler.CompiledNode frozen = new FSTCompiler.CompiledNode();
-      frozen.node = fstCompiler.addNode(node).nodeAddress();
+      frozen.node = fstCompiler.addNode(node, false).nodeAddress;
       rootNode.arcs[1].nextFinalOutput = nothing;
       rootNode.arcs[1].output = 42L;
       rootNode.arcs[1].target = frozen;
     }
 
-    fst.finish(fstCompiler.addNode(rootNode).nodeAddress());
+    fst.finish(fstCompiler.addNode(rootNode, false).nodeAddress);
 
     StringWriter w = new StringWriter();
     // Writer w = new OutputStreamWriter(new FileOutputStream("/x/tmp3/out.dot"));