Spark GraphX, in depth: the property graph, vertex-cut partitioning, aggregateMessages, Pregel and running it in production

The property graph model

A GraphX graph is a Graph[VD, ED]: a directed multigraph whose vertices carry an attribute of type VD and whose edges carry an attribute of type ED. Vertex ids are 64-bit Long values called VertexId. Parallel edges are allowed, and self-loops are allowed but rarely useful. The graph exposes three views:

• graph.vertices, a VertexRDD[VD] of (VertexId, VD) pairs, hash-partitioned and indexed by id so joins against it are fast.
• graph.edges, an EdgeRDD[ED] of Edge(srcId, dstId, attr), stored in columnar blocks per partition.
• graph.triplets, an RDD of EdgeTriplet that joins each edge with both endpoint attributes, srcAttr and dstAttr. Most graph computation is naturally written over triplets.

Graphs are immutable. Operators such as mapVertices, mapEdges, subgraph and joinVertices return new graphs that reuse unchanged structure and indexes, so changing a vertex attribute does not rebuild the edge partitions. That reuse is what makes iterative algorithms affordable.

Vertex-cut storage and the routing table

There are two ways to split a graph across machines. An edge cut assigns each vertex to one machine, and edges that cross machines must be duplicated or communicated. A vertex cut assigns each edge to one partition, and a vertex whose edges land in several partitions is replicated to each of them. Real-world graphs have power-law degree distributions, where a few vertices have millions of edges. Edge cuts handle them badly, because one machine owns the hub and all its edges. Vertex cuts spread the hub's edges across partitions and pay only for copying its attribute.

GraphX uses vertex cuts. Vertices live in the VertexRDD, hash-partitioned by id. Edges live in the EdgeRDD, partitioned by a PartitionStrategy. Each vertex partition keeps a routing table recording which edge partitions reference each of its vertices. When an operation needs vertex attributes next to edges, as triplets and aggregateMessages do, GraphX uses the routing table to ship each attribute only to the edge partitions that need it. It keeps these shipped copies as a replicated vertex view that is reused until the attributes change.

Partition strategies and the replication arithmetic

The partition strategy decides which edge partition each edge goes to, and therefore how many copies of each vertex exist. There are four built-in strategies:

Worked example. A payments graph has 50 million accounts and 1 billion transfers across 256 edge partitions. With EdgePartition2D, no vertex is replicated to more than about 2 x sqrt(256) = 32 partitions. With a random cut, a merchant receiving transfers from millions of customers is copied into nearly all 256. Suppose the average replication factor ends up at 6. Each iteration that ships a 16-byte id-plus-double attribute then moves about 50M x 6 x 16 bytes, around 4.5 GiB, in addition to the messages coming back. Halving the replication factor halves that traffic. This is why EdgePartition2D is a sound default for skewed graphs and why the partition count matters. Replication grows with partitions, so more partitions is not automatically faster.

Two rules apply. Call partitionBy once, straight after construction, and cache the result, because repartitioning is a full shuffle of the edges. Also check what the algorithm needs. The guide says triangleCount requires edges in canonical orientation (srcId < dstId) and a graph partitioned with partitionBy. The general guidance on partition counts in Spark partitioning, in depth applies to the edge RDD too.

Building a graph from tables

Real graphs start as tables with string or UUID keys, and GraphX needs Long ids. Hashing strings to longs is tempting but can collide, and a collision silently merges two accounts. The safe approach is to assign ids with a dictionary and join both edge endpoints against it:

import org.apache.spark.graphx._
import org.apache.spark.sql.functions._

// Transfers between accounts, from a DataFrame. GraphX needs Long vertex ids.
val tx = spark.table("payments.transfers")                 // src_acct, dst_acct, amount
val ids = tx.select(col("src_acct").as("acct"))
  .union(tx.select(col("dst_acct").as("acct"))).distinct()
  .withColumn("vid", monotonically_increasing_id())       // unique, not contiguous
  .cache()

val edges = tx
  .join(ids.withColumnRenamed("acct", "src_acct").withColumnRenamed("vid", "s"), "src_acct")
  .join(ids.withColumnRenamed("acct", "dst_acct").withColumnRenamed("vid", "d"), "dst_acct")
  .select("s", "d", "amount").rdd
  .map(r => Edge(r.getLong(0), r.getLong(1), r.getDouble(2)))

val vertices = ids.rdd.map(r => (r.getLong(1), r.getString(0)))
val graph = Graph(vertices, edges, defaultVertexAttr = "?")
  .partitionBy(PartitionStrategy.EdgePartition2D)
  .cache()

monotonically_increasing_id gives unique but non-contiguous ids, which is fine because GraphX does not need them dense. Cache the id table, and persist it if you need to map results back later, because recomputing it may assign different ids. The defaultVertexAttr argument fills in any vertex referenced by an edge but missing from the vertex RDD. A non-zero count of such vertices usually means a join bug.

aggregateMessages: the core primitive

Most GraphX algorithms reduce to one step. Each edge sends messages to its endpoints based on the edge and vertex attributes, and each vertex combines its incoming messages. The operator is:

def aggregateMessages[Msg: ClassTag](
    sendMsg: EdgeContext[VD, ED, Msg] => Unit,
    mergeMsg: (Msg, Msg) => Msg,
    tripletFields: TripletFields = TripletFields.All): VertexRDD[Msg]

Inside sendMsg you read ctx.srcAttr, ctx.dstAttr and ctx.attr, and call ctx.sendToSrc or ctx.sendToDst. mergeMsg must be commutative and associative, because GraphX applies it first inside each edge partition, like a map-side combiner, and then again after the shuffle. Vertices that receive no message are absent from the result, so join back with a default.

The tripletFields argument is the main performance lever. It tells GraphX which vertex attributes sendMsg reads, so it can skip shipping the others. The default TripletFields.All ships both sides. For the payments graph, total inbound amount and payer count need only the edge attribute:

// Total inbound amount and number of distinct payers per account, in one pass.
// Only the edge attribute is needed, so ship no vertex attributes at all.
val inbound: VertexRDD[(Double, Int)] = graph.aggregateMessages[(Double, Int)](
  ctx => ctx.sendToDst((ctx.attr, 1)),
  (a, b) => (a._1 + b._1, a._2 + b._2),
  TripletFields.EdgeOnly
)

With EdgeOnly no vertex attributes are shipped at all, removing the 4.5 GiB from the example above. Always pass the narrowest fields your function actually uses. If you pass Src but read dstAttr, you get stale or default values, not an error.

Pregel: iterating to convergence

Pregel is the bulk-synchronous pattern from Google's paper of that name. In each superstep, every vertex that received messages runs a vertex program and sends messages along its edges. The computation stops when no messages are sent or when maxIterations is reached. GraphX's pregel operator is built on aggregateMessages. Its third argument, an EdgeDirection, controls which edges of the vertices that received messages may send in the next superstep. Pass it explicitly rather than relying on a default. Only edges next to recently changed vertices do work, which is why convergence gets cheaper as it proceeds.

// Single-source cheapest path by summed edge attribute (amount) from one account.
val source: VertexId = 42L
val init = graph.mapVertices((id, _) => if (id == source) 0.0 else Double.PositiveInfinity)

val sssp = init.pregel(Double.PositiveInfinity, 20, EdgeDirection.Out)(
  (id, dist, newDist) => math.min(dist, newDist),                 // vertex program
  t => if (t.srcAttr + t.attr < t.dstAttr)                        // send only on improvement
         Iterator((t.dstId, t.srcAttr + t.attr)) else Iterator.empty,
  (a, b) => math.min(a, b)                                        // merge
)

// Built-in algorithms return graphs with the result as the vertex attribute.
val cc  = graph.connectedComponents().vertices                   // (vid, smallest vid in component)
val pr  = graph.pageRank(tol = 0.0001).vertices
val tri = graph.convertToCanonicalEdges()
  .partitionBy(PartitionStrategy.RandomVertexCut).triangleCount().vertices

The shortest-path program shows the discipline Pregel needs. The vertex program must be idempotent, sendMsg should send only when it would change the receiver, and the initial message is delivered to every vertex in superstep zero. A sendMsg that always sends never converges and runs until maxIterations, so always set that cap.

The built-in library covers the common cases. It includes connectedComponents (each vertex is labelled with the smallest id in its component), stronglyConnectedComponents(numIter), pageRank(tol) and staticPageRank(numIter), triangleCount, LabelPropagation.run for community detection, ShortestPaths.run for hop counts to landmark vertices, and SVDPlusPlus. All of them return a graph whose vertex attribute is the result.

Operating GraphX jobs

Lineage and checkpointing. Every superstep adds to the RDD lineage. After a few hundred iterations, planning slows, stack overflows become possible and a lost executor can trigger recomputation back to the start. Set a checkpoint directory and spark.graphx.pregel.checkpointInterval so Pregel checkpoints the graph periodically. For your own loops, call graph.checkpoint() every N iterations and materialize it with an action.

Caching and unpersisting. Iterative algorithms create a new graph every iteration. Cache the new one, force it with an action, then unpersist the previous one. Otherwise executor storage fills with old versions and useful data gets evicted. The storage and execution trade-off is covered in Spark memory management.

Skew. A hub vertex still concentrates work. Its merged messages are small, but the edge partitions holding its edges and the task that updates it can straggle. Use EdgePartition2D, pre-filter edges that carry no signal (for example, exclude a payment processor's settlement account from a fraud-ring analysis), and check the per-task time distribution in the Spark UI.

Shuffle volume. Each superstep performs two shuffles: vertex attributes out to edge partitions, and messages back. The general tuning in Spark shuffle architecture applies, but the bigger wins are a narrower TripletFields and a lower replication factor.

Failure modes

• Silent id collisions from hashing string keys to Long. Use a dictionary-assigned id.
• Never-converging Pregel because sendMsg sends unconditionally or the vertex program oscillates. Always cap iterations and log the active-vertex count.
• Out-of-memory on replicated views from large vertex attributes such as feature arrays or strings shipped every iteration. Keep attributes small during iteration and join rich data back at the end.
• Wrong answers from a mismatched TripletFields, where the function reads a side that was not shipped.
• Driver stack overflow or slow planning after long iteration without checkpointing.
• Triangle counts wrong or failing because edges were not canonical or the graph was not partitioned first.

GraphX, GraphFrames or a graph database

Use GraphX when the job is batch graph analytics in Scala or Java on an existing Spark cluster, when you need custom message-passing algorithms, and when the graph fits comfortably in cluster memory with its replication factor. Use GraphFrames, a separate package built on DataFrames, when your team works in Python or SQL. It offers motif finding and DataFrame integration, and some of its algorithms run on GraphX internally. Use a graph database when the workload is interactive traversal of a few hops from a starting point with low latency, which neither Spark library is designed for. For graph machine learning, compute structural features such as degrees, component ids and PageRank in GraphX and feed them to spark.ml models.

What to do next

1. Load a sample of your edges with GraphLoader.edgeListFile or from DataFrames using dictionary-assigned ids, and count vertices with the default attribute.
2. Partition with EdgePartition2D, cache and inspect edge-partition sizes in the Spark UI.
3. Rewrite one aggregation as aggregateMessages with the narrowest TripletFields and compare shuffle bytes with TripletFields.All.
4. Run connectedComponents and one Pregel job with maxIterations capped and the active-vertex count logged.
5. Set a checkpoint directory and spark.graphx.pregel.checkpointInterval before running anything with long iterations.
6. Decide whether Python users need GraphFrames, and keep GraphX for the Scala-side batch algorithms.