[WIP] Add isTransposed to pickle DenseMatrix

nemccarthy · Apr 17, 2015 · 151f3b6 · 151f3b6
1 parent cc0b90a
commit 151f3b6
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
@@ -985,6 +985,7 @@ private[spark] object SerDe extends Serializable {
       val m: DenseMatrix = obj.asInstanceOf[DenseMatrix]
       val bytes = new Array[Byte](8 * m.values.size)
       val order = ByteOrder.nativeOrder()
+      val isTransposed = if (m.isTransposed) 1 else 0
       ByteBuffer.wrap(bytes).order(order).asDoubleBuffer().put(m.values)
 
       out.write(Opcodes.BININT)
@@ -994,19 +995,22 @@ private[spark] object SerDe extends Serializable {
       out.write(Opcodes.BINSTRING)
       out.write(PickleUtils.integer_to_bytes(bytes.length))
       out.write(bytes)
-      out.write(Opcodes.TUPLE3)
+      out.write(Opcodes.BININT)
+      out.write(PickleUtils.integer_to_bytes(isTransposed))
+      out.write(Opcodes.TUPLE)
     }
 
     def construct(args: Array[Object]): Object = {
-      if (args.length != 3) {
-        throw new PickleException("should be 3")
+      if (args.length != 4) {
+        throw new PickleException("should be 4")
       }
       val bytes = getBytes(args(2))
       val n = bytes.length / 8
       val values = new Array[Double](n)
       val order = ByteOrder.nativeOrder()
+      val isTransposed = args(3).asInstanceOf[Boolean]
       ByteBuffer.wrap(bytes).order(order).asDoubleBuffer().get(values)
-      new DenseMatrix(args(0).asInstanceOf[Int], args(1).asInstanceOf[Int], values)
+      new DenseMatrix(args(0).asInstanceOf[Int], args(1).asInstanceOf[Int], values, isTransposed)
     }
   }