SemyonSinchenko · Sep 23, 2025 · Apr 24, 2017 · May 9, 2017 · May 10, 2017 · May 18, 2017
diff --git a/core/src/main/scala/org/graphframes/lib/AggregateMessages.scala b/core/src/main/scala/org/graphframes/lib/AggregateMessages.scala
@@ -21,6 +21,7 @@ import org.apache.spark.sql.Column
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.functions.expr
+import org.apache.spark.sql.functions.struct
 import org.graphframes.GraphFrame
 import org.graphframes.Logging
 import org.graphframes.WithIntermediateStorageLevel
@@ -47,8 +48,8 @@ import org.graphframes.WithIntermediateStorageLevel
  *   - [[AggregateMessages.msg]]: message sent to vertex (for aggregation function)
  *
  * Note: If you use this operation to write an iterative algorithm, you may want to use
- * `org.apache.spark.sql.Dataset.checkpoint` or `org.apache.spark.sql.Dataset.localCheckpoint`` as
- * a workaround for caching issues.
+ * [[AggregateMessages$.getCachedDataFrame getCachedDataFrame()]] as a workaround for caching
+ * issues.
  *
  * @example
  *   We can use this function to compute the in-degree of each vertex
@@ -68,27 +69,29 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame)
   import org.graphframes.GraphFrame.ID
   import org.graphframes.GraphFrame.SRC
 
-  private var msgToSrc: Option[Column] = None
+  private var msgToSrc: Seq[Column] = Vector()
 
   /** Send message to source vertex */
-  def sendToSrc(value: Column): this.type = {
-    msgToSrc = Some(value)
+  def sendToSrc(value: Column, values: Column*): this.type = {
+    msgToSrc = value +: values
     this
   }
 
   /** Send message to source vertex, specifying SQL expression as a String */
-  def sendToSrc(value: String): this.type = sendToSrc(expr(value))
+  def sendToSrc(value: String, values: String*): this.type =
+    sendToSrc(expr(value), values.map(expr): _*)
 
-  private var msgToDst: Option[Column] = None
+  private var msgToDst: Seq[Column] = Vector()
 
   /** Send message to destination vertex */
-  def sendToDst(value: Column): this.type = {
-    msgToDst = Some(value)
+  def sendToDst(value: Column, values: Column*): this.type = {
+    msgToDst = value +: values
     this
   }
 
   /** Send message to destination vertex, specifying SQL expression as a String */
-  def sendToDst(value: String): this.type = sendToDst(expr(value))
+  def sendToDst(value: String, values: String*): this.type =
+    sendToDst(expr(value), values.map(expr): _*)
 
   /**
    * Run the aggregation, returning the resulting DataFrame of aggregated messages. This is a lazy
@@ -97,6 +100,7 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame)
    * This returns a DataFrame with schema:
    *   - column "id": vertex ID
    *   - aggCol: aggregate result
+   *   - aggCols: one column with the result of each additional defined aggregation
    * If you need to join this with the original [[GraphFrame.vertices]], you can run an inner join
    * of the form:
    * {{{
@@ -105,29 +109,35 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame)
    *   aggResult.join(g.vertices, ID)
    * }}}
    */
-  def agg(aggCol: Column): DataFrame = {
+  def agg(aggCol: Column, aggCols: Column*): DataFrame = {
     require(
       msgToSrc.nonEmpty || msgToDst.nonEmpty,
       "To run GraphFrame.aggregateMessages," +
         " messages must be sent to src, dst, or both.  Set using sendToSrc(), sendToDst().")
     val triplets = g.triplets
+
+    def msgColumn(columns: Seq[Column], idColumn: Column): DataFrame = columns match {
+      case Seq(c) => triplets.select(idColumn.as(ID), c.as(AggregateMessages.MSG_COL_NAME))
+      case columns =>
+        triplets.select(idColumn.as(ID), struct(columns: _*).as(AggregateMessages.MSG_COL_NAME))
+    }
+
     val cachedVertices = g.vertices.persist(intermediateStorageLevel)
-    val sentMsgsToSrc = msgToSrc.map { msg =>
-      val msgsToSrc =
-        triplets.select(msg.as(AggregateMessages.MSG_COL_NAME), triplets(SRC)(ID).as(ID))
-      // Inner join: only send messages to vertices with edges
+
+    val sentMsgsToSrc = msgToSrc.headOption.map { _ =>
+      val msgsToSrc = msgColumn(msgToSrc, triplets(SRC)(ID))
       msgsToSrc
         .join(cachedVertices, ID)
         .select(msgsToSrc(AggregateMessages.MSG_COL_NAME), col(ID))
     }
-    val sentMsgsToDst = msgToDst.map { msg =>
-      val msgsToDst =
-        triplets.select(msg.as(AggregateMessages.MSG_COL_NAME), triplets(DST)(ID).as(ID))
+    val sentMsgsToDst = msgToDst.headOption.map { _ =>
+      val msgsToDst = msgColumn(msgToDst, triplets(DST)(ID))
+
       msgsToDst
         .join(cachedVertices, ID)
         .select(msgsToDst(AggregateMessages.MSG_COL_NAME), col(ID))
     }
-    val unionedMsgs = (sentMsgsToSrc, sentMsgsToDst) match {
+    val unionMsgs = (sentMsgsToSrc, sentMsgsToDst) match {
       case (Some(toSrc), Some(toDst)) =>
         toSrc.unionAll(toDst)
       case (Some(toSrc), None) => toSrc
@@ -136,7 +146,9 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame)
         // Should never happen. Specify this case to avoid compilation warnings.
         throw new RuntimeException("AggregateMessages: No messages were specified to be sent.")
     }
-    val cachedResult = unionedMsgs.groupBy(ID).agg(aggCol).persist(intermediateStorageLevel)
+
+    val cachedResult =
+      unionMsgs.groupBy(ID).agg(aggCol, aggCols: _*).persist(intermediateStorageLevel)
     // materialize
     cachedResult.count()
     cachedVertices.unpersist()
@@ -149,7 +161,8 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame)
    *
    * See the overloaded method documentation for more details.
    */
-  def agg(aggCol: String): DataFrame = agg(expr(aggCol))
+  def agg(aggCol: String, aggCols: String*): DataFrame =
+    agg(expr(aggCol), aggCols.map(expr(_)): _*)
 }
 
 object AggregateMessages extends Logging with Serializable {

diff --git a/core/src/test/scala/org/graphframes/lib/AggregateMessagesSuite.scala b/core/src/test/scala/org/graphframes/lib/AggregateMessagesSuite.scala
@@ -17,10 +17,14 @@
 
 package org.graphframes.lib
 
+import org.apache.spark.sql.Row
 import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types._
+import org.graphframes.GraphFrame
 import org.graphframes.GraphFrameTestSparkContext
 import org.graphframes.GraphFramesUnreachableException
 import org.graphframes.SparkFunSuite
+import org.graphframes.TestUtils
 import org.graphframes.examples.Graphs
 
 import scala.collection.mutable
@@ -98,4 +102,74 @@ class AggregateMessagesSuite extends SparkFunSuite with GraphFrameTestSparkConte
       assert(agg2Map(user) === trueAgg(user), s"Failure on user $user")
     }
   }
+
+  test("aggregateMessages with multiple message and aggregation columns") {
+    val AM = AggregateMessages
+    val vertices =
+      sqlContext.createDataFrame(List((1, 30), (2, 40), (3, 50), (4, 60))).toDF("id", "att1")
+    val edges =
+      sqlContext.createDataFrame(List((1, 2, 4), (2, 3, 5), (1, 4, 6))).toDF("src", "dst", "att2")
+    val expectedValues = Map(
+      1 -> Tuple2(100L, 5.0),
+      2 -> Tuple2(80L, 4.5),
+      3 -> Tuple2(40L, 5.0),
+      4 -> Tuple2(30L, 6.0))
+
+    val g = GraphFrame(vertices, edges)
+    // aggregateMessages with column aliases
+    val agg = g.aggregateMessages
+      .sendToDst(AM.src("att1").as("att1"), AM.edge("att2").as("att2"))
+      .sendToSrc(AM.dst("att1").as("att1"), AM.edge("att2").as("att2"))
+      .agg(sum(AM.msg("att1")).as("sum_att1"), avg(AM.msg("att2")).as("avg_att2"))
+
+    // aggregateMessages with columns and no aliases
+    val agg2 = g.aggregateMessages
+      .sendToDst(AM.src("att1"), AM.edge("att2"))
+      .sendToSrc(AM.dst("att1"), AM.edge("att2"))
+      .agg(sum(AM.msg("att1")).as("sum_att1"), avg(AM.msg("att2")).as("avg_att2"))
+
+    // aggregateMessages with column expressions
+    val agg3 = g.aggregateMessages
+      .sendToDst("src['att1'] as att1", "edge['att2'] as att2")
+      .sendToSrc("dst['att1'] as att1", "edge['att2'] as att2")
+      .agg("sum(MSG['att1']) AS sum_att1", "avg(MSG['att2']) AS avg_att2")
+
+    // validate schema
+    assert(agg.schema.size === 3)
+    TestUtils.checkColumnType(agg.schema, "id", IntegerType)
+    TestUtils.checkColumnType(agg.schema, "sum_att1", LongType)
+    TestUtils.checkColumnType(agg.schema, "avg_att2", DoubleType)
+
+    assert(agg.schema === agg2.schema)
+    assert(agg.schema === agg3.schema)
+
+    // validate content
+    val output1 = agg
+      .collect()
+      .map {
+        case Row(id: Int, sumAtt1: Long, avgAtt2: Double) =>
+          id -> Tuple2(sumAtt1, avgAtt2)
+        case _ => throw new GraphFramesUnreachableException()
+      }
+      .toMap
+    val output2 = agg2
+      .collect()
+      .map {
+        case Row(id: Int, sumAtt1: Long, avgAtt2: Double) =>
+          id -> Tuple2(sumAtt1, avgAtt2)
+        case _ => throw new GraphFramesUnreachableException()
+      }
+      .toMap
+    val output3 = agg3
+      .collect()
+      .map {
+        case Row(id: Int, sumAtt1: Long, avgAtt2: Double) =>
+          id -> Tuple2(sumAtt1, avgAtt2)
+        case _ => throw new GraphFramesUnreachableException()
+      }
+      .toMap
+    assert(output1 === expectedValues)
+    assert(output2 === expectedValues)
+    assert(output3 === expectedValues)
+  }
 }