Github user mengxr commented on a diff in the pull request:
https://github.com/apache/spark/pull/493#discussion_r12081031
--- Diff:
mllib/src/main/scala/org/apache/spark/mllib/recommendation/ALS.scala ---
@@ -708,6 +709,123 @@ object ALS {
trainImplicit(ratings, rank, iterations, 0.01, -1, 1.0)
}
+ @DeveloperApi
+ case class IterationCost(inboundBytes: Double, computation: Double,
outboundBytes: Double)
+
+ /**
+ * :: DeveloperApi ::
+ * Given an RDD of ratings, a rank, and two partitioners, compute rough
estimates of the
+ * computation time and communication cost of one iteration of ALS.
Returns a pair of pairs of
+ * Maps. The first pair of maps represents computation time in
unspecified units. The second
+ * pair of maps represents communication cost in uncompressed bytes.
The first element of each
+ * pair is the cost attributable to user partitioning, while the second
is the cost attributable
+ * to product partitioning.
+ *
+ * @param ratings RDD of Rating objects
+ * @param rank number of features to use
+ * @param userPartitioner partitioner for partitioning users
+ * @param productPartitioner partitioner for partitioning products
+ */
+ @DeveloperApi
+ def estimateCost(ratings: RDD[Rating], rank: Int, userPartitioner:
Partitioner,
+ productPartitioner: Partitioner):
+ (Map[Int, IterationCost], Map[Int, IterationCost]) = {
+ // user partition -> set of products
+ val utalk = ratings.mapPartitions(x => {
+ val ht = new mutable.HashSet[(Int, Int)]()
+ while (x.hasNext) {
+ val rat = x.next()
+ val u = userPartitioner.getPartition(rat.user)
+ val p = rat.product
+ ht += ((u, p))
+ }
+ ht.iterator
+ }
+ )
+
+ utalk.persist()
+
+ // user partition -> number of products
+ val userInbound =
+ utalk.groupByKey.map(x => (x._1, 8.0 * rank *
x._2.toList.distinct.length)).collectAsMap()
+
+ // product -> number of user partitions, summed over each partition.
+ val productOutbound = (utalk.distinct.map(x =>
(productPartitioner.getPartition(x._2), x._1))
+ .groupByKey.mapValues(x => 8.0 * rank *
x.toList.length).collectAsMap())
+
+ utalk.unpersist()
+
+ // product partition -> set of users
+ val ptalk = ratings.mapPartitions(x => {
+ val ht = new mutable.HashSet[(Int, Int)]()
+ while (x.hasNext) {
+ val rat = x.next()
+ val u = rat.user
+ val p = productPartitioner.getPartition(rat.product)
+ ht += ((p, u))
+ }
+ ht.iterator
--- End diff --
Ditto. Try
~~~
val ht = iter.map(rat => (userPartitioner.getPartition(rat.product),
rat.user)).toSet
~~~
---
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.
---
