[SPARK-14850][ML] convert primitive array from/to unsafe array directly in VectorUDT/MatrixUDT

mllib/src/main/scala/org/apache/spark/mllib/linalg/Matrices.scala

@@ -27,8 +27,7 @@ import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.apache.spark.annotation.Since
 import org.apache.spark.ml.{linalg => newlinalg}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.GenericMutableRow
-import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.catalyst.expressions.{GenericMutableRow, UnsafeArrayData}
 import org.apache.spark.sql.types._
 
 /**
@@ -194,9 +193,9 @@ private[spark] class MatrixUDT extends UserDefinedType[Matrix] {
         row.setByte(0, 0)
         row.setInt(1, sm.numRows)
         row.setInt(2, sm.numCols)
-        row.update(3, new GenericArrayData(sm.colPtrs.map(_.asInstanceOf[Any])))
-        row.update(4, new GenericArrayData(sm.rowIndices.map(_.asInstanceOf[Any])))
-        row.update(5, new GenericArrayData(sm.values.map(_.asInstanceOf[Any])))
+        row.update(3, UnsafeArrayData.fromPrimitiveArray(sm.colPtrs))
+        row.update(4, UnsafeArrayData.fromPrimitiveArray(sm.rowIndices))
+        row.update(5, UnsafeArrayData.fromPrimitiveArray(sm.values))
         row.setBoolean(6, sm.isTransposed)
 
       case dm: DenseMatrix =>
@@ -205,7 +204,7 @@ private[spark] class MatrixUDT extends UserDefinedType[Matrix] {
         row.setInt(2, dm.numCols)
         row.setNullAt(3)
         row.setNullAt(4)
-        row.update(5, new GenericArrayData(dm.values.map(_.asInstanceOf[Any])))
+        row.update(5, UnsafeArrayData.fromPrimitiveArray(dm.values))
         row.setBoolean(6, dm.isTransposed)
     }
     row
@@ -219,12 +218,21 @@ private[spark] class MatrixUDT extends UserDefinedType[Matrix] {
         val tpe = row.getByte(0)
         val numRows = row.getInt(1)
         val numCols = row.getInt(2)
-        val values = row.getArray(5).toDoubleArray()
+        val values = row.getArray(5) match {
+          case u: UnsafeArrayData => u.toDoubleArrayUnchecked
+          case a => a.toDoubleArray()
+        }
         val isTransposed = row.getBoolean(6)
         tpe match {
           case 0 =>
-            val colPtrs = row.getArray(3).toIntArray()
-            val rowIndices = row.getArray(4).toIntArray()
+            val colPtrs = row.getArray(3) match {
+              case u: UnsafeArrayData => u.toIntArrayUnchecked
+              case a => a.toIntArray()
+            }
+            val rowIndices = row.getArray(4) match {
+              case u: UnsafeArrayData => u.toIntArrayUnchecked
+              case a => a.toIntArray()
+            }
             new SparseMatrix(numRows, numCols, colPtrs, rowIndices, values, isTransposed)
           case 1 =>
             new DenseMatrix(numRows, numCols, values, isTransposed)

mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala

@@ -33,8 +33,7 @@ import org.apache.spark.annotation.{AlphaComponent, Since}
 import org.apache.spark.ml.{linalg => newlinalg}
 import org.apache.spark.mllib.util.NumericParser
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.GenericMutableRow
-import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.catalyst.expressions.{GenericMutableRow, UnsafeArrayData}
 import org.apache.spark.sql.types._
 
 /**
@@ -216,15 +215,15 @@ class VectorUDT extends UserDefinedType[Vector] {
         val row = new GenericMutableRow(4)
         row.setByte(0, 0)
         row.setInt(1, size)
-        row.update(2, new GenericArrayData(indices.map(_.asInstanceOf[Any])))
-        row.update(3, new GenericArrayData(values.map(_.asInstanceOf[Any])))
+        row.update(2, UnsafeArrayData.fromPrimitiveArray(indices))
+        row.update(3, UnsafeArrayData.fromPrimitiveArray(values))
         row
       case DenseVector(values) =>
         val row = new GenericMutableRow(4)
         row.setByte(0, 1)
         row.setNullAt(1)
         row.setNullAt(2)
-        row.update(3, new GenericArrayData(values.map(_.asInstanceOf[Any])))
+        row.update(3, UnsafeArrayData.fromPrimitiveArray(values))
         row
     }
   }
@@ -238,11 +237,20 @@ class VectorUDT extends UserDefinedType[Vector] {
         tpe match {
           case 0 =>
             val size = row.getInt(1)
-            val indices = row.getArray(2).toIntArray()
-            val values = row.getArray(3).toDoubleArray()
+            val indices = row.getArray(2) match {
+              case u: UnsafeArrayData => u.toIntArrayUnchecked
+              case a => a.toIntArray()
+            }
+            val values = row.getArray(3) match {
+              case u: UnsafeArrayData => u.toDoubleArrayUnchecked
+              case a => a.toDoubleArray()
+            }
             new SparseVector(size, indices, values)
           case 1 =>
-            val values = row.getArray(3).toDoubleArray()
+            val values = row.getArray(3) match {
+              case u: UnsafeArrayData => u.toDoubleArrayUnchecked
+              case a => a.toDoubleArray()
+            }
             new DenseVector(values)
         }
     }

mllib/src/test/scala/org/apache/spark/mllib/linalg/UDTSerializationBenchmark.scala

@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.linalg
+
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.util.Benchmark
+
+/**
+ * Serialization benchmark for VectorUDT.
+ */
+object UDTSerializationBenchmark {
+
+  def main(args: Array[String]): Unit = {
+    val iters = 1e2.toInt
+    val numRows = 1e3.toInt
+
+    val encoder = ExpressionEncoder[Vector].defaultBinding
+
+    val vectors = (1 to numRows).map { i =>
+      Vectors.dense(Array.fill(1e5.toInt)(1.0 * i))
+    }.toArray
+    val rows = vectors.map(encoder.toRow)
+
+    val benchmark = new Benchmark("VectorUDT de/serialization", numRows, iters)
+
+    benchmark.addCase("serialize") { _ =>
+      var sum = 0
+      var i = 0
+      while (i < numRows) {
+        sum += encoder.toRow(vectors(i)).numFields
+        i += 1
+      }
+    }
+
+    benchmark.addCase("deserialize") { _ =>
+      var sum = 0
+      var i = 0
+      while (i < numRows) {
+        sum += encoder.fromRow(rows(i)).numActives
+        i += 1
+      }
+    }
+
+    /*
+    Java HotSpot(TM) 64-Bit Server VM 1.8.0_60-b27 on Mac OS X 10.11.4
+    Intel(R) Core(TM) i7-4960HQ CPU @ 2.60GHz
+
+    VectorUDT de/serialization:         Best/Avg Time(ms)    Rate(M/s)   Per Row(ns)   Relative
+    -------------------------------------------------------------------------------------------
+    serialize                                 380 /  392          0.0      379730.0       1.0X
+    deserialize                               138 /  142          0.0      137816.6       2.8X
+    */
+    benchmark.run()
+  }
+}

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

@@ -47,7 +47,7 @@
  * Instances of `UnsafeArrayData` act as pointers to row data stored in this format.
  */
 // todo: there is a lof of duplicated code between UnsafeRow and UnsafeArrayData.
-public class UnsafeArrayData extends ArrayData {
+public final class UnsafeArrayData extends ArrayData {
 
   private Object baseObject;
   private long baseOffset;
@@ -81,7 +81,7 @@ private void assertIndexIsValid(int ordinal) {
   }
 
   public Object[] array() {
-    throw new UnsupportedOperationException("Only supported on GenericArrayData.");
+    throw new UnsupportedOperationException("Not supported on UnsafeArrayData.");
   }
 
   /**
@@ -336,4 +336,78 @@ public UnsafeArrayData copy() {
     arrayCopy.pointTo(arrayDataCopy, Platform.BYTE_ARRAY_OFFSET, sizeInBytes);
     return arrayCopy;
   }
+
+  /**
+   * A faster version of `toIntArray`, which use memory copy instead of iterating all elements.
+   * Note that, this method is dangerous if this array contains null elements.  We don't write
+   * null elements into the data region and memory copy will crash as the data size doesn't match.
+   */
+  public int[] toIntArrayUnchecked() {
+    int[] result = new int[numElements];
+    Platform.copyMemory(baseObject, baseOffset + 4 + 4L * numElements,
+      result, Platform.INT_ARRAY_OFFSET, 4L * numElements);
+    return result;
+  }
+
+  /**
+   * A faster version of `toDoubleArray`, which use memory copy instead of iterating all elements.
+   * Note that, this method is dangerous if this array contains null elements.  We don't write
+   * null elements into the data region and memory copy will crash as the data size doesn't match.
+   */
+  public double[] toDoubleArrayUnchecked() {
+    double[] result = new double[numElements];
+    Platform.copyMemory(baseObject, baseOffset + 4 + 4L * numElements,
+      result, Platform.DOUBLE_ARRAY_OFFSET, 8L * numElements);
+    return result;
+  }
+
+  public static UnsafeArrayData fromPrimitiveArray(int[] arr) {
+    int offsetRegionSize = 4 * arr.length;
+    int valueRegionSize = 4 * arr.length;
+    int totalSize = 4 + offsetRegionSize + valueRegionSize;
+    byte[] data = new byte[totalSize];
+
+    Platform.putInt(data, Platform.BYTE_ARRAY_OFFSET, arr.length);
+
+    int offsetPosition = Platform.BYTE_ARRAY_OFFSET + 4;
+    int valueOffset = 4 + offsetRegionSize;
+    for (int i = 0; i < arr.length; i++) {
+      Platform.putInt(data, offsetPosition, valueOffset);
+      offsetPosition += 4;
+      valueOffset += 4;
+    }
+
+    Platform.copyMemory(arr, Platform.INT_ARRAY_OFFSET, data,
+      Platform.BYTE_ARRAY_OFFSET + 4 + offsetRegionSize, valueRegionSize);
+
+    UnsafeArrayData result = new UnsafeArrayData();
+    result.pointTo(data, Platform.BYTE_ARRAY_OFFSET, totalSize);
+    return result;
+  }
+
+  public static UnsafeArrayData fromPrimitiveArray(double[] arr) {
+    int offsetRegionSize = 4 * arr.length;
+    int valueRegionSize = 8 * arr.length;
+    int totalSize = 4 + offsetRegionSize + valueRegionSize;
+    byte[] data = new byte[totalSize];
+
+    Platform.putInt(data, Platform.BYTE_ARRAY_OFFSET, arr.length);
+
+    int offsetPosition = Platform.BYTE_ARRAY_OFFSET + 4;
+    int valueOffset = 4 + offsetRegionSize;
+    for (int i = 0; i < arr.length; i++) {
+      Platform.putInt(data, offsetPosition, valueOffset);
+      offsetPosition += 4;
+      valueOffset += 8;
+    }
+
+    Platform.copyMemory(arr, Platform.DOUBLE_ARRAY_OFFSET, data,
+      Platform.BYTE_ARRAY_OFFSET + 4 + offsetRegionSize, valueRegionSize);
+
+    UnsafeArrayData result = new UnsafeArrayData();
+    result.pointTo(data, Platform.BYTE_ARRAY_OFFSET, totalSize);
+    return result;
+  }
+
+  // TODO: add more specialized methods.
 }

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

@@ -30,7 +30,7 @@
  * [unsafe key array numBytes] [unsafe key array] [unsafe value array]
  */
 // TODO: Use a more efficient format which doesn't depend on unsafe array.
-public class UnsafeMapData extends MapData {
+public final class UnsafeMapData extends MapData {
 
   private Object baseObject;
   private long baseOffset;

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/UnsafeArraySuite.scala

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.util
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.expressions.UnsafeArrayData
+
+class UnsafeArraySuite extends SparkFunSuite {
+
+  test("from primitive int array") {
+    val array = Array(1, 10, 100)
+    val unsafe = UnsafeArrayData.fromPrimitiveArray(array)
+    assert(unsafe.numElements == 3)
+    assert(unsafe.getSizeInBytes == 4 + 4 * 3 + 4 * 3)
+    assert(unsafe.getInt(0) == 1)
+    assert(unsafe.getInt(1) == 10)
+    assert(unsafe.getInt(2) == 100)
+  }
+
+  test("from primitive double array") {
+    val array = Array(1.1, 2.2, 3.3)
+    val unsafe = UnsafeArrayData.fromPrimitiveArray(array)
+    assert(unsafe.numElements == 3)
+    assert(unsafe.getSizeInBytes == 4 + 4 * 3 + 8 * 3)
+    assert(unsafe.getDouble(0) == 1.1)
+    assert(unsafe.getDouble(1) == 2.2)
+    assert(unsafe.getDouble(2) == 3.3)
+  }
+
+  test("to int array unchecked") {
+    val array = Array(1, 10, 100)
+    val unsafe = UnsafeArrayData.fromPrimitiveArray(array)
+    val array2 = unsafe.toIntArrayUnchecked
+    assert(array.toSeq == array2.toSeq)
+  }
+
+  test("to double array unchecked") {
+    val array = Array(1.1, 2.2, 3.3)
+    val unsafe = UnsafeArrayData.fromPrimitiveArray(array)
+    val array2 = unsafe.toDoubleArrayUnchecked
+    assert(array.toSeq == array2.toSeq)
+  }
+}