addressed comments

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeArrayData.java

@@ -43,8 +43,8 @@
  * In the `values or offset` region, we store the content of elements. For fields that hold
  * fixed-length primitive types, such as long, double, or int, we store the value directly
  * in the field. For fields with non-primitive or variable-length values, we store a relative
- * offset (w.r.t. the base address of the row) that points to the beginning of the variable-length
- * field, and length (they are combined into a long).
+ * offset (w.r.t. the base address of the array) that points to the beginning of
+ * the variable-length field, and length (they are combined into a long).
  *
  * Instances of `UnsafeArrayData` act as pointers to row data stored in this format.
  */
@@ -301,6 +301,7 @@ public boolean equals(Object other) {
     }
     return false;
   }
+
   public void writeToMemory(Object target, long targetOffset) {
     Platform.copyMemory(baseObject, baseOffset, target, targetOffset, sizeInBytes);
   }
@@ -387,50 +388,52 @@ public double[] toDoubleArray() {
     return values;
   }
 
-  private static UnsafeArrayData fromPrimitiveArray(Object arr, int length, final int elementSize) {
-    final int headerSize = calculateHeaderPortionInBytes(length);
-    if (length > (Integer.MAX_VALUE - headerSize) / elementSize) {
+  private static UnsafeArrayData fromPrimitiveArray(
+       Object arr, int offset, int length, int elementSize) {
+    final long headerSize = calculateHeaderPortionInBytes(length);
+    final long valueRegionSize = (long)elementSize * (long)length;
+    final long allocationSize = (headerSize + valueRegionSize + 7) / 8;
+    if (allocationSize > (long)Integer.MAX_VALUE) {
       throw new UnsupportedOperationException("Cannot convert this array to unsafe format as " +
         "it's too big.");
     }
 
-    final int valueRegionSize = elementSize * length;
-    final byte[] data = new byte[valueRegionSize + headerSize];
+    final long[] data = new long[(int)allocationSize];
 
     Platform.putInt(data, Platform.BYTE_ARRAY_OFFSET, length);
     Platform.copyMemory(arr, Platform.INT_ARRAY_OFFSET, data,
       Platform.BYTE_ARRAY_OFFSET + headerSize, valueRegionSize);
 
     UnsafeArrayData result = new UnsafeArrayData();
-    result.pointTo(data, Platform.BYTE_ARRAY_OFFSET, valueRegionSize + headerSize);
+    result.pointTo(data, Platform.BYTE_ARRAY_OFFSET, (int)allocationSize * 8);
     return result;
   }
 
   public static UnsafeArrayData fromPrimitiveArray(boolean[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 1);
+    return fromPrimitiveArray(arr, Platform.BYTE_ARRAY_OFFSET, arr.length, 1);
   }
 
   public static UnsafeArrayData fromPrimitiveArray(byte[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 1);
+    return fromPrimitiveArray(arr, Platform.BYTE_ARRAY_OFFSET, arr.length, 1);
   }
 
   public static UnsafeArrayData fromPrimitiveArray(short[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 2);
+    return fromPrimitiveArray(arr, Platform.SHORT_ARRAY_OFFSET, arr.length, 2);
   }
 
   public static UnsafeArrayData fromPrimitiveArray(int[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 4);
+    return fromPrimitiveArray(arr, Platform.INT_ARRAY_OFFSET, arr.length, 4);
   }
 
   public static UnsafeArrayData fromPrimitiveArray(long[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 8);
+    return fromPrimitiveArray(arr, Platform.LONG_ARRAY_OFFSET, arr.length, 8);
   }
 
   public static UnsafeArrayData fromPrimitiveArray(float[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 4);
+    return fromPrimitiveArray(arr, Platform.FLOAT_ARRAY_OFFSET, arr.length, 4);
   }
 
   public static UnsafeArrayData fromPrimitiveArray(double[] arr) {
-    return fromPrimitiveArray(arr, arr.length, 8);
+    return fromPrimitiveArray(arr, Platform.DOUBLE_ARRAY_OFFSET, arr.length, 8);
   }
 }

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeArrayWriter.java

@@ -49,6 +49,7 @@ private void assertIndexIsValid(int index) {
   }
 
   public void initialize(BufferHolder holder, int numElements, int fixedElementSize) {
+    // We need 4 bytes to store numElements in header
     this.numElements = numElements;
     this.headerInBytes = calculateHeaderPortionInBytes(numElements);
 
@@ -60,14 +61,14 @@ public void initialize(BufferHolder holder, int numElements, int fixedElementSiz
 
     // Initialize information in header
     Platform.putInt(holder.buffer, startingOffset, numElements);
-    Arrays.fill(holder.buffer, startingOffset + 4 - Platform.BYTE_ARRAY_OFFSET,
-      startingOffset + headerInBytes - Platform.BYTE_ARRAY_OFFSET, (byte)0);
-
+    for (int i = 4; i < headerInBytes; i += 8) {
+      Platform.putLong(holder.buffer, startingOffset + i, 0L);
+    }
     holder.cursor += (headerInBytes + fixedElementSize * numElements);
   }
 
-  private long getElementOffset(int ordinal, int scale) {
-    return startingOffset + headerInBytes + ordinal * scale;
+  private long getElementOffset(int ordinal, int elementSize) {
+    return startingOffset + headerInBytes + ordinal * elementSize;
   }
 
   public void setOffsetAndSize(int ordinal, long currentCursor, long size) {
@@ -77,10 +78,6 @@ public void setOffsetAndSize(int ordinal, long currentCursor, long size) {
     write(ordinal, offsetAndSize);
   }
 
-  public void setNullAt(int ordinal) {
-    throw new UnsupportedOperationException("setNullAt() is not supported");
-  }
-
   private void setNullBit(int ordinal) {
     assertIndexIsValid(ordinal);
     BitSetMethods.set(holder.buffer, startingOffset + 4, ordinal);
@@ -187,7 +184,7 @@ public void write(int ordinal, Decimal input, int precision, int scale) {
         // Write the bytes to the variable length portion.
         Platform.copyMemory(
           bytes, Platform.BYTE_ARRAY_OFFSET, holder.buffer, holder.cursor, bytes.length);
-        write(ordinal, ((long)(holder.cursor - startingOffset) << 32) | ((long) bytes.length));
+        setOffsetAndSize(ordinal, holder.cursor, bytes.length);
         holder.cursor += bytes.length;
       }
     } else {
@@ -204,7 +201,7 @@ public void write(int ordinal, UTF8String input) {
     // Write the bytes to the variable length portion.
     input.writeToMemory(holder.buffer, holder.cursor);
 
-    write(ordinal, ((long)(holder.cursor - startingOffset) << 32) | ((long) numBytes));
+    setOffsetAndSize(ordinal, holder.cursor, numBytes);
 
     // move the cursor forward.
     holder.cursor += numBytes;
@@ -218,7 +215,7 @@ public void write(int ordinal, byte[] input) {
     Platform.copyMemory(
       input, Platform.BYTE_ARRAY_OFFSET, holder.buffer, holder.cursor, input.length);
 
-    write(ordinal, ((long)(holder.cursor - startingOffset) << 32) | ((long) input.length));
+    setOffsetAndSize(ordinal, holder.cursor, input.length);
 
     // move the cursor forward.
     holder.cursor += input.length;
@@ -232,7 +229,7 @@ public void write(int ordinal, CalendarInterval input) {
     Platform.putLong(holder.buffer, holder.cursor, input.months);
     Platform.putLong(holder.buffer, holder.cursor + 8, input.microseconds);
 
-    write(ordinal, ((long)(holder.cursor - startingOffset) << 32) | ((long) 16));
+    setOffsetAndSize(ordinal, holder.cursor, 16);
 
     // move the cursor forward.
     holder.cursor += 16;

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala

@@ -189,10 +189,10 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
 
     val jt = ctx.javaType(et)
 
-    val fixedElementSize = et match {
+    val elementOrOffsetSize = et match {
       case t: DecimalType if t.precision <= Decimal.MAX_LONG_DIGITS => 8
       case _ if ctx.isPrimitiveType(jt) => et.defaultSize
-      case _ => 8
+      case _ => 8  // we need 8 bytes to store offset and length for variable-length types]
     }
 
     val tmpCursor = ctx.freshName("tmpCursor")
@@ -232,7 +232,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
         ${writeUnsafeData(ctx, s"((UnsafeArrayData) $input)", bufferHolder)}
       } else {
         final int $numElements = $input.numElements();
-        $arrayWriter.initialize($bufferHolder, $numElements, $fixedElementSize);
+        $arrayWriter.initialize($bufferHolder, $numElements, $elementOrOffsetSize);
 
         for (int $index = 0; $index < $numElements; $index++) {
           if ($input.isNullAt($index)) {

sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/UnsafeArrayDataBenchmark.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.execution.benchmark
 
 import org.apache.spark.SparkConf
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions.{UnsafeArrayData, UnsafeRow}
 import org.apache.spark.sql.catalyst.expressions.codegen.{BufferHolder, UnsafeArrayWriter}
 import org.apache.spark.unsafe.Platform
@@ -32,11 +33,6 @@ import org.apache.spark.util.Benchmark
  */
 class UnsafeArrayDataBenchmark extends BenchmarkBase {
 
-  new SparkConf()
-    .setMaster("local[1]")
-    .setAppName("microbenchmark")
-    .set("spark.driver.memory", "3g")
-
   def calculateHeaderPortionInBytes(count: Int) : Int = {
     // Use this assignment for SPARK-15962
     // val size = 4 + 4 * count
@@ -47,12 +43,11 @@ class UnsafeArrayDataBenchmark extends BenchmarkBase {
   def readUnsafeArray(iters: Int): Unit = {
     val count = 1024 * 1024 * 16
 
-    val intUnsafeArray = new UnsafeArrayData
     var intResult: Int = 0
-    val intSize = calculateHeaderPortionInBytes(count) + 4 * count
-    val intBuffer = new Array[Byte](intSize)
-    Platform.putInt(intBuffer, Platform.BYTE_ARRAY_OFFSET, count)
-    intUnsafeArray.pointTo(intBuffer, Platform.BYTE_ARRAY_OFFSET, intSize)
+    val intBuffer = new Array[Int](count)
+    val intEncoder = ExpressionEncoder[Array[Int]].resolveAndBind()
+    val intInternalRow = intEncoder.toRow(intBuffer)
+    val intUnsafeArray = intInternalRow.getArray(0)
     val readIntArray = { i: Int =>
       var n = 0
       while (n < iters) {
@@ -68,12 +63,11 @@ class UnsafeArrayDataBenchmark extends BenchmarkBase {
       }
     }
 
-    val doubleUnsafeArray = new UnsafeArrayData
     var doubleResult: Double = 0
-    val doubleSize = calculateHeaderPortionInBytes(count) + 8 * count
-    val doubleBuffer = new Array[Byte](doubleSize)
-    Platform.putInt(doubleBuffer, Platform.BYTE_ARRAY_OFFSET, count)
-    doubleUnsafeArray.pointTo(doubleBuffer, Platform.BYTE_ARRAY_OFFSET, doubleSize)
+    val doubleBuffer = new Array[Double](count)
+    val doubleEncoder = ExpressionEncoder[Array[Double]].resolveAndBind()
+    val doubleInternalRow = doubleEncoder.toRow(doubleBuffer)
+    val doubleUnsafeArray = doubleInternalRow.getArray(0)
     val readDoubleArray = { i: Int =>
       var n = 0
       while (n < iters) {
@@ -94,16 +88,6 @@ class UnsafeArrayDataBenchmark extends BenchmarkBase {
     benchmark.addCase("Double")(readDoubleArray)
     benchmark.run
     /*
-    Without SPARK-15962
-    OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
-    Intel Xeon E3-12xx v2 (Ivy Bridge)
-    Read UnsafeArrayData:                    Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
-    ------------------------------------------------------------------------------------------------
-    Int                                            370 /  471        454.0           2.2       1.0X
-    Double                                         351 /  466        477.5           2.1       1.1X
-    */
-    /*
-    With SPARK-15962
     OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
     Intel Xeon E3-12xx v2 (Ivy Bridge)
     Read UnsafeArrayData:                    Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
@@ -161,16 +145,6 @@ class UnsafeArrayDataBenchmark extends BenchmarkBase {
     benchmark.addCase("Double")(writeDoubleArray)
     benchmark.run
     /*
-    Without SPARK-15962
-    OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
-    Intel Xeon E3-12xx v2 (Ivy Bridge)
-    Write UnsafeArrayData:                   Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
-    ------------------------------------------------------------------------------------------------
-    Int                                            337 /  407        498.0           2.0       1.0X
-    Double                                         458 /  496        366.2           2.7       0.7X
-    */
-    /*
-    With SPARK-15962
     OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
     Intel Xeon E3-12xx v2 (Ivy Bridge)
     Write UnsafeArrayData:                   Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
@@ -216,16 +190,6 @@ class UnsafeArrayDataBenchmark extends BenchmarkBase {
     benchmark.addCase("Double")(readDoubleArray)
     benchmark.run
     /*
-    Without SPARK-15962
-    OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
-    Intel Xeon E3-12xx v2 (Ivy Bridge)
-    Get primitive array from UnsafeArrayData: Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)  Relative
-    ------------------------------------------------------------------------------------------------
-    Int                                            218 /  256        288.1           3.5       1.0X
-    Double                                         318 /  539        198.0           5.1       0.7X
-    */
-    /*
-    With SPARK-15962
     OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
     Intel Xeon E3-12xx v2 (Ivy Bridge)
     Get primitive array from UnsafeArrayData: Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)  Relative
@@ -263,16 +227,6 @@ class UnsafeArrayDataBenchmark extends BenchmarkBase {
     benchmark.addCase("Double")(createDoubleArray)
     benchmark.run
     /*
-    Without SPARK-15962
-    OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
-    Intel Xeon E3-12xx v2 (Ivy Bridge)
-    Create UnsafeArrayData from primitive array: Best/Avg Time(ms)   Rate(M/s)  Per Row(ns) Relative
-    ------------------------------------------------------------------------------------------------
-    Int                                            343 /  437        183.6           5.4       1.0X
-    Double                                         322 /  505        195.6           5.1       1.1X
-    */
-    /*
-    With SPARK-15962
     OpenJDK 64-Bit Server VM 1.8.0_91-b14 on Linux 4.0.4-301.fc22.x86_64
     Intel Xeon E3-12xx v2 (Ivy Bridge)
     Create UnsafeArrayData from primitive array: Best/Avg Time(ms)   Rate(M/s)  Per Row(ns) Relative