Github user squito commented on a diff in the pull request:
https://github.com/apache/spark/pull/3100#discussion_r20130288
--- Diff:
graphx/src/main/scala/org/apache/spark/graphx/impl/EdgePartition.scala ---
@@ -285,50 +337,126 @@ class EdgePartition[
}
/**
- * Upgrade the given edge iterator into a triplet iterator.
+ * Send messages along edges and aggregate them at the receiving
vertices. Implemented by scanning
+ * all edges sequentially and filtering them with `idPred`.
+ *
+ * @param sendMsg generates messages to neighboring vertices of an edge
+ * @param mergeMsg the combiner applied to messages destined to the same
vertex
+ * @param sendMsgUsesSrcAttr whether or not `mapFunc` uses the edge's
source vertex attribute
+ * @param sendMsgUsesDstAttr whether or not `mapFunc` uses the edge's
destination vertex attribute
+ * @param idPred a predicate to filter edges based on their source and
destination vertex ids
*
- * Be careful not to keep references to the objects from this iterator.
To improve GC performance
- * the same object is re-used in `next()`.
+ * @return iterator aggregated messages keyed by the receiving vertex id
*/
- def upgradeIterator(
- edgeIter: Iterator[Edge[ED]], includeSrc: Boolean = true,
includeDst: Boolean = true)
- : Iterator[EdgeTriplet[VD, ED]] = {
- new ReusingEdgeTripletIterator(edgeIter, this, includeSrc, includeDst)
+ def aggregateMessages[A: ClassTag](
+ sendMsg: EdgeContext[VD, ED, A] => Unit,
+ mergeMsg: (A, A) => A,
+ tripletFields: TripletFields,
+ idPred: (VertexId, VertexId) => Boolean): Iterator[(VertexId, A)] = {
+ val aggregates = new Array[A](vertexAttrs.length)
+ val bitset = new BitSet(vertexAttrs.length)
+
+ var ctx = new AggregatingEdgeContext[VD, ED, A](mergeMsg, aggregates,
bitset)
+ var i = 0
+ while (i < size) {
+ val localSrcId = localSrcIds(i)
+ val srcId = local2global(localSrcId)
+ val localDstId = localDstIds(i)
+ val dstId = local2global(localDstId)
+ if (idPred(srcId, dstId)) {
+ ctx.localSrcId = localSrcId
+ ctx.localDstId = localDstId
+ ctx.srcId = srcId
+ ctx.dstId = dstId
+ ctx.attr = data(i)
+ if (tripletFields.useSrc) { ctx.srcAttr = vertexAttrs(localSrcId) }
+ if (tripletFields.useDst) { ctx.dstAttr = vertexAttrs(localDstId) }
+ sendMsg(ctx)
+ }
+ i += 1
+ }
+
+ bitset.iterator.map { localId => (local2global(localId),
aggregates(localId)) }
}
/**
- * Get an iterator over the edges in this partition whose source vertex
ids match srcIdPred. The
- * iterator is generated using an index scan, so it is efficient at
skipping edges that don't
- * match srcIdPred.
+ * Send messages along edges and aggregate them at the receiving
vertices. Implemented by
+ * filtering the source vertex index with `srcIdPred`, then scanning
edge clusters and filtering
+ * with `dstIdPred`. Both `srcIdPred` and `dstIdPred` must match for an
edge to run.
*
- * Be careful not to keep references to the objects from this iterator.
To improve GC performance
- * the same object is re-used in `next()`.
- */
- def indexIterator(srcIdPred: VertexId => Boolean): Iterator[Edge[ED]] =
- index.iterator.filter(kv =>
srcIdPred(kv._1)).flatMap(Function.tupled(clusterIterator))
-
- /**
- * Get an iterator over the cluster of edges in this partition with
source vertex id `srcId`. The
- * cluster must start at position `index`.
+ * @param sendMsg generates messages to neighboring vertices of an edge
+ * @param mergeMsg the combiner applied to messages destined to the same
vertex
+ * @param srcIdPred a predicate to filter edges based on their source
vertex id
+ * @param dstIdPred a predicate to filter edges based on their
destination vertex id
*
- * Be careful not to keep references to the objects from this iterator.
To improve GC performance
- * the same object is re-used in `next()`.
+ * @return iterator aggregated messages keyed by the receiving vertex id
*/
- private def clusterIterator(srcId: VertexId, index: Int) = new
Iterator[Edge[ED]] {
- private[this] val edge = new Edge[ED]
- private[this] var pos = index
+ def aggregateMessagesWithIndex[A: ClassTag](
+ sendMsg: EdgeContext[VD, ED, A] => Unit,
+ mergeMsg: (A, A) => A,
+ tripletFields: TripletFields,
+ srcIdPred: VertexId => Boolean,
+ dstIdPred: VertexId => Boolean): Iterator[(VertexId, A)] = {
+ val aggregates = new Array[A](vertexAttrs.length)
+ val bitset = new BitSet(vertexAttrs.length)
- override def hasNext: Boolean = {
- pos >= 0 && pos < EdgePartition.this.size && srcIds(pos) == srcId
+ var ctx = new AggregatingEdgeContext[VD, ED, A](mergeMsg, aggregates,
bitset)
+ index.iterator.foreach { cluster =>
+ val clusterSrcId = cluster._1
+ val clusterPos = cluster._2
+ val clusterLocalSrcId = localSrcIds(clusterPos)
+ if (srcIdPred(clusterSrcId)) {
+ var pos = clusterPos
+ ctx.srcId = clusterSrcId
+ ctx.localSrcId = clusterLocalSrcId
+ if (tripletFields.useSrc) { ctx.srcAttr =
vertexAttrs(clusterLocalSrcId) }
+ while (pos < size && localSrcIds(pos) == clusterLocalSrcId) {
+ val localDstId = localDstIds(pos)
+ val dstId = local2global(localDstId)
+ if (dstIdPred(dstId)) {
+ ctx.dstId = dstId
+ ctx.localDstId = localDstId
+ ctx.attr = data(pos)
+ if (tripletFields.useDst) { ctx.dstAttr =
vertexAttrs(localDstId) }
+ sendMsg(ctx)
+ }
+ pos += 1
+ }
+ }
}
- override def next(): Edge[ED] = {
- assert(srcIds(pos) == srcId)
- edge.srcId = srcIds(pos)
- edge.dstId = dstIds(pos)
- edge.attr = data(pos)
- pos += 1
- edge
+ bitset.iterator.map { localId => (local2global(localId),
aggregates(localId)) }
+ }
+}
+
+private class AggregatingEdgeContext[VD, ED, A](
+ mergeMsg: (A, A) => A,
+ aggregates: Array[A],
+ bitset: BitSet)
+ extends EdgeContext[VD, ED, A] {
+
+ var srcId: VertexId = _
--- End diff --
@rxin just curious -- do you know whether the JIT will take care of for
you? I had assumed it would, but I've been wrong plenty of times before ...
---
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
