Connected components at scale in PySpark | by Schaun Wheeler | Towards Data Science

A while ago, I had a network of nodes for which I needed to calculate connected components. That’s not a particularly difficult thing to do. The Python networkx library has a nice implementation that makes it particularly easy, but even if you wanted to roll your own function, it’s a straightforward breadth-first-search. (Khan Academy gives a nice little overview of how that works if you want a gentle introduction). Conceptually, it just involves a set of queues. You pick a node, find all the nodes connected to that node, find all the nodes connected to those nodes, and so forth. Each time you find a node’s connections, you move that node into a “done” list. When you can’t find any connections that aren’t in the “done” list, then you’ve found your connected components. Then you repeat the process for all the rest of the nodes in the graph.
While a breadth first search is conceptually simple, it’s a beast to implement at scale. That’s because “big data” systems like Spark differentiate transformations from actions. A transformation is a plan — a set of intended actions. You can build out a long list of transformations and the call an action — say, the action of collecting the final results — at the end. That’s the cost of working with large datasets: you can store way more information than you can in your individual computer, but accessing that information is much more computationally intensive than it would be if you were just pulling information from your local machine’s memory. That’s why calculating connected components at scale is complicated: breadth first search determines future transformations based on the results of previous transformations, which means you have to call actions all the way through. That can be a real resource burden. At the time I needed to do this, I was dealing with a graph with tens of millions of nodes and hundreds of millions of edges.
Spark’s GraphX library has a connected components function, but at the time I was looking for a way to do this, my entire workflow was is Python and GraphX is only implemented in Scala. There is a GraphFrames project that provides a Python wrapper for GraphX and in some cases an alternative implementation, but as far as I can tell that project isn’t updated all too regularly. I was concerned about my project being reliant on a big complicated dependency that wasn’t consistently supported. Besides, I didn’t need everything in that big, complicated dependency. I only needed a way to calculate connected components.
That set me off on an internet search that gave me some ideas that were, frankly, a little beyond my capability to translate into working code (my fault, not theirs — for example: here, here, and here.) But looking at those implementations gave me at least one really useful perspective: for a breadth-first search, you don’t need queues; you just need to know whether the relationship between two nodes has already been explored. A queue is a simple way to do that, but it’s not the only way.
So here’s what you can do: assign each of your nodes a numeric hash. Assuming an undirected graph (where, if node a is connected to node b, then node b is also connected to node a), you can create an adjacency list (two columns, each full of node hashes, where each row indicates an edge), then group by the hashes of the left-hand column, and take the minimum of the hashes in the right-hand column. Join those back in to your adjacency list and replace the left-column hashes where a new, lower one exists. Then just keep repeating that until nothing changes anymore.
Below, you can see a proof-of-concept using PySpark. This requires you to checkpoint your dataframe by saving it to file. Normally, I find it preferable to persist (cache) the dataframe in memory, because then I don’t have to remember to clean up after myself, but if you cache too many times, Spark starts to spend so much time tracing out the entire lineage of transformations that it never gets around to calling the action. In my use cases, I found that checkpointing every 3–4 iterations, and persisting in all the others, led to acceptable performance.
import pyspark.sql.functions as fimport subprocessdef pyspark_connected_components(    ss, adj, a, b, label, checkpoint_dir,     checkpoint_every=2, max_n=None):    """    This function takes an adjacency list stored in a Spark     data frame and calculates connected components. This     implementation only deals with use cases: assuming an     undirected graph (if a is connected to b then b is    connected to a).    Args:    ss (pyspark.sql.session.SparkSession): an active SparkSession    adj (pyspark.sql.DataFrame): A data frame with at least two        columns, where each entry is a node of a graph and each row        represents an edge connecting two nodes.    a (str): the column name indicating one of the node pairs in the        adjacency list. The column must be numeric.    b (str): the column name indicating the other of the node pairs        in the adjacency list. The column must be numeric.    label (str): a label for the output column of nodes    checkpoint_dir (str): a location on HDFS in which to save         checkpoints    checkpoint every (int): how many iterations to wait before        checkpointing the dataframe    max_n (int): the maximum number of iterations to run    Returns:    A PySpark DataFrame with one column (named by the label         argument) listing all nodes, and another columns         ("component") listing component labels.    """    # set up a few strings for use in the code further down    b2 = b + 'join'    filepattern = 'pyspark_cc_iter{}.parquet'    # cache the adjacency matrix because it will be used many times    adj_cache = adj.persist()    # for each id in the `a` column, take the minimum of the     # ids in the `b` column    mins = (        adj_cache        .groupby(a)        .agg(f.min(f.col(b)).alias('min1'))        .withColumn('change', f.lit(0))        .persist()    )    if max_n is not None:        # ensure a global minimum id that is less than         # any ids in the dataset        minimum_id = mins.select(            f.min(f.col('min1')).alias('n')        ).collect()[0]['n'] - 1000000    # transformation used repeatedly to replace old labels with new    fill_in = f.coalesce(f.col('min2'), f.col('min1'))    # join conditions for comparing nodes    criteria = [        f.col(b) == f.col(b2), f.col('min1') > f.col('min3')    ]    # initial conditions for the loop    not_done = mins.count()    old_done = not_done + 1    new_done = not_done    niter = 0    used_filepaths = []    filepath = None    # repeat until done    while not_done > 0:        niter += 1                if max_n is not None:            min_a = mins.filter(f.col('min1') != minimum_id)        else:            min_a = mins            # get new minimum ids            newmins = (                adj_cache                .join(min_a, a, 'inner')                .join(                    mins.select(                        f.col(a).alias(b2),                         f.col('min1').alias('min3')                    ),                     criteria,                     'inner'                )                .groupby(f.col('min1').alias('min1'))                .agg(f.min(f.col('min3')).alias('min2'))            )        # reconcile new minimum ids with the         # minimum ids from the previous iteration        mins = (            mins            .join(newmins, 'min1', 'left')            .select(                a,                fill_in.alias('min1'),                (f.col('min1') != fill_in).alias('change')            )        )        # if there is a max_n, assign the global minimum id to         # any components with more than max_n nodes        if max_n is not None:            mins = (                mins                .withColumn(                    'n',                    f.count(f.col(a)).over(                        w.partitionBy(                            f.col('min1')).orderBy(f.lit(None)                        )                    )                )                .withColumn(                    'min1',                     f.when(f.col('n') >= max_n, f.lit(minimum_id))                    .otherwise(f.col('min1'))                )                .drop('n')            )            # logic for deciding whether to persist of checkpoint            if (niter % checkpoint_every) == 0:                filepath = (                    checkpoint_dir + filepattern.format(niter)                )                used_filepaths.append(filepath)                mins.write.parquet(filepath, mode='overwrite')                mins = ss.read.parquet(filepath)            else:                mins = mins.persist()        # update inputs for stopping logic        not_done = mins.filter(f.col('change') == True).count()        old_done = new_done        new_done = not_done        n_components = mins.select(            f.countDistinct(f.col(m1)).alias('n')        ).collect()[0]['n']        print(niter, not_done, n_components)    output = mins.select(        f.col(a).alias(label),         f.col(m1).alias('component')    )    if max_n is not None:        output = (            output            .withColumn(                'component',                f.when(                    f.col('component') == minimum_id,                     f.lit(None)                ).otherwise(f.col('component'))            )        )    output = output.persist()    n_components = output.select(        f.countDistinct(f.col('component')).alias('n')    ).collect()[0]['n']    print('total components:', n_components)    # clean up checkpointed files    for filepath in used_filepaths:        subprocess.call(["hdfs", "dfs", "-rm", "-r", filepath])return output