Code review follow up.

witgo · Apr 9, 2014 · 5a33f1d · 5a33f1d
1 parent e8741a7
commit 5a33f1d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala b/core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala
@@ -63,6 +63,8 @@ class BernoulliSampler[T](lb: Double, ub: Double, complement: Boolean = false)
     }
   }
 
+  def cloneComplement() = new BernoulliSampler[T](lb, ub, !complement)
+
   override def clone = new BernoulliSampler[T](lb, ub, complement)
 }
 

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala b/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala
@@ -187,12 +187,12 @@ object MLUtils {
    */
   def kFold[T : ClassTag](rdd: RDD[T], folds: Int, seed: Int): List[Pair[RDD[T], RDD[T]]] = {
     val foldsF = folds.toFloat
-    1.to(folds).map(fold => ((
-      new PartitionwiseSampledRDD(rdd, new BernoulliSampler[T]((fold-1)/foldsF,fold/foldsF,
-        complement = false), seed),
-      new PartitionwiseSampledRDD(rdd, new BernoulliSampler[T]((fold-1)/foldsF,fold/foldsF,
-        complement = true), seed)
-    ))).toList
+    1.to(folds).map  { fold =>
+      val sampler = new BernoulliSampler[T]((fold-1)/foldsF,fold/foldsF, complement = false)
+      val train = new PartitionwiseSampledRDD(rdd, sampler, seed)
+      val test = new PartitionwiseSampledRDD(rdd, sampler.cloneComplement(), seed)
+      (train, test)
+    }.toList
   }
 
   /**
-Original file line number
+Diff line change
@@ Expand Up @@
         }
       }
+      def cloneComplement() = new BernoulliSampler[T](lb, ub, !complement)
       override def clone = new BernoulliSampler[T](lb, ub, complement)
     }
@@ Expand Down @@