Big arrays sliced from netty buffers (double) (#90745)

martijnvg · web-flow · commit d19603d266e3 · 2022-10-14T11:30:28.000-04:00
Based on #89668 but for doubles. This should allow aggregations down the road to read doubles values directly from netty buffer, rather than copying it from the netty buffer. Relates to #89437
diff --git a/server/src/main/java/org/elasticsearch/common/bytes/AbstractBytesReference.java b/server/src/main/java/org/elasticsearch/common/bytes/AbstractBytesReference.java
@@ -31,6 +31,18 @@ public int getIntLE(int index) {
         return (get(index + 3) & 0xFF) << 24 | (get(index + 2) & 0xFF) << 16 | (get(index + 1) & 0xFF) << 8 | get(index) & 0xFF;
     }
 
+    @Override
+    public long getLongLE(int index) {
+        return (long) (get(index + 7) & 0xFF) << 56 | (long) (get(index + 6) & 0xFF) << 48 | (long) (get(index + 5) & 0xFF) << 40
+            | (long) (get(index + 4) & 0xFF) << 32 | (long) (get(index + 3) & 0xFF) << 24 | (get(index + 2) & 0xFF) << 16 | (get(index + 1)
+                & 0xFF) << 8 | get(index) & 0xFF;
+    }
+
+    @Override
+    public double getDoubleLE(int index) {
+        return Double.longBitsToDouble(getLongLE(index));
+    }
+
     @Override
     public int indexOf(byte marker, int from) {
         final int to = length();
diff --git a/server/src/main/java/org/elasticsearch/common/bytes/BytesArray.java b/server/src/main/java/org/elasticsearch/common/bytes/BytesArray.java
@@ -136,4 +136,9 @@ public void writeTo(OutputStream os) throws IOException {
     public int getIntLE(int index) {
         return ByteUtils.readIntLE(bytes, offset + index);
     }
+
+    @Override
+    public double getDoubleLE(int index) {
+        return ByteUtils.readDoubleLE(bytes, offset + index);
+    }
 }
diff --git a/server/src/main/java/org/elasticsearch/common/bytes/BytesReference.java b/server/src/main/java/org/elasticsearch/common/bytes/BytesReference.java
@@ -129,6 +129,16 @@ static BytesReference fromByteArray(ByteArray byteArray, int length) {
      */
     int getIntLE(int index);
 
+    /**
+     * Returns the long read from the 8 bytes (LE) starting at the given index.
+     */
+    long getLongLE(int index);
+
+    /**
+     * Returns the double read from the 8 bytes (LE) starting at the given index.
+     */
+    double getDoubleLE(int index);
+
     /**
      * Finds the index of the first occurrence of the given marker between within the given bounds.
      * @param marker marker byte to search
diff --git a/server/src/main/java/org/elasticsearch/common/bytes/CompositeBytesReference.java b/server/src/main/java/org/elasticsearch/common/bytes/CompositeBytesReference.java
@@ -237,4 +237,16 @@ public int getIntLE(int index) {
         }
         return super.getIntLE(index);
     }
+
+    @Override
+    public double getDoubleLE(int index) {
+        int i = getOffsetIndex(index);
+        int idx = index - offsets[i];
+        int end = idx + 8;
+        BytesReference wholeDoublesLivesHere = references[i];
+        if (end <= wholeDoublesLivesHere.length()) {
+            return wholeDoublesLivesHere.getDoubleLE(idx);
+        }
+        return super.getDoubleLE(index);
+    }
 }
diff --git a/server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java b/server/src/main/java/org/elasticsearch/common/bytes/ReleasableBytesReference.java
@@ -110,6 +110,18 @@ public int getIntLE(int index) {
         return delegate.getIntLE(index);
     }
 
+    @Override
+    public long getLongLE(int index) {
+        assert hasReferences();
+        return delegate.getLongLE(index);
+    }
+
+    @Override
+    public double getDoubleLE(int index) {
+        assert hasReferences();
+        return delegate.getDoubleLE(index);
+    }
+
     @Override
     public int indexOf(byte marker, int from) {
         assert hasReferences();
diff --git a/server/src/main/java/org/elasticsearch/common/util/BigArrays.java b/server/src/main/java/org/elasticsearch/common/util/BigArrays.java
@@ -312,6 +312,13 @@ public void set(long index, byte[] buf, int offset, int len) {
             assert index >= 0 && index < size();
             System.arraycopy(buf, offset << 3, array, (int) index << 3, len << 3);
         }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            int size = (int) size();
+            out.writeVInt(size * 8);
+            out.write(array, 0, size * Double.BYTES);
+        }
     }
 
     private static class ByteArrayAsFloatArrayWrapper extends AbstractArrayWrapper implements FloatArray {
diff --git a/server/src/main/java/org/elasticsearch/common/util/BigDoubleArray.java b/server/src/main/java/org/elasticsearch/common/util/BigDoubleArray.java
@@ -10,7 +10,9 @@
 
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.common.io.stream.StreamOutput;
 
+import java.io.IOException;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.VarHandle;
 import java.nio.ByteOrder;
@@ -24,6 +26,12 @@
  */
 final class BigDoubleArray extends AbstractBigArray implements DoubleArray {
 
+    static {
+        if (ByteOrder.nativeOrder() != ByteOrder.LITTLE_ENDIAN) {
+            throw new Error("The deserialization assumes this class is written with little-endian numbers.");
+        }
+    }
+
     private static final BigDoubleArray ESTIMATOR = new BigDoubleArray(0, BigArrays.NON_RECYCLING_INSTANCE, false);
 
     static final VarHandle VH_PLATFORM_NATIVE_DOUBLE = MethodHandles.byteArrayViewVarHandle(double[].class, ByteOrder.nativeOrder());
@@ -123,4 +131,21 @@ public static long estimateRamBytes(final long size) {
     public void set(long index, byte[] buf, int offset, int len) {
         set(index, buf, offset, len, pages, 3);
     }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        int size = (int) this.size;
+        out.writeVInt(size * Double.BYTES);
+        int lastPageEnd = size % DOUBLE_PAGE_SIZE;
+        if (lastPageEnd == 0) {
+            for (byte[] page : pages) {
+                out.write(page);
+            }
+            return;
+        }
+        for (int i = 0; i < pages.length - 1; i++) {
+            out.write(pages[i]);
+        }
+        out.write(pages[pages.length - 1], 0, lastPageEnd * Double.BYTES);
+    }
 }
diff --git a/server/src/main/java/org/elasticsearch/common/util/DoubleArray.java b/server/src/main/java/org/elasticsearch/common/util/DoubleArray.java
@@ -8,10 +8,19 @@
 
 package org.elasticsearch.common.util;
 
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.Writeable;
+
+import java.io.IOException;
+
 /**
  * Abstraction of an array of double values.
  */
-public interface DoubleArray extends BigArray {
+public interface DoubleArray extends BigArray, Writeable {
+
+    static DoubleArray readFrom(StreamInput in) throws IOException {
+        return new ReleasableDoubleArray(in);
+    }
 
     /**
      * Get an element given its index.
diff --git a/server/src/main/java/org/elasticsearch/common/util/ReleasableDoubleArray.java b/server/src/main/java/org/elasticsearch/common/util/ReleasableDoubleArray.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.common.util;
+
+import org.apache.lucene.util.RamUsageEstimator;
+import org.elasticsearch.common.bytes.ReleasableBytesReference;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+
+import java.io.IOException;
+
+class ReleasableDoubleArray implements DoubleArray {
+    private static final long SHALLOW_SIZE = RamUsageEstimator.shallowSizeOfInstance(ReleasableDoubleArray.class);
+
+    private final ReleasableBytesReference ref;
+
+    ReleasableDoubleArray(StreamInput in) throws IOException {
+        ref = in.readReleasableBytesReference();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeBytesReference(ref);
+    }
+
+    @Override
+    public long size() {
+        return ref.length() / Long.BYTES;
+    }
+
+    @Override
+    public double get(long index) {
+        if (index > Integer.MAX_VALUE / Long.BYTES) {
+            // We can't serialize messages longer than 2gb anyway
+            throw new ArrayIndexOutOfBoundsException();
+        }
+        return ref.getDoubleLE((int) index * Long.BYTES);
+    }
+
+    @Override
+    public double set(long index, double value) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public double increment(long index, double inc) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void fill(long fromIndex, long toIndex, double value) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void set(long index, byte[] buf, int offset, int len) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        /*
+         * If we return the size of the buffer that we've sliced
+         * we're likely to double count things.
+         */
+        return SHALLOW_SIZE;
+    }
+
+    @Override
+    public void close() {
+        ref.decRef();
+    }
+}
diff --git a/server/src/test/java/org/elasticsearch/common/bytes/BytesArrayTests.java b/server/src/test/java/org/elasticsearch/common/bytes/BytesArrayTests.java
@@ -70,4 +70,34 @@ public void testGetIntLE() {
          * the number of bytes in an int. Get it? I'm not sure I do either....
          */
     }
+
+    public void testGetLongLE() {
+        // first 8 bytes = 888, second 8 bytes = Long.MAX_VALUE
+        // tag::noformat
+        byte[] array = new byte[] {
+            0x78, 0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+            -0x1, -0x1, -0x1, -0x1, -0x1, -0x1, -0x1, 0x7F
+        };
+        // end::noformat
+        BytesReference ref = new BytesArray(array, 0, array.length);
+        assertThat(ref.getLongLE(0), equalTo(888L));
+        assertThat(ref.getLongLE(8), equalTo(Long.MAX_VALUE));
+        Exception e = expectThrows(ArrayIndexOutOfBoundsException.class, () -> ref.getLongLE(9));
+        assertThat(e.getMessage(), equalTo("Index 16 out of bounds for length 16"));
+    }
+
+    public void testGetDoubleLE() {
+        // first 8 bytes = 1.2, second 8 bytes = 1.4
+        // tag::noformat
+        byte[] array = new byte[] {
+            0x33, 0x33, 0x33, 0x33, 0x33, 0x33, -0xD, 0x3F,
+            0x66, 0x66, 0x66, 0x66, 0x66, 0x66, -0xA, 0x3F
+        };
+        // end::noformat
+        BytesReference ref = new BytesArray(array, 0, array.length);
+        assertThat(ref.getDoubleLE(0), equalTo(1.2));
+        assertThat(ref.getDoubleLE(8), equalTo(1.4));
+        Exception e = expectThrows(ArrayIndexOutOfBoundsException.class, () -> ref.getDoubleLE(9));
+        assertThat(e.getMessage(), equalTo("Index 9 out of bounds for length 9"));
+    }
 }
diff --git a/server/src/test/java/org/elasticsearch/common/bytes/CompositeBytesReferenceTests.java b/server/src/test/java/org/elasticsearch/common/bytes/CompositeBytesReferenceTests.java
@@ -160,7 +160,34 @@ public void testGetIntLE() {
         assertThat(comp.getIntLE(2), equalTo(0x02010012));
         assertThat(comp.getIntLE(3), equalTo(0x03020100));
         assertThat(comp.getIntLE(4), equalTo(0x04030201));
-        Exception e = expectThrows(ArrayIndexOutOfBoundsException.class, () -> comp.getIntLE(5));
-        assertThat(e.getMessage(), equalTo("Index 4 out of bounds for length 4"));
+        // The jvm can optimize throwing ArrayIndexOutOfBoundsException by reusing the same exception,
+        // but these reused exceptions have no message or stack trace. This sometimes happens when running this test case.
+        // We can assert the exception message if -XX:-OmitStackTraceInFastThrow is set in gradle test task.
+        expectThrows(ArrayIndexOutOfBoundsException.class, () -> comp.getIntLE(5));
+    }
+
+    public void testGetDoubleLE() {
+        // first double = 1.2, second double = 1.4, third double = 1.6
+        // tag::noformat
+        byte[] data = new byte[] {
+            0x33, 0x33, 0x33, 0x33, 0x33, 0x33, -0xD, 0x3F,
+            0x66, 0x66, 0x66, 0x66, 0x66, 0x66, -0xA, 0x3F,
+            -0x66, -0x67, -0x67, -0x67, -0x67, -0x67, -0x7, 0x3F};
+        // end::noformat
+
+        List<BytesReference> refs = new ArrayList<>();
+        int bytesPerChunk = randomFrom(4, 16);
+        for (int offset = 0; offset < data.length; offset += bytesPerChunk) {
+            int length = Math.min(bytesPerChunk, data.length - offset);
+            refs.add(new BytesArray(data, offset, length));
+        }
+        BytesReference comp = CompositeBytesReference.of(refs.toArray(BytesReference[]::new));
+        assertThat(comp.getDoubleLE(0), equalTo(1.2));
+        assertThat(comp.getDoubleLE(8), equalTo(1.4));
+        assertThat(comp.getDoubleLE(16), equalTo(1.6));
+        // The jvm can optimize throwing ArrayIndexOutOfBoundsException by reusing the same exception,
+        // but these reused exceptions have no message or stack trace. This sometimes happens when running this test case.
+        // We can assert the exception message if -XX:-OmitStackTraceInFastThrow is set in gradle test task.
+        expectThrows(IndexOutOfBoundsException.class, () -> comp.getDoubleLE(17));
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/common/bytes/ReleasableBytesReferenceStreamInputTests.java b/server/src/test/java/org/elasticsearch/common/bytes/ReleasableBytesReferenceStreamInputTests.java
@@ -12,6 +12,7 @@
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.DoubleArray;
 import org.elasticsearch.common.util.IntArray;
 import org.junit.After;
 
@@ -80,4 +81,23 @@ public void testBigIntArrayLivesAfterReleasableIsDecremented() throws IOExceptio
         assertThat(ref.hasReferences(), equalTo(false));
     }
 
+    public void testBigDoubleArrayLivesAfterReleasableIsDecremented() throws IOException {
+        DoubleArray testData = BigArrays.NON_RECYCLING_INSTANCE.newDoubleArray(1, false);
+        testData.set(0, 1);
+
+        BytesStreamOutput out = new BytesStreamOutput();
+        testData.writeTo(out);
+
+        ReleasableBytesReference ref = ReleasableBytesReference.wrap(out.bytes());
+
+        try (DoubleArray in = DoubleArray.readFrom(ref.streamInput())) {
+            ref.decRef();
+            assertThat(ref.hasReferences(), equalTo(true));
+
+            assertThat(in.size(), equalTo(testData.size()));
+            assertThat(in.get(0), equalTo(1.0));
+        }
+        assertThat(ref.hasReferences(), equalTo(false));
+    }
+
 }
diff --git a/server/src/test/java/org/elasticsearch/common/io/stream/AbstractStreamTests.java b/server/src/test/java/org/elasticsearch/common/io/stream/AbstractStreamTests.java
@@ -14,6 +14,7 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.settings.SecureString;
 import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.DoubleArray;
 import org.elasticsearch.common.util.IntArray;
 import org.elasticsearch.common.util.PageCacheRecycler;
 import org.elasticsearch.common.util.set.Sets;
@@ -245,6 +246,31 @@ private void assertBigIntArray(int size) throws IOException {
         }
     }
 
+    public void testSmallBigDoubleArray() throws IOException {
+        assertBigDoubleArray(between(0, PageCacheRecycler.DOUBLE_PAGE_SIZE));
+    }
+
+    public void testLargeBigDoubleArray() throws IOException {
+        assertBigDoubleArray(between(PageCacheRecycler.DOUBLE_PAGE_SIZE, 10000));
+    }
+
+    private void assertBigDoubleArray(int size) throws IOException {
+        DoubleArray testData = BigArrays.NON_RECYCLING_INSTANCE.newDoubleArray(size, false);
+        for (int i = 0; i < size; i++) {
+            testData.set(i, randomDouble());
+        }
+
+        BytesStreamOutput out = new BytesStreamOutput();
+        testData.writeTo(out);
+
+        try (DoubleArray in = DoubleArray.readFrom(getStreamInput(out.bytes()))) {
+            assertThat(in.size(), equalTo(testData.size()));
+            for (int i = 0; i < size; i++) {
+                assertThat(in.get(i), equalTo(testData.get(i)));
+            }
+        }
+    }
+
     public void testCollection() throws IOException {
         class FooBar implements Writeable {
 
diff --git a/test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java b/test/framework/src/main/java/org/elasticsearch/common/util/MockBigArrays.java
@@ -611,6 +611,11 @@ public void set(long index, byte[] buf, int offset, int len) {
         public Collection<Accountable> getChildResources() {
             return Collections.singleton(Accountables.namedAccountable("delegate", in));
         }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            in.writeTo(out);
+        }
     }
 
     private class ObjectArrayWrapper<T> extends AbstractArrayWrapper implements ObjectArray<T> {

Original file line number	Diff line number	Diff line change
`@@ -136,4 +136,9 @@ public void writeTo(OutputStream os) throws IOException {`
`136`	`136`	`public int getIntLE(int index) {`
`137`	`137`	`return ByteUtils.readIntLE(bytes, offset + index);`
`138`	`138`	`}`
	`139`	`+`
	`140`	`+ @Override`
	`141`	`+ public double getDoubleLE(int index) {`
	`142`	`+ return ByteUtils.readDoubleLE(bytes, offset + index);`
	`143`	`+ }`
`139`	`144`	`}`
Original file line number	Diff line number	Diff line change
`@@ -237,4 +237,16 @@ public int getIntLE(int index) {`
`237`	`237`	`}`
`238`	`238`	`return super.getIntLE(index);`
`239`	`239`	`}`
	`240`	`+`
	`241`	`+ @Override`
	`242`	`+ public double getDoubleLE(int index) {`
	`243`	`+ int i = getOffsetIndex(index);`
	`244`	`+ int idx = index - offsets[i];`
	`245`	`+ int end = idx + 8;`
	`246`	`+ BytesReference wholeDoublesLivesHere = references[i];`
	`247`	`+ if (end <= wholeDoublesLivesHere.length()) {`
	`248`	`+ return wholeDoublesLivesHere.getDoubleLE(idx);`
	`249`	`+ }`
	`250`	`+ return super.getDoubleLE(index);`
	`251`	`+ }`
`240`	`252`	`}`
Original file line number	Diff line number	Diff line change
`@@ -312,6 +312,13 @@ public void set(long index, byte[] buf, int offset, int len) {`
`312`	`312`	`assert index >= 0 && index < size();`
`313`	`313`	`System.arraycopy(buf, offset << 3, array, (int) index << 3, len << 3);`
`314`	`314`	`}`
	`315`	`+`
	`316`	`+ @Override`
	`317`	`+ public void writeTo(StreamOutput out) throws IOException {`
	`318`	`+ int size = (int) size();`
	`319`	`+ out.writeVInt(size * 8);`
	`320`	`+ out.write(array, 0, size * Double.BYTES);`
	`321`	`+ }`
`315`	`322`	`}`
`316`	`323`
`317`	`324`	`private static class ByteArrayAsFloatArrayWrapper extends AbstractArrayWrapper implements FloatArray {`