Skip to content

revert ConnectedComponents back to 0.9.3 version to fix perf regression #772

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
james-willis merged 1 commit into from
Jan 14, 2026
Merged

revert ConnectedComponents back to 0.9.3 version to fix perf regression #772

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
166 changes: 79 additions & 87 deletions core/src/main/scala/org/graphframes/lib/ConnectedComponents.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

package org.graphframes.lib

import org.apache.hadoop.fs.Path
import org.apache.spark.graphframes.graphx
import org.apache.spark.sql.Column
import org.apache.spark.sql.DataFrame
Expand Down Expand Up @@ -93,6 +94,7 @@ object ConnectedComponents extends Logging {
private val ORIG_ID = "orig_id"
private val MIN_NBR = "min_nbr"
private val CNT = "cnt"
private val CHECKPOINT_NAME_PREFIX = "connected-components"

/**
* Returns the symmetric directed graph of the graph specified by input edges.
Expand All @@ -119,8 +121,11 @@ object ConnectedComponents extends Logging {
* `dst`.
*/
private def prepare(graph: GraphFrame): GraphFrame = {
// TODO: This assignment job might fail if the graph is skewed.
val vertices = graph.indexedVertices
.select(col(LONG_ID).as(ID), col(ATTR))
// TODO: confirm the contract for a graph and decide whether we need distinct here
// .distinct()
val edges = graph.indexedEdges
.select(col(LONG_SRC).as(SRC), col(LONG_DST).as(DST))
val orderedEdges = edges
Expand All @@ -136,36 +141,20 @@ object ConnectedComponents extends Logging {
* - `min_nbr`, the min vertex ID among itself and its neighbors
* - `cnt`, the total number of neighbors
*/
private def minNbrs(
ee: DataFrame,
computeCount: Boolean,
includeSelf: Boolean,
doSymmetrize: Boolean): DataFrame = {
val ee2 = if (doSymmetrize) {
symmetrize(ee)
} else {
ee
}
val res = if (computeCount) {
ee2
.groupBy(SRC)
.agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT))
} else {
ee2
.groupBy(SRC)
.agg(min(col(DST)).as(MIN_NBR))
}
if (includeSelf) {
res.withColumn(MIN_NBR, minValue(col(SRC), col(MIN_NBR)))
} else { res }
private def minNbrs(ee: DataFrame): DataFrame = {
symmetrize(ee)
.groupBy(SRC)
.agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT))
.withColumn(MIN_NBR, minValue(col(SRC), col(MIN_NBR)))
}

private def minValue(x: Column, y: Column): Column = {
when(x < y, x).otherwise(y)
}

private def maxValue(x: Column, y: Column): Column =
private def maxValue(x: Column, y: Column): Column = {
when(x > y, x).otherwise(y)
}

/**
* Performs a possibly skewed join between edges and current component assignments. The skew
Expand Down Expand Up @@ -196,12 +185,7 @@ object ConnectedComponents extends Logging {
private def runGraphX(graph: GraphFrame, maxIter: Int): DataFrame = {
val components =
graphx.lib.ConnectedComponents.run(graph.cachedTopologyGraphX, maxIter)
val res =
GraphXConversions.fromGraphX(graph, components, vertexNames = Seq(COMPONENT)).vertices
res.persist(StorageLevel.MEMORY_AND_DISK_SER)
res.count()
components.unpersist()
res
GraphXConversions.fromGraphX(graph, components, vertexNames = Seq(COMPONENT)).vertices
}

private def run(
Expand All @@ -218,30 +202,40 @@ object ConnectedComponents extends Logging {
}

val spark = graph.spark
val sc = spark.sparkContext
// Store original AQE setting
val originalAQE = spark.conf.get("spark.sql.adaptive.enabled")

try {
val shouldDoSkewedJoin = broadcastThreshold != -1

if (shouldDoSkewedJoin) {
spark.conf.set("spark.sql.adaptive.enabled", "false")
}
spark.conf.set("spark.sql.adaptive.enabled", "false")

val runId = UUID.randomUUID().toString.takeRight(8)
val logPrefix = s"[CC $runId]"
logInfo(s"$logPrefix Start connected components with run ID $runId.")

val shouldCheckpoint = checkpointInterval > 0
if (shouldCheckpoint && !useLocalCheckpoints && spark.sparkContext.getCheckpointDir.isEmpty) {
// Spark-Connect workaround
spark.sparkContext.setCheckpointDir(spark.conf.getOption("spark.checkpoint.dir") match {
case Some(d) => 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'.")
})
val checkpointDir: Option[String] = if (useLocalCheckpoints) { None }
else if (shouldCheckpoint) {
val dir = sc.getCheckpointDir
.map { d =>
new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString
}
.getOrElse {
// Spark-Connect workaround
spark.conf.getOption("spark.checkpoint.dir") match {
case Some(d) => new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString
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'.")
}
}
logInfo(s"$logPrefix Using $dir for checkpointing with interval $checkpointInterval.")
Some(dir)
} else {
logInfo(
s"$logPrefix Checkpointing is disabled because checkpointInterval=$checkpointInterval.")
None
}

logInfo(s"$logPrefix Preparing the graph for connected component computation ...")
Expand All @@ -254,24 +248,29 @@ object ConnectedComponents extends Logging {
var iteration = 1

def _calcMinNbrSum(minNbrsDF: DataFrame): BigDecimal = {
if (shouldDoSkewedJoin) {
minNbrsDF
.select(sum(col(MIN_NBR).cast(DecimalType(38, 0))), count("*"))
.first()
.getDecimal(0)
} else {
minNbrsDF.select(sum(col(MIN_NBR).cast(DecimalType(38, 0)))).first().getDecimal(0)
// Taking the sum in DecimalType to preserve precision.
// We use 20 digits for long values and Spark SQL will add 10 digits for the sum.
// It should be able to handle 200 billion edges without overflow.
val (minNbrSum, cnt) = minNbrsDF
.select(sum(col(MIN_NBR).cast(DecimalType(20, 0))), count("*"))
.rdd
.map { r =>
(r.getAs[BigDecimal](0), r.getLong(1))
}
.first()
if (cnt != 0L && minNbrSum == null) {
throw new ArithmeticException(s"""
|The total sum of edge src IDs is used to determine convergence during iterations.
|However, the total sum at iteration $iteration exceeded 30 digits (1e30),
Comment thread
james-willis marked this conversation as resolved.
|which should happen only if the graph contains more than 200 billion edges.
|If not, please file a bug report at https://github.com/graphframes/graphframes/issues.
""".stripMargin)
}
minNbrSum
}
// compute min neighbors (including self-min)
var minNbrs1: DataFrame =
minNbrs(
ee,
computeCount = shouldDoSkewedJoin,
includeSelf = true,
doSymmetrize = true
) // src >= min_nbr
.persist(intermediateStorageLevel)
var minNbrs1: DataFrame = minNbrs(ee) // src >= min_nbr
.persist(intermediateStorageLevel)

var prevSum: BigDecimal = _calcMinNbrSum(minNbrs1)

Expand All @@ -280,37 +279,23 @@ object ConnectedComponents extends Logging {
var currRoundPersistedDFs = Seq[DataFrame]()
// large-star step
// connect all strictly larger neighbors to the min neighbor (including self)
ee = {
if (shouldDoSkewedJoin) {
skewedJoin(ee, minNbrs1, broadcastThreshold, logPrefix)

} else {
ee.join(minNbrs1, SRC)
}
}.select(col(DST).as(SRC), col(MIN_NBR).as(DST)) // src > dst
ee = skewedJoin(ee, minNbrs1, broadcastThreshold, logPrefix)
.select(col(DST).as(SRC), col(MIN_NBR).as(DST)) // src > dst
.distinct()
.persist(intermediateStorageLevel)

currRoundPersistedDFs = currRoundPersistedDFs :+ ee

// small-star step
// compute min neighbors (excluding self-min)
val minNbrs2 = minNbrs(
ee,
computeCount = shouldDoSkewedJoin,
includeSelf = false,
doSymmetrize = false)
val minNbrs2 = ee
.groupBy(col(SRC))
.agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT)) // src > min_nbr
.persist(intermediateStorageLevel)
currRoundPersistedDFs = currRoundPersistedDFs :+ minNbrs2

// connect all smaller neighbors to the min neighbor
ee = {
if (shouldDoSkewedJoin) {
skewedJoin(ee, minNbrs2, broadcastThreshold, logPrefix)
} else {
ee.join(minNbrs2, SRC)
}
}.select(col(MIN_NBR).as(SRC), col(DST)) // src <= dst
ee = skewedJoin(ee, minNbrs2, broadcastThreshold, logPrefix)
.select(col(MIN_NBR).as(SRC), col(DST)) // src <= dst
.filter(col(SRC) =!= col(DST)) // src < dst
// connect self to the min neighbor
ee = ee
Expand All @@ -322,19 +307,26 @@ object ConnectedComponents extends Logging {
if (useLocalCheckpoints) {
ee = ee.localCheckpoint(eager = true)
} else {
ee = ee.checkpoint(eager = true)
// TODO: remove this after DataFrame.checkpoint is implemented
val out = s"${checkpointDir.get}/$iteration"
ee.write.parquet(out)
// may hit S3 eventually consistent issue
ee = spark.read.parquet(out)

// remove previous checkpoint
if (iteration > checkpointInterval) {
val path = new Path(s"${checkpointDir.get}/${iteration - checkpointInterval}")
path.getFileSystem(sc.hadoopConfiguration).delete(path, true)
}

System.gc() // hint Spark to clean shuffle directories
}
}

ee = ee.persist(intermediateStorageLevel)
ee.persist(intermediateStorageLevel)
currRoundPersistedDFs = currRoundPersistedDFs :+ ee

minNbrs1 = minNbrs(
ee,
computeCount = shouldDoSkewedJoin,
includeSelf = true,
doSymmetrize = true
) // src >= min_nbr
minNbrs1 = minNbrs(ee) // src >= min_nbr
.persist(intermediateStorageLevel)
currRoundPersistedDFs = currRoundPersistedDFs :+ minNbrs1

Expand Down