use long array

Author: cloud-fan

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/ColumnVector.java

@@ -520,11 +520,11 @@ private void throwUnsupportedException(int requiredCapacity, Throwable cause) {
 
   /**
    * After writing array elements to the child column vector, call this method to set the offset and
-   * length of the written array.
+   * size of the written array.
    */
-  public void putArrayOffsetAndLength(int rowId, int offset, int length) {
-    putInt(2 * rowId, offset);
-    putInt(2 * rowId + 1, length);
+  public void putArrayOffsetAndSize(int rowId, int offset, int size) {
+    long offsetAndSize = (offset << 32) | size;
+    putLong(rowId, offsetAndSize);
   }
 
   /**
@@ -548,8 +548,9 @@ public ColumnarBatch.Row getStruct(int rowId, int size) {
    * Returns the array at rowid.
    */
   public final Array getArray(int rowId) {
-    resultArray.offset = getInt(2 * rowId);
-    resultArray.length = getInt(2 * rowId + 1);
+    long offsetAndSize = getLong(rowId);
+    resultArray.offset = (int) (offsetAndSize >> 32);
+    resultArray.length = (int) offsetAndSize;
     return resultArray;
   }
 
@@ -563,7 +564,7 @@ public final Array getArray(int rowId) {
    */
   public int putByteArray(int rowId, byte[] value, int offset, int length) {
     int result = arrayData().appendBytes(length, value, offset);
-    putArrayOffsetAndLength(rowId, result, length);
+    putArrayOffsetAndSize(rowId, result, length);
     return result;
   }
 
@@ -829,13 +830,13 @@ public final int appendDoubles(int length, double[] src, int offset) {
   public final int appendByteArray(byte[] value, int offset, int length) {
     int copiedOffset = arrayData().appendBytes(length, value, offset);
     reserve(elementsAppended + 1);
-    putArrayOffsetAndLength(elementsAppended, copiedOffset, length);
+    putArrayOffsetAndSize(elementsAppended, copiedOffset, length);
     return elementsAppended++;
   }
 
   public final int appendArray(int length) {
     reserve(elementsAppended + 1);
-    putArrayOffsetAndLength(elementsAppended, arrayData().elementsAppended, length);
+    putArrayOffsetAndSize(elementsAppended, arrayData().elementsAppended, length);
     return elementsAppended++;
   }
 

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/OffHeapColumnVector.java

@@ -401,7 +401,7 @@ public void loadBytes(ColumnVector.Array array) {
   protected void reserveInternal(int newCapacity) {
     int oldCapacity = (this.data == 0L) ? 0 : capacity;
     if (this.resultArray != null) {
-      // need 2 ints as offset and length for each array.
+      // need a long as offset and length for each array.
       this.data = Platform.reallocateMemory(data, oldCapacity * 8, newCapacity * 8);
     } else if (type instanceof ByteType || type instanceof BooleanType) {
       this.data = Platform.reallocateMemory(data, oldCapacity, newCapacity);

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/OnHeapColumnVector.java

@@ -42,9 +42,9 @@ public final class OnHeapColumnVector extends ColumnVector {
   // Array for each type. Only 1 is populated for any type.
   private byte[] byteData;
   private short[] shortData;
+  private int[] intData;
   // This is not only used to store data for int column vector, but also can store offsets and
   // lengths for array column vector.
-  private int[] intData;
   private long[] longData;
   private float[] floatData;
   private double[] doubleData;
@@ -374,11 +374,11 @@ public void loadBytes(ColumnVector.Array array) {
   @Override
   protected void reserveInternal(int newCapacity) {
     if (this.resultArray != null || DecimalType.isByteArrayDecimalType(type)) {
-      // need 2 ints as offset and length for each array.
-      if (intData == null || intData.length < newCapacity * 2) {
-        int[] newData = new int[newCapacity * 2];
-        if (intData != null) System.arraycopy(intData, 0, newData, 0, intData.length);
-        intData = newData;
+      // need 1 long as offset and length for each array.
+      if (longData == null || longData.length < newCapacity) {
+        long[] newData = new long[newCapacity];
+        if (longData != null) System.arraycopy(longData, 0, newData, 0, capacity);
+        longData = newData;
       }
     } else if (type instanceof BooleanType) {
       if (byteData == null || byteData.length < newCapacity) {

sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala

@@ -631,7 +631,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
       assert(column.arrayData().elementsAppended == 17)
 
       // Put the same "ll" at offset. This should not allocate more memory in the column.
-      column.putArrayOffsetAndLength(idx, offset, 2)
+      column.putArrayOffsetAndSize(idx, offset, 2)
       reference += "ll"
       idx += 1
       assert(column.arrayData().elementsAppended == 17)
@@ -667,10 +667,10 @@ class ColumnarBatchSuite extends SparkFunSuite {
       }
 
       // Populate it with arrays [0], [1, 2], [], [3, 4, 5]
-      column.putArrayOffsetAndLength(0, 0, 1)
-      column.putArrayOffsetAndLength(1, 1, 2)
-      column.putArrayOffsetAndLength(2, 3, 0)
-      column.putArrayOffsetAndLength(3, 3, 3)
+      column.putArrayOffsetAndSize(0, 0, 1)
+      column.putArrayOffsetAndSize(1, 1, 2)
+      column.putArrayOffsetAndSize(2, 3, 0)
+      column.putArrayOffsetAndSize(3, 3, 3)
 
       val a1 = ColumnVectorUtils.toPrimitiveJavaArray(column.getArray(0)).asInstanceOf[Array[Int]]
       val a2 = ColumnVectorUtils.toPrimitiveJavaArray(column.getArray(1)).asInstanceOf[Array[Int]]
@@ -703,7 +703,7 @@ class ColumnarBatchSuite extends SparkFunSuite {
       data.reserve(array.length)
       assert(data.capacity == array.length * 2)
       data.putInts(0, array.length, array, 0)
-      column.putArrayOffsetAndLength(0, 0, array.length)
+      column.putArrayOffsetAndSize(0, 0, array.length)
       assert(ColumnVectorUtils.toPrimitiveJavaArray(column.getArray(0)).asInstanceOf[Array[Int]]
         === array)
     }}