From 8b9bb8fb53bb22a8181ac85e50ee7b2eb9269ee0 Mon Sep 17 00:00:00 2001
From: Ivan Sopov <i.sopov@corp.mail.ru>
Date: Fri, 26 Apr 2019 06:23:50 +0300
Subject: [PATCH] [core] Optimize toArray of all ByteIterators except
 RandomByteIterator (#1112)

Method nextBuf is a clever hack that outperforms Random.nextBytes
but performs poorly for all other ByteIterator implementations

This commit moves it to RandomByteIterator and adds efficient
toArray implementations for other ByteIterator classes.

Also InputStreamByteIterator.reset method that unconditionally
throws UnsupportedOperationException is fixed
---
 .../com/yahoo/ycsb/ByteArrayByteIterator.java | 13 +++++++++--
 .../java/com/yahoo/ycsb/ByteIterator.java     |  5 ++--
 .../yahoo/ycsb/InputStreamByteIterator.java   | 23 ++++++++++++++++---
 .../com/yahoo/ycsb/RandomByteIterator.java    | 20 +++++++++++++---
 .../com/yahoo/ycsb/StringByteIterator.java    | 12 +++++++++-
 5 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/core/src/main/java/com/yahoo/ycsb/ByteArrayByteIterator.java b/core/src/main/java/com/yahoo/ycsb/ByteArrayByteIterator.java
index 8f8762b4..cf1470cb 100644
--- a/core/src/main/java/com/yahoo/ycsb/ByteArrayByteIterator.java
+++ b/core/src/main/java/com/yahoo/ycsb/ByteArrayByteIterator.java
@@ -21,7 +21,7 @@ package com.yahoo.ycsb;
  */
 public class ByteArrayByteIterator extends ByteIterator {
   private final int originalOffset;
-  private byte[] str;
+  private final byte[] str;
   private int off;
   private final int len;
 
@@ -60,5 +60,14 @@ public class ByteArrayByteIterator extends ByteIterator {
   public void reset() {
     off = originalOffset;
   }
-  
+
+  @Override
+  public byte[] toArray() {
+    int size = (int) bytesLeft();
+    byte[] bytes = new byte[size];
+    System.arraycopy(str, off, bytes, 0, size);
+    off = len;
+    return bytes;
+  }
+
 }
diff --git a/core/src/main/java/com/yahoo/ycsb/ByteIterator.java b/core/src/main/java/com/yahoo/ycsb/ByteIterator.java
index 9be84d5a..616675ce 100644
--- a/core/src/main/java/com/yahoo/ycsb/ByteIterator.java
+++ b/core/src/main/java/com/yahoo/ycsb/ByteIterator.java
@@ -96,9 +96,8 @@ public abstract class ByteIterator implements Iterator<Byte> {
       throw new ArrayIndexOutOfBoundsException("Too much data to fit in one array!");
     }
     byte[] ret = new byte[(int) left];
-    int off = 0;
-    while (off < ret.length) {
-      off = nextBuf(ret, off);
+    for (int i = 0; i < ret.length; i++) {
+      ret[i] = nextByte();
     }
     return ret;
   }
diff --git a/core/src/main/java/com/yahoo/ycsb/InputStreamByteIterator.java b/core/src/main/java/com/yahoo/ycsb/InputStreamByteIterator.java
index 02ca3800..31c37065 100644
--- a/core/src/main/java/com/yahoo/ycsb/InputStreamByteIterator.java
+++ b/core/src/main/java/com/yahoo/ycsb/InputStreamByteIterator.java
@@ -23,8 +23,8 @@ import java.io.InputStream;
  *  A ByteIterator that iterates through an inputstream of bytes.
  */
 public class InputStreamByteIterator extends ByteIterator {
-  private long len;
-  private InputStream ins;
+  private final long len;
+  private final InputStream ins;
   private long off;
   private final boolean resetable;
 
@@ -63,17 +63,34 @@ public class InputStreamByteIterator extends ByteIterator {
     return len - off;
   }
 
+  @Override
+  public byte[] toArray() {
+    int size = (int) bytesLeft();
+    byte[] bytes = new byte[size];
+    try {
+      if (ins.read(bytes) < size) {
+        throw new IllegalStateException("Past EOF!");
+      }
+    } catch (IOException e) {
+      throw new IllegalStateException(e);
+    }
+    off = len;
+    return bytes;
+  }
+  
   @Override
   public void reset() {
     if (resetable) {
       try {
         ins.reset();
         ins.mark((int) len);
+        off = 0;
       } catch (IOException e) {
         throw new IllegalStateException("Failed to reset the input stream", e);
       }
+    } else {
+      throw new UnsupportedOperationException();
     }
-    throw new UnsupportedOperationException();
   }
   
 }
diff --git a/core/src/main/java/com/yahoo/ycsb/RandomByteIterator.java b/core/src/main/java/com/yahoo/ycsb/RandomByteIterator.java
index e21689cc..33632661 100644
--- a/core/src/main/java/com/yahoo/ycsb/RandomByteIterator.java
+++ b/core/src/main/java/com/yahoo/ycsb/RandomByteIterator.java
@@ -22,10 +22,10 @@ import java.util.concurrent.ThreadLocalRandom;
  *  A ByteIterator that generates a random sequence of bytes.
  */
 public class RandomByteIterator extends ByteIterator {
-  private long len;
+  private final long len;
   private long off;
   private int bufOff;
-  private byte[] buf;
+  private final byte[] buf;
 
   @Override
   public boolean hasNext() {
@@ -100,5 +100,19 @@ public class RandomByteIterator extends ByteIterator {
   public void reset() {
     off = 0;
   }
-  
+
+  /** Consumes remaining contents of this object, and returns them as a byte array. */
+  public byte[] toArray() {
+    long left = bytesLeft();
+    if (left != (int) left) {
+      throw new ArrayIndexOutOfBoundsException("Too much data to fit in one array!");
+    }
+    byte[] ret = new byte[(int) left];
+    int bufOffset = 0;
+    while (bufOffset < ret.length) {
+      bufOffset = nextBuf(ret, bufOffset);
+    }
+    return ret;
+  }
+
 }
diff --git a/core/src/main/java/com/yahoo/ycsb/StringByteIterator.java b/core/src/main/java/com/yahoo/ycsb/StringByteIterator.java
index 82235491..63dba894 100644
--- a/core/src/main/java/com/yahoo/ycsb/StringByteIterator.java
+++ b/core/src/main/java/com/yahoo/ycsb/StringByteIterator.java
@@ -100,7 +100,17 @@ public class StringByteIterator extends ByteIterator {
   public void reset() {
     off = 0;
   }
-  
+
+  @Override
+  public byte[] toArray() {
+    byte[] bytes = new byte[(int) bytesLeft()];
+    for (int i = 0; i < bytes.length; i++) {
+      bytes[i] = (byte) str.charAt(off + i);
+    }
+    off = str.length();
+    return bytes;
+  }
+
   /**
    * Specialization of general purpose toString() to avoid unnecessary
    * copies.
-- 
GitLab