🌐 Graph Processing in PySpark with GraphX

Architecture Diagram

┌─────────────────────────────────────────────────────────────────────────┐
│                    GRAPHX ARCHITECTURE IN SPARK                         │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   ┌──────────────┐     ┌──────────────┐     ┌──────────────┐          │
│   │   Vertex     │────▶│   Edge       │────▶│   Graph      │          │
│   │   RDD        │     │   RDD        │     │   Object     │          │
│   └──────────────┘     └──────────────┘     └──────┬───────┘          │
│                                                     │                   │
│                              ┌──────────────────────┤                   │
│                              │                      │                   │
│                              ▼                      ▼                   │
│                   ┌──────────────────┐   ┌──────────────────┐         │
│                   │  Graph Ops       │   │  Pregel API      │         │
│                   │  ─────────────   │   │  (Iterative)     │         │
│                   │  connectedComp   │   │  ─────────────   │         │
│                   │  triangleCount   │   │  PageRank        │         │
│                   │  shortestPaths   │   │  Label Propagation│         │
│                   └────────┬─────────┘   └────────┬─────────┘         │
│                            │                      │                    │
│                            ▼                      ▼                    │
│                   ┌──────────────────┐   ┌──────────────────┐         │
│                   │  AggregateMessages│  │  Vertex Program  │         │
│                   │  (Map-reduce on  │   │  (Apply function │         │
│                   │   edges)         │   │   to each vertex)│         │
│                   └──────────────────┘   └──────────────────┘         │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────┐
│              PREGEL ITERATIVE COMPUTATION MODEL                         │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   Superstep 0         Superstep 1         Superstep 2                   │
│   ┌─────────┐         ┌─────────┐         ┌─────────┐                 │
│   │  (A,∞)  │────────▶│  (A,0)  │────────▶│  (A,1)  │                 │
│   │  (B,∞)  │         │  (B,1)  │         │  (B,2)  │                 │
│   │  (C,∞)  │         │  (C,2)  │         │  (C,3)  │                 │
│   │  (D,∞)  │         │  (D,3)  │         │  (D,4)  │                 │
│   └─────────┘         └─────────┘         └─────────┘                 │
│        │                   │                   │                        │
│        ▼                   ▼                   ▼                        │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  Each superstep:                                            │     │
│   │  1. Each vertex computes message from incoming edges        │     │
│   │  2. Messages sent along edges                               │     │
│   │  3. Each vertex aggregates received messages                 │     │
│   │  4. Vertex updates its value based on aggregated messages    │     │
│   │  5. Repeat until convergence or max iterations              │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  Convergence Detection:                                     │     │
│   │  ┌─────────────────────────────────────────────────────┐   │     │
│   │  │ if no messages sent OR maxIterations reached:       │   │     │
│   │  │     return final graph                              │   │     │
│   │  └─────────────────────────────────────────────────────┘   │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────┐
│              GRAPH ALGORITHMS COMPARISON                                │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   Algorithm          Complexity      Use Case            Parallelism   │
│   ────────────────────────────────────────────────────────────────────  │
│                                                                         │
│   PageRank           O(V × E × k)    Web ranking         High          │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  V=vertices, E=edges, k=iterations                         │     │
│   │  Each iteration: O(E) for message passing                  │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
│   Connected Comp     O(V + E)        Social networks     Very High     │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  Union-Find with path compression and union by rank        │     │
│   │  Near-linear in practice with good partitioning            │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
│   Shortest Paths     O(V × E)        Routing             Medium        │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  Dijkstra-like relaxation across partitions                │     │
│   │  Requires careful message passing to handle edge cuts      │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
│   Triangle Count     O(E^1.5)        Network analysis    Low-Medium    │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  Counts triangles incident on each vertex                 │     │
│   │  Requires edge orientation for efficient counting          │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
│   Label Propagation   O(V × E × k)   Community detection  High         │
│   ┌─────────────────────────────────────────────────────────────┐     │
│   │  Semi-supervised learning on graph structure               │     │
│   │  Converges when labels stabilize                           │     │
│   └─────────────────────────────────────────────────────────────┘     │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

Detailed Explanation

GraphX is Spark's API for graph-parallel computation. It extends the RDD abstraction by introducing the Graph abstraction, which consists of a Vertex RDD (representing vertices with attributes) and an Edge RDD (representing directed edges with attributes). GraphX leverages Spark's distributed computing model to process graphs that are too large to fit in memory on a single machine, partitioning vertices and edges across the cluster.

The Pregel API is the core abstraction for iterative graph algorithms. Pregel implements the "think like a vertex" paradigm, where each vertex executes the same user-defined function at each superstep. The vertex function receives messages from the previous superstep, updates the vertex's state, and sends messages to neighboring vertices. This model naturally expresses algorithms like PageRank, shortest paths, and label propagation. The computation terminates when no more messages are sent or when a maximum iteration count is reached.

PageRank measures the importance of vertices in a graph by counting the number and quality of links to a page. The algorithm assigns a score to each vertex, where vertices with more incoming high-score edges receive higher scores. In GraphX, PageRank is implemented as a iterative message-passing algorithm where each vertex distributes its rank equally among its outgoing neighbors, then sums the contributions received from its incoming neighbors. The damping factor (typically 0.85) accounts for the probability of following a random link versus jumping to a random page.

Connected components identify groups of vertices that are reachable from each other. In GraphX, this is implemented using a label propagation approach where each vertex starts with its own ID as its component label, then iteratively adopts the minimum label among its neighbors. This converges when all vertices in the same component share the same label. The algorithm is highly parallelizable and achieves near-linear time complexity with good graph partitioning.

Graph partitioning is critical for performance in distributed graph processing. GraphX uses edge partitioning by default, where edges are distributed across partitions to minimize cross-partition communication. The partition strategy can be configured using PartitionStrategy, with options like RandomVertexCut (random edge assignment), EdgePartition2D (2D grid partitioning), and CanonicalRandomVertexCut (consistent partitioning for undirected graphs). The quality of partitioning directly impacts the number of messages that must be shuffled between partitions during computation.

The aggregateMessages API is the low-level building block for graph algorithms. It allows users to define a sendMsg function that sends messages along edges and a mergeMsg function that combines messages arriving at the same vertex. This API provides maximum flexibility for implementing custom algorithms that cannot be expressed through the built-in Pregel API. The sendMsg function receives the source and destination vertex attributes along with the edge, and returns an optional message to send to either vertex.

Mathematical Foundations

Definition: Property Graph

A property graph $G = (V, E, P_V, P_E)$ consists of vertex set $V$, edge set $E \subseteq V \times V$, vertex properties $P_V: V \rightarrow \mathcal{P}$, and edge properties $P_E: E \rightarrow \mathcal{P}$ where $\mathcal{P}$ is the property domain.

PageRank Iteration

PageRank at iteration $k+1$ is computed as:

where $d$ is the damping factor (typically 0.85) and $\text{in}(v)$ is the set of vertices with edges into $v$.

Connected Components Theorem

In an undirected graph $G$, two vertices $u$ and $v$ are in the same connected component if and only if there exists a path $p = (u = v_0, v_1, \ldots, v_k = v)$ where $\forall i: (v_i, v_{i+1}) \in E$. The number of components equals the number of equivalence classes under the reachability relation.

Betweenness Centrality

For vertex $v$, betweenness centrality is:

where $\sigma_{st}$ is the total number of shortest paths from $s$ to $t$, and $\sigma_{st}(v)$ is the number passing through $v$.

Graph Partitioning Load Balance

For $k$ partitions, load imbalance is:

Target: Imbalance $< 0.1$ for efficient distributed computation.

Key Insight

GraphX uses edge-centric partitioning where edges are assigned to partitions by hashing source vertex IDs. This creates communication patterns dominated by vertex messages across partition boundaries, making vertex-cut partitioning critical for performance.

Summary

Graph processing on Spark leverages BSP (Bulk Synchronous Parallel) computation. PageRank converges via power iteration, connected components use label propagation, and graph partitioning quality directly impacts communication overhead in distributed settings.

Key Concepts Table

Code Examples

Creating and Analyzing Graphs

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import *

spark = SparkSession.builder \
    .appName("GraphProcessing") \
    .getOrCreate()

# Create edge DataFrame (social network)
edges_data = [
    (1, 2, "friend"),
    (1, 3, "friend"),
    (2, 3, "colleague"),
    (3, 4, "friend"),
    (4, 5, "colleague"),
    (5, 6, "friend"),
    (6, 1, "friend"),
    (2, 5, "friend"),
    (3, 6, "colleague"),
    (4, 7, "friend"),
    (7, 8, "friend"),
    (8, 9, "colleague"),
    (9, 10, "friend"),
    (10, 7, "friend"),
]

edges_df = spark.createDataFrame(edges_data, ["src_id", "dst_id", "relationship"])

# Create vertex DataFrame
vertices_data = [
    (1, "Alice", 28, "Engineering"),
    (2, "Bob", 35, "Marketing"),
    (3, "Charlie", 42, "Engineering"),
    (4, "Diana", 31, "Sales"),
    (5, "Eve", 38, "Marketing"),
    (6, "Frank", 45, "Engineering"),
    (7, "Grace", 29, "Sales"),
    (8, "Hank", 33, "Engineering"),
    (9, "Ivy", 27, "Marketing"),
    (10, "Jack", 40, "Sales"),
]

vertices_df = spark.createDataFrame(vertices_data, ["id", "name", "age", "department"])

# Convert to GraphX graph (using GraphFrames alternative)
# Since GraphX is Scala-only, we use GraphFrames API
from graphframes import GraphFrame

graph = GraphFrame(vertices_df, edges_df)

# Basic graph statistics
print("Vertices:", graph.vertices.count())
print("Edges:", graph.edges.count())

# Find strongly connected components
scc = graph.connectedComponents()
scc.groupBy("component").count().orderBy("count", ascending=False).show()

# Triangle count
triangle_count = graph.triangleCount()
triangle_count.select("id", "name", "count").orderBy("count", ascending=False).show()

# PageRank
pagerank = graph.pageRank(resetProbability=0.15, maxIter=20)
pagerank.vertices.select("id", "name", "pagerank") \
    .orderBy("pagerank", ascending=False).show()

# Shortest paths from vertex 1
shortest_paths = graph.shortestPaths(landmarks=["1", "5"])
shortest_paths.select("id", "name", "distances").show(truncate=False)

Custom Graph Algorithm with Pregel

# Implementing personalized PageRank using GraphX-style Pregel
# (Using RDD-based approach for true GraphX simulation)

from pyspark import RDD
from typing import Tuple, Dict, List
import math

# Define vertex and edge types
VertexData = Tuple[int, Dict]  # (vertex_id, {attr1: val1, ...})
EdgeData = Tuple[int, int, Dict]  # (src_id, dst_id, {attr1: val1, ...})

def create_graph(vertices: List[Tuple], edges: List[Tuple]) -> Tuple[RDD, RDD]:
    """Create vertex and edge RDDs."""
    vertex_rdd = spark.sparkContext.parallelize(vertices)
    edge_rdd = spark.sparkContext.parallelize(edges)
    return vertex_rdd, edge_rdd

# Pregel implementation for personalized PageRank
def personalized_pagerank(
    vertex_rdd: RDD,
    edge_rdd: RDD,
    source_vertex: int,
    damping: float = 0.85,
    max_iterations: int = 20,
    tolerance: float = 0.001
) -> RDD:
    """
    Compute personalized PageRank from source vertex.
    
    Each vertex maintains:
    - rank: current PageRank score
    - residual: accumulated rank from neighbors
    """
    # Initialize: source vertex gets rank 1.0, others get 0.0
    def init_vertex(vid, attrs):
        rank = 1.0 if vid == source_vertex else 0.0
        return {**attrs, 'rank': rank, 'residual': 0.0}
    
    vertices = vertex_rdd.map(lambda v: (v[0], init_vertex(v[0], v[1])))
    
    # Create adjacency list (outgoing edges)
    adjacency = edge_rdd.map(lambda e: (e[0], e[1])) \
        .groupByKey() \
        .mapValues(list)
    
    # Pregel iterations
    for iteration in range(max_iterations):
        # Send messages: each vertex sends rank/out_degree to neighbors
        messages = vertices.join(adjacency) \
            .flatMap(lambda x: [
                (dst, x[1][0]['rank'] / len(x[1][1]))
                for dst in x[1][1]
            ])
        
        # Aggregate messages at destination vertices
        aggregated = messages.reduceByKey(lambda a, b: a + b)
        
        # Update vertex ranks
        new_vertices = vertices.leftOuterJoin(aggregated) \
            .map(lambda x: (
                x[0],
                {
                    **x[1][0],
                    'residual': x[1][1] if x[1][1] is not None else 0.0
                }
            ))
        
        # Apply rank update formula
        def update_rank(vid, attrs):
            new_rank = (1 - damping) + damping * attrs['residual']
            return {**attrs, 'rank': new_rank, 'residual': 0.0}
        
        updated_vertices = new_vertices.map(lambda v: (v[0], update_rank(v[0], v[1])))
        
        # Check convergence
        rank_diff = updated_vertices.join(vertices) \
            .map(lambda x: abs(x[1][0]['rank'] - x[1][1]['rank'])) \
            .max()
        
        vertices = updated_vertices
        
        if rank_diff < tolerance:
            print(f"Converged after {iteration + 1} iterations")
            break
    
    return vertices

# Run personalized PageRank from vertex 1
vertex_rdd, edge_rdd = create_graph(vertices_data, edges_data)
result = personalized_pagerank(vertex_rdd, edge_rdd, source_vertex=1)

# Display results
result.map(lambda x: (x[0], x[1]['rank'])) \
    .sortBy(lambda x: -x[1]) \
    .collect()

Graph Analytics Pipeline

# Complete graph analytics pipeline
def graph_analytics_pipeline(edges_df, vertices_df, source_vertex):
    """End-to-end graph analytics pipeline."""
    
    graph = GraphFrame(vertices_df, edges_df)
    
    # 1. Basic statistics
    stats = {
        'vertices': graph.vertices.count(),
        'edges': graph.edges.count(),
        'in_degree': graph.inDegrees.orderBy("inDegree", ascending=False).first(),
        'out_degree': graph.outDegrees.orderBy("outDegree", ascending=False).first(),
    }
    
    # 2. Connected components
    components = graph.connectedComponents()
    num_components = components.select("component").distinct().count()
    largest_component = components.groupBy("component") \
        .count() \
        .orderBy("count", ascending=False) \
        .first()
    
    # 3. PageRank
    pagerank = graph.pageRank(resetProbability=0.15, maxIter=20)
    top_pagerank = pagerank.vertices \
        .select("id", "name", "pagerank") \
        .orderBy("pagerank", ascending=False) \
        .first()
    
    # 4. Triangle count
    triangles = graph.triangleCount()
    total_triangles = triangles.agg(sum("count")).first()[0]
    
    # 5. Label propagation (community detection)
    communities = graph.labelPropagation(maxIter=5)
    num_communities = communities.select("label").distinct().count()
    
    return {
        'basic_stats': stats,
        'num_components': num_components,
        'largest_component_size': largest_component['count'],
        'top_pagerank_vertex': top_pagerank,
        'total_triangles': total_triangles,
        'num_communities': num_communities,
    }

results = graph_analytics_pipeline(edges_df, vertices_df, source_vertex=1)
for key, value in results.items():
    print(f"{key}: {value}")

Performance Metrics

Best Practices

1. Use edge partitioning (EdgePartition2D) for graphs with high vertex degree to minimize cross-partition communication
2. Cache vertex and edge RDDs when performing multiple graph operations to avoid recomputation
3. Tune partition count to approximately 2-3x the number of cores in your cluster for optimal parallelism
4. Use persist(StorageLevel.MEMORY_AND_DISK) for iterative algorithms to handle memory pressure gracefully
5. Implement early stopping in iterative algorithms based on convergence tolerance to avoid unnecessary iterations
6. Pre-compute adjacency lists when performing multiple traversals from the same graph structure
7. Use canonical vertex ordering for undirected graphs to avoid double-counting edges
8. Monitor shuffle spill during graph operations—if spill is high, increase executor memory or reduce partition count
9. Use GraphFrames (Scala API) when possible for better performance than Python-based graph processing
10. Implement checkpointing for long-running iterative algorithms to enable fault recovery without full restart
11. Partition by edge direction when algorithms require different traversal patterns for incoming vs outgoing edges
12. Use aggregateMessages instead of sendToSrc/sendToDst for better performance when messages flow in one direction

See also: Snowflake Time Travel (snowflake/02), Kafka CDC (kafka/04), Airflow DAGs (airflow/02)