Avoid BytesRef's copying in ScriptDocValues's Strings

server/src/main/java/org/elasticsearch/index/fielddata/ScriptDocValues.java

@@ -35,6 +35,7 @@
 import org.joda.time.ReadableDateTime;
 
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import java.util.AbstractList;
@@ -570,90 +571,115 @@ private static boolean[] grow(boolean[] array, int minSize) {
             } else
                 return array;
         }
-
     }
 
-    abstract static class BinaryScriptDocValues<T> extends ScriptDocValues<T> {
-
+    public static final class Strings extends ScriptDocValues<String> {
         private final SortedBinaryDocValues in;
-        protected BytesRefBuilder[] values = new BytesRefBuilder[0];
-        protected int count;
+        private String[] values = new String[0];
+        private int count;
+        private boolean valuesSet = false;
 
-        BinaryScriptDocValues(SortedBinaryDocValues in) {
+        public Strings(SortedBinaryDocValues in) {
             this.in = in;
         }
 
         @Override
-        public void setNextDocId(int docId) throws IOException {
-            if (in.advanceExact(docId)) {
-                resize(in.docValueCount());
-                for (int i = 0; i < count; i++) {
-                    // We need to make a copy here, because BytesBinaryDVAtomicFieldData's SortedBinaryDocValues
-                    // implementation reuses the returned BytesRef. Otherwise we would end up with the same BytesRef
-                    // instance for all slots in the values array.
-                    values[i].copyBytes(in.nextValue());
+        public String get(int index) {
+            if (valuesSet == false) {
+                try {
+                    fillValues();
+                } catch (IOException e) {
+                    throw new UncheckedIOException(e);
                 }
-            } else {
-                resize(0);
             }
+            return values[index];
+        }
+
+        public String getValue() {
+            return count == 0 ? null : values[0];
         }
 
         /**
          * Set the {@link #size()} and ensure that the {@link #values} array can
          * store at least that many entries.
-         */
-        protected void resize(int newSize) {
+        */
+        private void resize(int newSize) {
+            values = ArrayUtil.grow(values, newSize);
             count = newSize;
-            if (newSize > values.length) {
-                final int oldLength = values.length;
-                values = ArrayUtil.grow(values, count);
-                for (int i = oldLength; i < values.length; ++i) {
-                    values[i] = new BytesRefBuilder();
-                }
+        }
+
+        private void fillValues() throws IOException {
+            assert valuesSet == false;
+            resize(count);
+            for (int i = 0; i < count; i++) {
+                values[i] = in.nextValue().utf8ToString();
+            }
+            valuesSet = true;
+        }
+
+        @Override
+        public void setNextDocId(int docId) throws IOException {
+            if (in.advanceExact(docId)) {
+                count = in.docValueCount();
+                valuesSet = false;
+            } else {
+                resize(0);
+                valuesSet = true;
             }
         }
 
         @Override
         public int size() {
             return count;
         }
-
     }
 
-    public static final class Strings extends BinaryScriptDocValues<String> {
+    public static final class BytesRefs extends ScriptDocValues<BytesRef> {
+        private final SortedBinaryDocValues in;
+        private BytesRef[] values = new BytesRef[0];
+        private int count;
 
-        public Strings(SortedBinaryDocValues in) {
-            super(in);
+        public BytesRefs(SortedBinaryDocValues in) {
+            this.in = in;
         }
 
         @Override
-        public String get(int index) {
-            return values[index].get().utf8ToString();
+        public void setNextDocId(int docId) throws IOException {
+            if (in.advanceExact(docId)) {
+                resize(in.docValueCount());
+                for (int i = 0; i < count; i++) {
+                    /**
+                     * We need to make a copy here because {@link SortedBinaryDocValues} might reuse the returned value
+                     * and the same instance might be used to return values from multiple documents.
+                     **/
+                    values[i] = BytesRef.deepCopyOf(in.nextValue());
+                }
+            } else {
+                resize(0);
+            }
         }
 
-        public String getValue() {
-            return count == 0 ? null : get(0);
+        @Override
+        public BytesRef get(int index) {
+            return values[index];
         }
-    }
-
-    public static final class BytesRefs extends BinaryScriptDocValues<BytesRef> {
 
-        public BytesRefs(SortedBinaryDocValues in) {
-            super(in);
+        public BytesRef getValue() {
+            return count == 0 ? new BytesRef() : values[0];
         }
 
-        @Override
-        public BytesRef get(int index) {
-            /**
-             * We need to make a copy here because {@link BinaryScriptDocValues} might reuse the
-             * returned value and the same instance might be used to
-             * return values from multiple documents.
-             **/
-            return values[index].toBytesRef();
+        /**
+         * Set the {@link #size()} and ensure that the {@link #values} array can
+         * store at least that many entries.
+         */
+        protected void resize(int newSize) {
+            values = ArrayUtil.grow(values, newSize);
+            count = newSize;
         }
 
-        public BytesRef getValue() {
-            return count == 0 ? new BytesRef() : get(0);
+        @Override
+        public int size() {
+            return count;
         }
 
     }