[SPARK-16043][SQL] Prepare GenericArrayData implementation specialized for a primitive array

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.expressions.codegen;
 
+import org.apache.spark.sql.catalyst.util.ArrayData;
+import org.apache.spark.sql.catalyst.util.GenericArrayData;
 import org.apache.spark.sql.types.Decimal;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.array.ByteArrayMethods;
@@ -259,4 +261,116 @@ public void write(int ordinal, CalendarInterval input) {
     // move the cursor forward.
     holder.cursor += 16;
   }
+
+  public void writePrimitiveBooleanArray(ArrayData arrayData) {
+    boolean[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).booleanArray()) != null) {
+      int length = input.length;
+      Platform.copyMemory(input, Platform.BOOLEAN_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putBoolean(holder.buffer, startingOffset + headerInBytes + i,
+                arrayData.getBoolean(i));
+      }
+    }
+  }
+
+  public void writePrimitiveByteArray(ArrayData arrayData) {
+    byte[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).byteArray()) != null) {
+      int length = input.length;
+      Platform.copyMemory(input, Platform.BYTE_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putByte(holder.buffer, startingOffset + headerInBytes + i,
+                arrayData.getByte(i));
+      }
+    }
+  }
+
+  public void writePrimitiveShortArray(ArrayData arrayData) {
+    short[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).shortArray()) != null) {
+      int length = input.length * 2;
+      Platform.copyMemory(input, Platform.SHORT_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putShort(holder.buffer, startingOffset + headerInBytes + i * 2,
+                arrayData.getShort(i));
+      }
+    }
+  }
+
+  public void writePrimitiveIntArray(ArrayData arrayData) {
+    int[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).intArray()) != null) {
+      int length = input.length * 4;
+      Platform.copyMemory(input, Platform.INT_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putInt(holder.buffer, startingOffset + headerInBytes + i * 4,
+                arrayData.getInt(i));
+      }
+    }
+  }
+
+  public void writePrimitiveLongArray(ArrayData arrayData) {
+    long[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).longArray()) != null) {
+      int length = input.length * 8;
+      Platform.copyMemory(input, Platform.LONG_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putLong(holder.buffer, startingOffset + headerInBytes + i * 8,
+                arrayData.getLong(i));
+      }
+    }
+  }
+
+  public void writePrimitiveFloatArray(ArrayData arrayData) {
+    float[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).floatArray()) != null) {
+      int length = input.length * 4;
+      Platform.copyMemory(input, Platform.FLOAT_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putFloat(holder.buffer, startingOffset + headerInBytes + i * 4,
+                arrayData.getFloat(i));
+      }
+    }
+  }
+
+  public void writePrimitiveDoubleArray(ArrayData arrayData) {
+    double[] input;
+    if (arrayData instanceof GenericArrayData &&
+            (input = ((GenericArrayData)arrayData).doubleArray()) != null) {
+      int length = input.length * 8;
+      Platform.copyMemory(input, Platform.DOUBLE_ARRAY_OFFSET,
+              holder.buffer, startingOffset + headerInBytes, length);
+    } else {
+      int length = arrayData.numElements();
+      for (int i = 0; i < length; i++) {
+        Platform.putDouble(holder.buffer, startingOffset + headerInBytes + i * 8,
+                arrayData.getDouble(i));
+      }
+    }
+  }
 }

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

@@ -117,12 +117,12 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
               $rowWriter.setOffsetAndSize($index, $tmpCursor, $bufferHolder.cursor - $tmpCursor);
             """
 
-          case a @ ArrayType(et, _) =>
+          case a @ ArrayType(et, cn) =>
             s"""
               // Remember the current cursor so that we can calculate how many bytes are
               // written later.
               final int $tmpCursor = $bufferHolder.cursor;
-              ${writeArrayToBuffer(ctx, input.value, et, bufferHolder)}
+              ${writeArrayToBuffer(ctx, input.value, et, cn, bufferHolder)}
               $rowWriter.setOffsetAndSize($index, $tmpCursor, $bufferHolder.cursor - $tmpCursor);
             """
 
@@ -171,6 +171,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
       ctx: CodegenContext,
       input: String,
       elementType: DataType,
+      containsNull: Boolean,
       bufferHolder: String): String = {
     val arrayWriterClass = classOf[UnsafeArrayWriter].getName
     val arrayWriter = ctx.freshName("arrayWriter")
@@ -202,10 +203,10 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
           $arrayWriter.setOffsetAndSize($index, $tmpCursor, $bufferHolder.cursor - $tmpCursor);
         """
 
-      case a @ ArrayType(et, _) =>
+      case a @ ArrayType(et, cn) =>
         s"""
           final int $tmpCursor = $bufferHolder.cursor;
-          ${writeArrayToBuffer(ctx, element, et, bufferHolder)}
+          ${writeArrayToBuffer(ctx, element, et, cn, bufferHolder)}
           $arrayWriter.setOffsetAndSize($index, $tmpCursor, $bufferHolder.cursor - $tmpCursor);
         """
 
@@ -225,21 +226,39 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     }
 
     val primitiveTypeName = if (ctx.isPrimitiveType(jt)) ctx.primitiveTypeName(et) else ""
+    val storeElements = if (containsNull) {
+      s"""
+       for (int $index = 0; $index < $numElements; $index++) {
+         if ($input.isNullAt($index)) {
+           $arrayWriter.setNull$primitiveTypeName($index);
+         } else {
+           final $jt $element = ${ctx.getValue(input, et, index)};
+           $writeElement
+         }
+       }
+     """
+    } else {
+      if (ctx.isPrimitiveType(et)) {
+        val typeName = ctx.primitiveTypeName(et)
+        s"$arrayWriter.writePrimitive${typeName}Array($input);"
+      } else {
+        s"""
+         for (int $index = 0; $index < $numElements; $index++) {
+           final $jt $element = ${ctx.getValue(input, et, index)};
+           $writeElement
+         }
+       """
+      }
+    }
+
     s"""
       if ($input instanceof UnsafeArrayData) {
         ${writeUnsafeData(ctx, s"((UnsafeArrayData) $input)", bufferHolder)}
       } else {
         final int $numElements = $input.numElements();
         $arrayWriter.initialize($bufferHolder, $numElements, $elementOrOffsetSize);
 
-        for (int $index = 0; $index < $numElements; $index++) {
-          if ($input.isNullAt($index)) {
-            $arrayWriter.setNull$primitiveTypeName($index);
-          } else {
-            final $jt $element = ${ctx.getValue(input, et, index)};
-            $writeElement
-          }
-        }
+        $storeElements
       }
     """
   }
@@ -271,11 +290,11 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
         // Remember the current cursor so that we can write numBytes of key array later.
         final int $tmpCursor = $bufferHolder.cursor;
 
-        ${writeArrayToBuffer(ctx, keys, keyType, bufferHolder)}
+        ${writeArrayToBuffer(ctx, keys, keyType, false, bufferHolder)}
         // Write the numBytes of key array into the first 8 bytes.
         Platform.putLong($bufferHolder.buffer, $tmpCursor - 8, $bufferHolder.cursor - $tmpCursor);
 
-        ${writeArrayToBuffer(ctx, values, valueType, bufferHolder)}
+        ${writeArrayToBuffer(ctx, values, valueType, true, bufferHolder)}
       }
     """
   }