Github user sethah commented on a diff in the pull request:
https://github.com/apache/spark/pull/8112#discussion_r40339926
--- Diff: core/src/main/scala/org/apache/spark/rdd/PairRDDFunctions.scala
---
@@ -263,6 +263,80 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
}
/**
+ * ::Experimental::
+ * Return random, non-overlapping splits of this RDD sampled by key (via
stratified sampling)
+ * with each split containing exactly math.ceil(numItems * samplingRate)
for each stratum.
+ *
+ * This method differs from [[sampleByKey]] and [[sampleByKeyExact]] in
that it provides random
+ * splits (and their complements) instead of just a subsample of the
data. This requires
+ * segmenting random keys into ranges with upper and lower bounds
instead of segmenting the keys
+ * into a high/low bisection of the entire dataset.
+ *
+ * @param weights array of maps of (key -> samplingRate) pairs for each
split, normed by key
+ * @param exact boolean specifying whether to use exact subsampling
+ * @param seed seed for the random number generator
+ * @return array of tuples containing the subsample and complement RDDs
for each split
+ */
+ @Experimental
+ def randomSplitByKey(
+ weights: Array[Map[K, Double]],
+ exact: Boolean = false,
+ seed: Long = Utils.random.nextLong): Array[(RDD[(K, V)], RDD[(K,
V)])] = self.withScope {
+
+ require(weights.flatMap(_.values).forall(v => v >= 0.0), "Negative
sampling rates.")
+ if (weights.length > 1) {
+ require(weights.map(m => m.keys.toList).sliding(2).forall(t => t(0)
== t(1)),
+ "Inconsistent keys between splits.")
+ }
+
+ // normalize and cumulative sum
+ val baseFold = weights(0).map(x => (x._1, 0.0))
+ val cumWeightsByKey = weights.scanLeft(baseFold){ case (accMap,
iterMap) =>
+ accMap.map { case (k, v) => (k, v + iterMap(k)) }
+ }.drop(1)
+
+ val weightSumsByKey = cumWeightsByKey.last
+ val normedCumWeightsByKey = cumWeightsByKey.dropRight(1).map(_.map {
case (key, threshold) =>
+ (key, threshold / weightSumsByKey(key))
+ })
+
+ // compute exact thresholds for each stratum if required
+ val splits = if (exact) {
--- End diff --
Renamed some of the variables to make this more readable for other
developers.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
