update blockify strategy

Author: zhengruifeng

mllib/src/main/scala/org/apache/spark/ml/feature/Instance.scala

@@ -142,55 +142,32 @@ private[spark] object InstanceBlock {
     new Iterator[InstanceBlock] {
       private var numCols = -1L
       private val buff = mutable.ArrayBuilder.make[Instance]
-      private var buffCnt = 0L
-      private var buffNnz = 0L
-      private var buffUnitWeight = true
-      private var block = Option.empty[InstanceBlock]
 
-      private def flush(): Unit = {
-        block = Some(InstanceBlock.fromInstances(buff.result()))
-        buff.clear()
-        buffCnt = 0L
-        buffNnz = 0L
-        buffUnitWeight = true
-      }
+      override def hasNext: Boolean = iterator.hasNext
 
-      private def blockify(): Unit = {
-        block = None
+      override def next(): InstanceBlock = {
+        buff.clear()
+        var buffCnt = 0L
+        var buffNnz = 0L
+        var buffUnitWeight = true
+        var blockMemUsage = 0L
 
-        while (block.isEmpty && iterator.hasNext) {
+        while (iterator.hasNext && blockMemUsage < maxMemUsage) {
           val instance = iterator.next()
           if (numCols < 0L) numCols = instance.features.size
           require(numCols == instance.features.size)
           val nnz = instance.features.numNonzeros
 
-          // Check if enough memory remains to add this instance to the block.
-          if (getBlockMemUsage(numCols, buffCnt + 1L, buffNnz + nnz,
-            buffUnitWeight && (instance.weight == 1)) > maxMemUsage) {
-            // Check if this instance is too large
-            require(buffCnt > 0, s"instance $instance exceeds memory limit $maxMemUsage, " +
-              s"please increase block size")
-            flush()
-          }
-
           buff += instance
           buffCnt += 1L
           buffNnz += nnz
           buffUnitWeight &&= (instance.weight == 1)
+          blockMemUsage = getBlockMemUsage(numCols, buffCnt, buffNnz, buffUnitWeight)
         }
 
-        if (block.isEmpty && buffCnt > 0) flush()
-      }
-
-      override def hasNext: Boolean = {
-        block.nonEmpty || { blockify(); block.nonEmpty }
-      }
-
-      override def next(): InstanceBlock = {
-        if (block.isEmpty) blockify()
-        val ret = block.get
-        blockify()
-        ret
+        // the block mem usage may slightly exceed threshold, not a big issue.
+        // and this ensure even if one row exceed block limit, each block has one row
+        InstanceBlock.fromInstances(buff.result())
       }
     }
   }

mllib/src/test/scala/org/apache/spark/ml/feature/InstanceSuite.scala

@@ -97,13 +97,13 @@ class InstanceSuite extends SparkFunSuite{
       assert(vec.toArray === instances(i).features.toArray)
     }
 
+    // instances larger than maxMemUsage
     val bigInstance = Instance(-1.0, 2.0, Vectors.dense(Array.fill(10000)(1.0)))
-    val inputIter1 = Iterator.apply(bigInstance)
-    val inputIter2 = Iterator.apply(instance1, instance2, bigInstance)
-    Seq(inputIter1, inputIter2).foreach { inputIter =>
-      intercept[IllegalArgumentException] {
-        InstanceBlock.blokifyWithMaxMemUsage(inputIter, 1024).toArray
-      }
+    InstanceBlock.blokifyWithMaxMemUsage(Iterator.fill(10)(bigInstance), 64).size
+
+    // different numFeatures
+    intercept[IllegalArgumentException] {
+      InstanceBlock.blokifyWithMaxMemUsage(Iterator.apply(instance1, bigInstance), 64).size
     }
   }
 }