extends Logging

with WithLocalCheckpoints

with WithIntermediateStorageLevel {

private val withVertexColumnList = collection.mutable.ListBuffer.empty[(String, Column, Column)]

private var maxIter: Int = 10

private var checkpointInterval = 2

private var earlyStopping = false

private var stopIfAllNonActiveVertices = false

private var skipMessagesFromNonActiveVertices = false

private var initialActiveVertexExpression = lit(true)

private var updateActiveVertexExpression = lit(true)

private val sendMsgs = collection.mutable.ListBuffer.empty[(Column, Column)]

private var aggMsgsCol: Column = null

// Required columns for source and destination vertices in triplets

// When empty, all columns are selected (default behavior)

private val requiredSrcColumnsList = collection.mutable.ListBuffer.empty[String]

private val requiredDstColumnsList = collection.mutable.ListBuffer.empty[String]

// Required columns for edges

// When empty, only src and dst are selected

private val requiredEdgeColumnsList = collection.mutable.ListBuffer.empty[String]

def setMaxIter(value: Int): this.type = {

maxIter = value

this

}

def setCheckpointInterval(value: Int): this.type = {

checkpointInterval = value

this

}

def setEarlyStopping(value: Boolean): this.type = {

earlyStopping = value

this

}

def setStopIfAllNonActiveVertices(value: Boolean): this.type = {

stopIfAllNonActiveVertices = value

this

}

def setInitialActiveVertexExpression(expression: Column): this.type = {

initialActiveVertexExpression = expression

this

}

def setUpdateActiveVertexExpression(expression: Column): this.type = {

updateActiveVertexExpression = expression

this

}

def setSkipMessagesFromNonActiveVertices(value: Boolean): this.type = {

skipMessagesFromNonActiveVertices = value

this

}

def withVertexColumn(

colName: String,

initialExpr: Column,

updateAfterAggMsgsExpr: Column): this.type = {

// TODO: check if this column exists.

require(

colName != null && colName != ID && colName != Pregel.MSG_COL_NAME,

"additional column name cannot be null and cannot be the same name with ID column or " +

"msg column.")

require(initialExpr != null, "additional column should provide a nonnull initial expression.")

require(

updateAfterAggMsgsExpr != null,

"additional column should provide a nonnull " +

"updateAfterAggMsgs expression.")

withVertexColumnList += Tuple3(colName, initialExpr, updateAfterAggMsgsExpr)

this

}

def sendMsgToSrc(msgExpr: Column): this.type = {

sendMsgs += Tuple2(Pregel.src(ID), msgExpr)

this

}

def sendMsgToDst(msgExpr: Column): this.type = {

sendMsgs += Tuple2(Pregel.dst(ID), msgExpr)

this

}

def requiredSrcColumns(colName: String, colNames: String*): this.type = {

requiredSrcColumnsList.clear()

requiredSrcColumnsList += colName

requiredSrcColumnsList ++= colNames

this

}

def requiredDstColumns(colName: String, colNames: String*): this.type = {

requiredDstColumnsList.clear()

requiredDstColumnsList += colName

requiredDstColumnsList ++= colNames

this

}

def requiredEdgeColumns(colName: String, colNames: String*): this.type = {

requiredEdgeColumnsList.clear()

requiredEdgeColumnsList += colName

requiredEdgeColumnsList ++= colNames

this

}

def aggMsgs(aggExpr: Column): this.type = {

aggMsgsCol = aggExpr

this

}

def run(): DataFrame = {

require(

sendMsgs.length > 0,

"We need to set at least one message expression for pregel running.")

require(aggMsgsCol != null, "We need to set aggMsgs for pregel running.")

require(maxIter >= 1, "The max iteration number should be >= 1.")

require(

checkpointInterval >= 0,

"The checkpoint interval should be >= 0, 0 indicates no checkpoint.")

require(

withVertexColumnList.size > 0,

"There should be at least one additional vertex columns for updating.")

val sendMsgsColList = sendMsgs.toList.map { case (id, msg) =>

struct(id.as(ID), msg.as("msg"))

}

val initVertexCols = withVertexColumnList.toList.map { case (colName, initExpr, _) =>

initExpr.as(colName)

}

val updateVertexCols = withVertexColumnList.toList.map { case (colName, _, updateExpr) =>

updateExpr.as(colName)

}

var lastRoundPersistent: scala.collection.mutable.Queue[DataFrame] =

scala.collection.mutable.Queue[DataFrame]()

val initialAttributes = graph.vertices.columns.map(col).toSeq

var currentVertices = graph.vertices.select(

((initialAttributes :+ initialActiveVertexExpression.alias(

Pregel.ACTIVE_FLAG_COL)) ++ initVertexCols): _*)

// Automatic optimization: detect if destination vertex state is needed by analyzing

// the MESSAGE expressions only (not the target ID expressions, since dst.id is always

// available from the edge). If no message expression references dst.* columns,

// we can skip the second join entirely.

// Additionally, if the only dst field referenced is "id", we can still skip since

// dst.id is available from the edge's dst column.

val messageExpressions = sendMsgs.toList.map { case (_, msgExpr) => msgExpr }

val allDstRefs = messageExpressions.flatMap { expr =>

SparkShims.extractColumnReferences(graph.spark, expr).get(DST)

}

val dstPrefixReferenced = allDstRefs.nonEmpty

val dstFieldsReferenced = allDstRefs.flatten.toSet

// We need the dst join if dst is referenced AND fields other than just "id" are accessed

val needsDstState =

dstPrefixReferenced && (dstFieldsReferenced.isEmpty || dstFieldsReferenced != Set(ID))

if (!needsDstState) {

logDebug(

"Optimization: skipping second join (dst state not required by message expressions)")

}

val edges = (if (requiredEdgeColumnsList.isEmpty) {

graph.edges

.select(col(SRC).alias("edge_src"), col(DST).alias("edge_dst"))

} else {

graph.edges

.select(

col(SRC).alias("edge_src"),

col(DST).alias("edge_dst"),

struct(

requiredEdgeColumnsList.head,

requiredEdgeColumnsList.tail.toSeq: _*).as(EDGE))

}).repartition(col("edge_src")).persist(intermediateStorageLevel)

var iteration = 1

val shouldCheckpoint = checkpointInterval > 0

if (shouldCheckpoint && graph.spark.sparkContext.getCheckpointDir.isEmpty && !useLocalCheckpoints) {

// Spark Connect workaround

graph.spark.conf.getOption("spark.checkpoint.dir") match {

case Some(d) => graph.spark.sparkContext.setCheckpointDir(d)

case None =>

throw new IOException(

"Checkpoint directory is not set. Please set it first using sc.setCheckpointDir()" +

"or by specifying the conf 'spark.checkpoint.dir'.")

}

}

// Columns to include in triplet structs (ID + active flag always included if specified)

val srcCols =

if (requiredSrcColumnsList.isEmpty) Seq(col("*"))

else (Seq(ID, Pregel.ACTIVE_FLAG_COL) ++ requiredSrcColumnsList).distinct.map(col)

val dstCols =

if (requiredDstColumnsList.isEmpty) Seq(col("*"))

else (Seq(ID, Pregel.ACTIVE_FLAG_COL) ++ requiredDstColumnsList).distinct.map(col)

breakable {

while (iteration <= maxIter) {

logInfo(s"start Pregel iteration $iteration / $maxIter")

val currRoundPersistent = scala.collection.mutable.Queue[DataFrame]()

currRoundPersistent.enqueue(currentVertices.persist(intermediateStorageLevel))

// Prune non-active vertices early if skipMessagesFromNonActiveVertices

// is enabled and we don't need the dst state.

val srcVertices =

if (!needsDstState && skipMessagesFromNonActiveVertices)

currentVertices.filter(col(Pregel.ACTIVE_FLAG_COL))

else currentVertices

// Build triplets: start with src vertex state joined with edges

val srcWithEdges = srcVertices

.select(struct(srcCols: _*).as(SRC))

.join(edges, Pregel.src(ID) === col("edge_src"))

// Only perform the second join (adding dst vertex state) if needed

var tripletsDF = if (needsDstState) {

srcWithEdges

.join(

currentVertices.select(struct(dstCols: _*).as(DST)),

col("edge_dst") === Pregel.dst(ID))

.drop(col("edge_src"), col("edge_dst"))

} else {

// Skip second join - dst state not needed by any message expression.

// Create a minimal dst struct with just the id from edge_dst for sendMsgToDst to work.

srcWithEdges

.withColumn(DST, struct(col("edge_dst").as(ID)))

.drop(col("edge_src"), col("edge_dst"))

}

// Only prune here if we didn't prune above.

if (needsDstState && skipMessagesFromNonActiveVertices) {

tripletsDF = tripletsDF.filter(

Pregel.src(Pregel.ACTIVE_FLAG_COL) || Pregel.dst(Pregel.ACTIVE_FLAG_COL))

}

val msgDF: DataFrame = tripletsDF

.select(explode(array(sendMsgsColList: _*)).as("msg"))

.select(col("msg.id"), col("msg.msg").as(Pregel.MSG_COL_NAME))

.filter(Pregel.msg.isNotNull)

if (earlyStopping && msgDF.isEmpty) {

logInfo(

s"there are no more non-null messages; Pregel stops earlier at iteration $iteration")

while (lastRoundPersistent.nonEmpty) {

lastRoundPersistent.dequeue().unpersist()

}

lastRoundPersistent = currRoundPersistent

break()

}

val newAggMsgDF = msgDF

.groupBy(ID)

.agg(aggMsgsCol.as(Pregel.MSG_COL_NAME))

val verticesWithMsg = currentVertices.join(newAggMsgDF, Seq(ID), "left_outer")

currentVertices = verticesWithMsg.select(

((initialAttributes :+ updateActiveVertexExpression.alias(

Pregel.ACTIVE_FLAG_COL)) ++ updateVertexCols): _*)

if (shouldCheckpoint && iteration % checkpointInterval == 0) {

if (useLocalCheckpoints) {

currentVertices = currentVertices.localCheckpoint(eager = false)

} else {

currentVertices = currentVertices.checkpoint(eager = false)

}

} else {

// checkpointing do persistence and we do not need to do it again

currRoundPersistent.enqueue(currentVertices.persist(intermediateStorageLevel))

}

if (stopIfAllNonActiveVertices) {

if (currentVertices.filter(col(Pregel.ACTIVE_FLAG_COL)).isEmpty) {

logInfo(

s"all the verties are non-active; Pregel stops earlier at iteration $iteration")

while (lastRoundPersistent.nonEmpty) {

lastRoundPersistent.dequeue().unpersist()

}

lastRoundPersistent = currRoundPersistent

break()

}

}

if (!earlyStopping && !stopIfAllNonActiveVertices) {

// we need to call materialize

currentVertices.count()

}

while (lastRoundPersistent.nonEmpty) {

lastRoundPersistent.dequeue().unpersist()

}

lastRoundPersistent = currRoundPersistent

iteration += 1

}

}

val res = currentVertices.persist(intermediateStorageLevel)

res.count()

while (lastRoundPersistent.nonEmpty) {

lastRoundPersistent.dequeue().unpersist()

}

edges.unpersist()

System.gc()

res

}

val ACTIVE_FLAG_COL = "_pregel_is_active"

val msg: Column = col(MSG_COL_NAME)

def src(colName: String): Column = col(GraphFrame.SRC + "." + colName)

def dst(colName: String): Column = col(GraphFrame.DST + "." + colName)

def edge(colName: String): Column = col(GraphFrame.EDGE + "." + colName)