[lucene] branch branch_9x updated: LUCENE-10657: CopyBytes now saves one memory copy on ByteBuffersDataOutput (#1034)

[jp...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jpountz pushed a commit to branch branch_9x
in repository https://gitbox.apache.org/repos/asf/lucene.git


The following commit(s) were added to refs/heads/branch_9x by this push:
     new 7328ad2dafc LUCENE-10657: CopyBytes now saves one memory copy on ByteBuffersDataOutput (#1034)
7328ad2dafc is described below

commit 7328ad2dafc2e78bf2950cb4bdd2c8785f31f7b9
Author: luyuncheng <lu...@bytedance.com>
AuthorDate: Tue Jul 19 15:37:07 2022 +0800

    LUCENE-10657: CopyBytes now saves one memory copy on ByteBuffersDataOutput (#1034)
    
    Abstract method copyBytes need to copy from input to a buffer and then write into ByteBuffersDataOutput, i think there is unnecessary, we can override it, copy directly from input into output
---
 lucene/CHANGES.txt                                 |  2 ++
 .../apache/lucene/store/ByteBuffersDataOutput.java | 23 ++++++++++++++
 .../lucene/store/TestByteBuffersDataOutput.java    | 36 ++++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 508a4ae092d..f422c5da0d8 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -60,6 +60,8 @@ Optimizations
 
 * GITHUB#1010: Specialize ordinal encoding for common case in SortedSetDocValues. (Greg Miller)
 
+* LUCENE-10657: CopyBytes now saves one memory copy on ByteBuffersDataOutput. (luyuncheng)
+
 Changes in runtime behavior
 ---------------------
 
diff --git a/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java b/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java
index 8639ac5522e..84a7b317d88 100644
--- a/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java
+++ b/lucene/core/src/java/org/apache/lucene/store/ByteBuffersDataOutput.java
@@ -309,6 +309,29 @@ public final class ByteBuffersDataOutput extends DataOutput implements Accountab
     return arr;
   }
 
+  @Override
+  public void copyBytes(DataInput input, long numBytes) throws IOException {
+    assert numBytes >= 0 : "numBytes=" + numBytes;
+    int length = (int) numBytes;
+    while (length > 0) {
+      if (!currentBlock.hasRemaining()) {
+        appendBlock();
+      }
+      if (!currentBlock.hasArray()) {
+        break;
+      }
+      int chunk = Math.min(currentBlock.remaining(), length);
+      final int pos = currentBlock.position();
+      input.readBytes(currentBlock.array(), Math.addExact(currentBlock.arrayOffset(), pos), chunk);
+      length -= chunk;
+      currentBlock.position(pos + chunk);
+    }
+    // if current block is Direct, we fall back to super.copyBytes for remaining bytes
+    if (length > 0) {
+      super.copyBytes(input, length);
+    }
+  }
+
   /** Copy the current content of this object into another {@link DataOutput}. */
   public void copyTo(DataOutput output) throws IOException {
     for (ByteBuffer bb : blocks) {
diff --git a/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java b/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java
index 065606c3ba3..50eefc3f3bf 100644
--- a/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java
+++ b/lucene/core/src/test/org/apache/lucene/store/TestByteBuffersDataOutput.java
@@ -204,6 +204,42 @@ public final class TestByteBuffersDataOutput extends BaseDataOutputTestCase<Byte
         ArrayUtil.copyOfSubArray(bytes, offset, offset + len), o.toArrayCopy());
   }
 
+  @Test
+  public void testCopyBytesOnHeap() throws IOException {
+    byte[] bytes = new byte[1024 * 8 + 10];
+    random().nextBytes(bytes);
+    int offset = TestUtil.nextInt(random(), 0, 100);
+    int len = bytes.length - offset;
+    ByteArrayDataInput in = new ByteArrayDataInput(bytes, offset, len);
+    ByteBuffersDataOutput o =
+        new ByteBuffersDataOutput(
+            ByteBuffersDataOutput.DEFAULT_MIN_BITS_PER_BLOCK,
+            ByteBuffersDataOutput.DEFAULT_MAX_BITS_PER_BLOCK,
+            ByteBuffersDataOutput.ALLOCATE_BB_ON_HEAP,
+            ByteBuffersDataOutput.NO_REUSE);
+    o.copyBytes(in, len);
+    Assert.assertArrayEquals(
+        o.toArrayCopy(), ArrayUtil.copyOfSubArray(bytes, offset, offset + len));
+  }
+
+  @Test
+  public void testCopyBytesOnDirectByteBuffer() throws IOException {
+    byte[] bytes = new byte[1024 * 8 + 10];
+    random().nextBytes(bytes);
+    int offset = TestUtil.nextInt(random(), 0, 100);
+    int len = bytes.length - offset;
+    ByteArrayDataInput in = new ByteArrayDataInput(bytes, offset, len);
+    ByteBuffersDataOutput o =
+        new ByteBuffersDataOutput(
+            ByteBuffersDataOutput.DEFAULT_MIN_BITS_PER_BLOCK,
+            ByteBuffersDataOutput.DEFAULT_MAX_BITS_PER_BLOCK,
+            ByteBuffer::allocateDirect,
+            ByteBuffersDataOutput.NO_REUSE);
+    o.copyBytes(in, len);
+    Assert.assertArrayEquals(
+        o.toArrayCopy(), ArrayUtil.copyOfSubArray(bytes, offset, offset + len));
+  }
+
   @Test
   public void testToBufferListReturnsReadOnlyBuffers() throws Exception {
     ByteBuffersDataOutput dst = new ByteBuffersDataOutput();