Skip to content

feat: fix memory leaks in GraphX #687

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
SemyonSinchenko merged 2 commits into from
Sep 15, 2025
Merged

feat: fix memory leaks in GraphX #687

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
54fbf17
fix memory leaks in GraphX
SemyonSinchenko Sep 15, 2025
0c44f12
remove two more hidden `cache` calls in GraphX
SemyonSinchenko Sep 15, 2025
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
7 changes: 6 additions & 1 deletion 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 @@ -185,7 +185,12 @@ object ConnectedComponents extends Logging {
private def runGraphX(graph: GraphFrame, maxIter: Int): DataFrame = {
val components =
graphx.lib.ConnectedComponents.run(graph.cachedTopologyGraphX, maxIter)
GraphXConversions.fromGraphX(graph, components, vertexNames = Seq(COMPONENT)).vertices
val res =
GraphXConversions.fromGraphX(graph, components, vertexNames = Seq(COMPONENT)).vertices
res.persist(StorageLevel.MEMORY_AND_DISK_SER)
res.count()
components.unpersist()
res
}

private def run(
Expand Down
7 changes: 6 additions & 1 deletion core/src/main/scala/org/graphframes/lib/LabelPropagation.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.MapType
import org.apache.spark.storage.StorageLevel
import org.graphframes.GraphFrame
import org.graphframes.WithAlgorithmChoice
import org.graphframes.WithCheckpointInterval
Expand Down Expand Up @@ -67,7 +68,11 @@ class LabelPropagation private[graphframes] (private val graph: GraphFrame)
private object LabelPropagation {
private def runInGraphX(graph: GraphFrame, maxIter: Int): DataFrame = {
val gx = graphxlib.LabelPropagation.run(graph.cachedTopologyGraphX, maxIter)
GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(LABEL_ID)).vertices
val res = GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(LABEL_ID)).vertices
res.persist(StorageLevel.MEMORY_AND_DISK_SER)
res.count()
gx.unpersist()
res
}

private def keyWithMaxValue(column: Column): Column = {
Expand Down
7 changes: 6 additions & 1 deletion core/src/main/scala/org/graphframes/lib/ShortestPaths.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import org.apache.spark.sql.functions.transform_values
import org.apache.spark.sql.functions.when
import org.apache.spark.sql.types.IntegerType
import org.apache.spark.sql.types.MapType
import org.apache.spark.storage.StorageLevel
import org.graphframes.GraphFrame
import org.graphframes.GraphFrame.quote
import org.graphframes.GraphFramesUnreachableException
Expand Down Expand Up @@ -109,7 +110,11 @@ private object ShortestPaths extends Logging {
transform_keys(col(DISTANCE_ID), (longId: Column, _) => longIdToLandmarkColumn(longId))
}
val cols = graph.vertices.columns.map(quote).map(col) :+ distanceCol.as(DISTANCE_ID)
g.vertices.select(cols.toSeq: _*)
val res = g.vertices.select(cols.toSeq: _*)
res.persist(StorageLevel.MEMORY_AND_DISK_SER)
res.count()
gx.unpersist()
res
}

private def runInGraphFrames(
Expand Down
7 changes: 6 additions & 1 deletion core/src/main/scala/org/graphframes/lib/StronglyConnectedComponents.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ package org.graphframes.lib

import org.apache.spark.graphframes.graphx.{lib => graphxlib}
import org.apache.spark.sql.DataFrame
import org.apache.spark.storage.StorageLevel
import org.graphframes.GraphFrame
import org.graphframes.WithMaxIter

Expand All @@ -42,7 +43,11 @@ class StronglyConnectedComponents private[graphframes] (private val graph: Graph
private object StronglyConnectedComponents {
private def run(graph: GraphFrame, numIter: Int): DataFrame = {
val gx = graphxlib.StronglyConnectedComponents.run(graph.cachedTopologyGraphX, numIter)
GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(COMPONENT_ID)).vertices
val res = GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(COMPONENT_ID)).vertices
res.persist(StorageLevel.MEMORY_AND_DISK_SER)
res.count()
gx.unpersist()
res
}

private[graphframes] val COMPONENT_ID = "component"
Expand Down
3 changes: 2 additions & 1 deletion graphx/src/main/scala/org/apache/spark/graphframes/graphx/Pregel.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,8 +135,9 @@ object Pregel extends Logging {

// compute the messages
var messages = GraphXUtils.mapReduceTriplets(g, sendMsg, mergeMsg)
// It is absolutely enough to checkpoint only graph itself.
val messageCheckpointer =
new PeriodicRDDCheckpointer[(VertexId, A)](checkpointInterval, graph.vertices.sparkContext)
new PeriodicRDDCheckpointer[(VertexId, A)](-1, graph.vertices.sparkContext)
messageCheckpointer.update(messages.asInstanceOf[RDD[(VertexId, A)]])
var isActiveMessagesNonEmpty = !messages.isEmpty()

Expand Down
9 changes: 1 addition & 8 deletions graphx/src/main/scala/org/apache/spark/graphframes/graphx/impl/GraphImpl.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,8 +200,6 @@ class GraphImpl[VD: ClassTag, ED: ClassTag] protected (
mergeMsg: (A, A) => A,
tripletFields: TripletFields,
activeSetOpt: Option[(VertexRDD[_], EdgeDirection)]): VertexRDD[A] = {

vertices.cache()
// For each vertex, replicate its attribute only to partitions where it is
// in the relevant position in an edge.
replicatedVertexView.upgrade(vertices, tripletFields.useSrc, tripletFields.useDst)
Expand Down Expand Up @@ -267,7 +265,6 @@ class GraphImpl[VD: ClassTag, ED: ClassTag] protected (
// The implicit parameter eq will be populated by the compiler if VD and VD2 are equal, and left
// null if not
if (eq != null) {
vertices.cache()
// updateF preserves type, so we can use incremental replication
val newVerts = vertices.leftJoin(other)(updateF).cache()
val changedVerts = vertices.asInstanceOf[VertexRDD[VD2]].diff(newVerts)
Expand Down Expand Up @@ -337,14 +334,10 @@ object GraphImpl {
def apply[VD: ClassTag, ED: ClassTag](
vertices: VertexRDD[VD],
edges: EdgeRDD[ED]): GraphImpl[VD, ED] = {

vertices.cache()

// Convert the vertex partitions in edges to the correct type
val newEdges = edges
.asInstanceOf[EdgeRDDImpl[ED, _]]
.mapEdgePartitions((_, part) => part.withoutVertexAttributes[VD]())
.cache()

GraphImpl.fromExistingRDDs(vertices, newEdges)
}
Expand All @@ -369,7 +362,7 @@ object GraphImpl {
defaultVertexAttr: VD,
edgeStorageLevel: StorageLevel,
vertexStorageLevel: StorageLevel): GraphImpl[VD, ED] = {
val edgesCached = edges.withTargetStorageLevel(edgeStorageLevel).cache()
val edgesCached = edges.withTargetStorageLevel(edgeStorageLevel)
val vertices =
VertexRDD
.fromEdges(edgesCached, edgesCached.partitions.length, defaultVertexAttr)
Expand Down