minor enhance

jackylk · Feb 28, 2015 · 946a35b · 946a35b
1 parent e0e64ba
commit 946a35b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/core/src/main/scala/org/apache/spark/util/collection/PrimitiveVector.scala b/core/src/main/scala/org/apache/spark/util/collection/PrimitiveVector.scala
@@ -71,12 +71,16 @@ class PrimitiveVector[@specialized(Long, Int, Double) V: ClassTag](initialSize:
 
   /** Resizes the array, dropping elements if the total length decreases. */
   def resize(newLength: Int): PrimitiveVector[V] = {
-    val newArray = new Array[V](newLength)
-    _array.copyToArray(newArray)
-    _array = newArray
+    _array = copyArrayWithLength(newLength)
     if (newLength < _numElements) {
       _numElements = newLength
     }
     this
   }
+
+  protected def copyArrayWithLength(length: Int): Array[V] = {
+    val copy = new Array[V](length)
+    _array.copyToArray(copy)
+    copy
+  }
 }
diff --git a/core/src/main/scala/org/apache/spark/util/collection/SizeTrackingVector.scala b/core/src/main/scala/org/apache/spark/util/collection/SizeTrackingVector.scala
@@ -41,6 +41,6 @@ private[spark] class SizeTrackingVector[T: ClassTag]
    * Return a trimmed version of the underlying array.
    */
   def toArray: Array[T] = {
-    super.iterator.toArray
+    super.copyArrayWithLength(size)
   }
 }