Skip to content

feat(pregel): automatically skip second join when dst columns not needed #795

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 13 commits into from
Mar 12, 2026
+690 −7
Merged

feat(pregel): automatically skip second join when dst columns not needed #795

Changes from 1 commit
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
2d3c72d
feat(pregel): automatically skip second join when dst columns not needed
james-willis Feb 6, 2026
65b7264
fix: only analyze message expressions for dst detection, provide dst.…
james-willis Feb 6, 2026
ffbfc03
style: run scalafmt and update SparkShims docs to be implementation-a…
james-willis Feb 6, 2026
ea67cfc
feat: skip dst join when only dst.id is referenced
james-willis Feb 6, 2026
9b90e54
refactor: address reviewer feedback on dst join optimization
james-willis Feb 6, 2026
5bd03f5
test: add comprehensive tests for extractColumnReferences and dst joi…
james-willis Feb 7, 2026
7e34f89
refactor: remove Pregel references from SparkShims comments
Feb 28, 2026
f7ccf17
perf: optimize caching and partitioning when skipping dst join
Feb 28, 2026
d3e5b5a
fix: correct repartition syntax for conditional partitioning
Feb 28, 2026
d92982c
fix: correct repartition syntax for conditional partitioning
Feb 28, 2026
1fc9b3c
style: apply scalafmt formatting to Pregel.scala
Feb 28, 2026
b24f943
address Sem's PR comments
james-willis Mar 11, 2026
1be371c
push triplet filtering when no dst join
james-willis Mar 12, 2026
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
Prev Previous commit
Next Next commit
address Sem's PR comments
  • Loading branch information
@james-willis
james-willis committed Mar 11, 2026
commit b24f943387779d610ecae4c9d0a2e38d0cecef94
11 changes: 2 additions & 9 deletions core/src/main/scala/org/graphframes/lib/Pregel.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -424,8 +424,7 @@ class Pregel(val graph: GraphFrame)

val edges = graph.edges
.select(col(SRC).alias("edge_src"), col(DST).alias("edge_dst"), struct(col("*")).as(EDGE))
.repartition(
(if (needsDstState) Seq(col("edge_src"), col("edge_dst")) else Seq(col("edge_src"))): _*)
.repartition(col("edge_src"))
.persist(intermediateStorageLevel)

var iteration = 1
Expand Down Expand Up @@ -458,16 +457,10 @@ class Pregel(val graph: GraphFrame)
currRoundPersistent.enqueue(currentVertices.persist(intermediateStorageLevel))

// Build triplets: start with src vertex state joined with edges
var srcWithEdges = currentVertices
val srcWithEdges = currentVertices
.select(struct(srcCols: _*).as(SRC))
.join(edges, Pregel.src(ID) === col("edge_src"))

// Optimization: persist srcWithEdges when skipping dst join to avoid recomputation
if (!needsDstState) {
srcWithEdges = srcWithEdges.persist(intermediateStorageLevel)
currRoundPersistent.enqueue(srcWithEdges)
}

// Only perform the second join (adding dst vertex state) if needed
var tripletsDF = if (needsDstState) {
srcWithEdges
Expand Down
Loading