LUCENE-10125: Optimize primitive writes in OutputStreamIndexOutput (#321)

lucene/CHANGES.txt

@@ -282,6 +282,9 @@ Improvements
 * LUCENE-10112: Improve LZ4 Compression performance with direct primitive read/writes.
   (Tim Brooks, Uwe Schindler, Robert Muir, Adrien Grand)
 
+* LUCENE-10125: Optimize primitive writes in OutputStreamIndexOutput.
+  (Uwe Schindler, Robert Muir, Adrien Grand)
+
 Bug fixes
 
 * LUCENE-10070 Skip deleted docs when accumulating facet counts for all docs. (Ankur Goel, Greg Miller)

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

@@ -21,12 +21,13 @@ import java.io.IOException;
 import java.io.OutputStream;
 import java.util.zip.CRC32;
 import java.util.zip.CheckedOutputStream;
+import org.apache.lucene.util.BitUtil;
 
 /** Implementation class for buffered {@link IndexOutput} that writes to an {@link OutputStream}. */
 public class OutputStreamIndexOutput extends IndexOutput {
 
   private final CRC32 crc = new CRC32();
-  private final BufferedOutputStream os;
+  private final XBufferedOutputStream os;
 
   private long bytesWritten = 0L;
   private boolean flushedOnClose = false;
@@ -35,12 +36,16 @@ public class OutputStreamIndexOutput extends IndexOutput {
    * Creates a new {@link OutputStreamIndexOutput} with the given buffer size.
    *
    * @param bufferSize the buffer size in bytes used to buffer writes internally.
-   * @throws IllegalArgumentException if the given buffer size is less or equal to <code>0</code>
+   * @throws IllegalArgumentException if the given buffer size is less than <code>
+   *     {@value Long#BYTES}</code>
    */
   public OutputStreamIndexOutput(
       String resourceDescription, String name, OutputStream out, int bufferSize) {
     super(resourceDescription, name);
-    this.os = new BufferedOutputStream(new CheckedOutputStream(out, crc), bufferSize);
+    if (bufferSize < Long.BYTES) {
+      throw new IllegalArgumentException("Buffer size too small, need: " + Long.BYTES);
+    }
+    this.os = new XBufferedOutputStream(new CheckedOutputStream(out, crc), bufferSize);
   }
 
   @Override
@@ -55,6 +60,24 @@ public class OutputStreamIndexOutput extends IndexOutput {
     bytesWritten += length;
   }
 
+  @Override
+  public void writeShort(short i) throws IOException {
+    os.writeShort(i);
+    bytesWritten += Short.BYTES;
+  }
+
+  @Override
+  public void writeInt(int i) throws IOException {
+    os.writeInt(i);
+    bytesWritten += Integer.BYTES;
+  }
+
+  @Override
+  public void writeLong(long i) throws IOException {
+    os.writeLong(i);
+    bytesWritten += Long.BYTES;
+  }
+
   @Override
   public void close() throws IOException {
     try (final OutputStream o = os) {
@@ -81,4 +104,36 @@ public class OutputStreamIndexOutput extends IndexOutput {
     os.flush();
     return crc.getValue();
   }
+
+  /** This subclass is an optimization for writing primitives. Don't use outside of this class! */
+  private static final class XBufferedOutputStream extends BufferedOutputStream {
+
+    XBufferedOutputStream(OutputStream out, int size) {
+      super(out, size);
+    }
+
+    private void flushIfNeeded(int len) throws IOException {
+      if (len > buf.length - count) {
+        flush();
+      }
+    }
+
+    void writeShort(short i) throws IOException {
+      flushIfNeeded(Short.BYTES);
+      BitUtil.VH_LE_SHORT.set(buf, count, i);
+      count += Short.BYTES;
+    }
+
+    void writeInt(int i) throws IOException {
+      flushIfNeeded(Integer.BYTES);
+      BitUtil.VH_LE_INT.set(buf, count, i);
+      count += Integer.BYTES;
+    }
+
+    void writeLong(long i) throws IOException {
+      flushIfNeeded(Long.BYTES);
+      BitUtil.VH_LE_LONG.set(buf, count, i);
+      count += Long.BYTES;
+    }
+  }
 }

lucene/core/src/test/org/apache/lucene/store/TestOutputStreamIndexOutput.java

@@ -0,0 +1,54 @@
+/*
+ * 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.ByteArrayOutputStream;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import org.apache.lucene.util.LuceneTestCase;
+
+public class TestOutputStreamIndexOutput extends LuceneTestCase {
+
+  public void testDataTypes() throws Exception {
+    for (int i = 0; i < 12; i++) {
+      doTestDataTypes(i);
+    }
+  }
+
+  private void doTestDataTypes(int offset) throws Exception {
+    final ByteArrayOutputStream bos = new ByteArrayOutputStream();
+    final IndexOutput out = new OutputStreamIndexOutput("test" + offset, "test", bos, 12);
+    for (int i = 0; i < offset; i++) {
+      out.writeByte((byte) i);
+    }
+    out.writeShort((short) 12345);
+    out.writeInt(1234567890);
+    out.writeLong(1234567890123456789L);
+    assertEquals(offset + 14, out.getFilePointer());
+    out.close();
+
+    // read the primitives using ByteBuffer to ensure encoding in byte array is LE:
+    final ByteBuffer buf = ByteBuffer.wrap(bos.toByteArray()).order(ByteOrder.LITTLE_ENDIAN);
+    for (int i = 0; i < offset; i++) {
+      assertEquals(i, buf.get());
+    }
+    assertEquals(12345, buf.getShort());
+    assertEquals(1234567890, buf.getInt());
+    assertEquals(1234567890123456789L, buf.getLong());
+    assertEquals(0, buf.remaining());
+  }
+}