A Deep Dive into Parallelizing Clustering Algorithms with PySpark

From MapReduce patterns in K-Means to the complex neighborhood problem in DBSCAN, here’s a developer’s guide to building scalable clustering from scratch.

While libraries like scikit-learn provide powerful, easy-to-use clustering implementations, they are fundamentally limited to data that can be processed on a single machine. When datasets grow beyond the capacity of a single node’s memory, we must turn to distributed computing frameworks like Apache Spark.

This article provides a deep dive into implementing three fundamental clustering algorithms—K-Means, Gaussian Mixture Models (GMM), and DBSCAN—using PySpark. We won’t just use a high-level API; we’ll build them from the ground up to understand the specific parallelization strategies and challenges inherent to each. We will focus on the core logic of distributing the computation, using Spark’s Resilient Distributed Datasets (RDDs) as our foundation.

All the code discussed is part of a single, self-contained script. Let’s begin by setting up our Spark environment.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs, make_moons
from scipy.stats import multivariate_normal
import findspark

# Make sure pyspark is found
findspark.init()

from pyspark import SparkContext

# Initialize SparkContext to use 4 local cores for parallel execution
sc = SparkContext('local[4]', 'DistributedClustering')

# --- Create Sample Data RDDs ---
# Data for K-Means and GMM (well-separated blobs)
X_blobs, _ = make_blobs(n_samples=2000, centers=3, n_features=2, random_state=42, cluster_std=1.2)
rdd_blobs = sc.parallelize(X_blobs.tolist())

# Data for DBSCAN (non-linear shapes)
X_moons, _ = make_moons(n_samples=1000, noise=0.1, random_state=42)
rdd_moons = sc.parallelize(X_moons.tolist())

# Helper function for distance, used across algorithms
def euclidean_distance(p1, p2):
    return np.sqrt(np.sum((np.array(p1) - np.array(p2))**2))

print(f"Number of partitions for blobs RDD: {rdd_blobs.getNumPartitions()}")
# >>> Number of partitions for blobs RDD: 4

The RDD is the core data structure, and in this case, it’s been split into 4 partitions. Our goal is to perform as much computation as possible on these partitions independently before shuffling any data.

1. K-Means: The MapReduce Archetype

K-Means aims to partition $n$ observations into $K$ clusters by minimizing the within-cluster sum of squares (inertia). The algorithm iterates between two steps: the Expectation (E-step) and the Maximization (M-step).

E-Step: Assign each data point to its nearest centroid.
M-Step: Recalculate each centroid as the mean of all points assigned to it.

This process lends itself perfectly to the MapReduce paradigm, which is a cornerstone of Spark’s execution model.

K-Means Code Implementation

The KMeans_PySpark class encapsulates this logic. The fit method contains the main iterative loop.

class KMeans_PySpark:
    def __init__(self, K=3, max_iters=100):
        self.K = K
        self.max_iters = max_iters
        self.centroids = []

    def _closest_centroid(self, point, centroids):
        distances = [euclidean_distance(point, c) for c in centroids]
        return np.argmin(distances)

    def _is_converged(self, old, new):
        dist = sum([euclidean_distance(old[i], new[i]) for i in range(self.K)])
        return dist == 0

    def fit(self, rdd):
        # --- Step 1: Initialize centroids on the driver ---
        # takeSample is an action, returning K random points to the driver.
        self.centroids = rdd.takeSample(False, self.K, seed=42)

        for i in range(self.max_iters):
            # --- Step 2: Broadcast centroids and assign points (E-step) ---
            # Broadcast read-only centroids to all worker nodes efficiently.
            centroids_bc = sc.broadcast(self.centroids)
            
            # map: each worker executes this lambda on every point in its partition.
            # It returns a (key, value) pair: (cluster_index, (point_vector, 1)).
            closest_rdd = rdd.map(
                lambda point: (self._closest_centroid(point, centroids_bc.value), (point, 1))
            )
            
            # --- Step 3: Update centroids (M-step) ---
            # reduceByKey: performs a parallel aggregation on the workers.
            # For each key (cluster_index), it aggregates the tuples.
            # p1 is the accumulator, p2 is the next value.
            # It sums the point vectors and the counts (the '1's) element-wise.
            sum_counts_rdd = closest_rdd.reduceByKey(
                lambda p1, p2: ([p1[0][i] + p2[0][i] for i in range(len(p1[0]))], p1[1] + p2[1])
            )
            
            # The collect() action brings the K aggregated results to the driver.
            # This is a small amount of data (K vectors and K counts).
            cluster_sums = sum_counts_rdd.collect()
            
            # --- Step 4: Finalize new centroids on the driver ---
            new_centroids = self.centroids[:] # Copy
            for cluster_idx, (point_sum, count) in cluster_sums:
                new_centroids[cluster_idx] = [coord / count for coord in point_sum]
                
            # Check for convergence on the driver
            if self._is_converged(self.centroids, new_centroids):
                print(f"Converged at iteration {i+1}")
                break
            
            self.centroids = new_centroids

    def predict(self, rdd):
        # Broadcast the final centroids for prediction
        centroids_bc = sc.broadcast(self.centroids)
        # Simple map operation to assign a final cluster to each point
        labels = rdd.map(lambda point: self._closest_centroid(point, centroids_bc.value))
        return labels

Key parallel concepts in action:

Broadcasting: We use sc.broadcast() to avoid sending a new copy of the centroids with every task. The broadcast variable is cached on each worker.
Map: The assignment step is “embarrassingly parallel.” Each point’s assignment is independent of all others, making map the perfect tool.
ReduceByKey: This is the workhorse of the M-step. It performs the aggregation on the worker nodes before shuffling any data. Only the final K aggregated results are sent over the network to the driver, minimizing traffic.

K-Means Result

The result is a clean partitioning of the data, computed in a scalable fashion.

2. GMM & The Expectation-Maximization (EM) Algorithm

Gaussian Mixture Models (GMMs) assume that the data points are generated from a mixture of several Gaussian distributions. The goal is to learn the parameters of these distributions: the weights ($\pi\_k$), means ($\mu\_k$), and covariances ($\Sigma\_k$). This is achieved with the Expectation-Maximization (EM) algorithm.

E-Step (Expectation): We calculate the responsibility $\gamma(z\_{nk})$, which is the posterior probability that point $x\_n$ was generated by component $k$. $\gamma(z_{nk}) = \frac{\pi_k \mathcal{N}(x_n | \mu_k, \Sigma_k)}{\sum_{j=1}^{K} \pi_j \mathcal{N}(x_n | \mu_j, \Sigma_j)}$
M-Step (Maximization): We use the responsibilities to update the model parameters. The total responsibility for a cluster $k$ is $N_k = \sum_{n=1}^{N} \gamma(z_{nk})$. The parameters are updated as follows: $\mu_k^{new} = \frac{1}{N_k} \sum_{n=1}^{N} \gamma(z_{nk})x_n$ $\Sigma_k^{new} = \frac{1}{N_k} \sum_{n=1}^{N} \gamma(z_{nk})(x_n - \mu_k^{new})(x_n - \mu_k^{new})^T$ $\pi_k^{new} = \frac{N_k}{N}$

GMM Code Implementation

This process is more computationally intensive than K-Means. To optimize, we use mapPartitions instead of map. This allows us to perform the entire E-step calculation for a full partition of data at once, reducing overhead and enabling more efficient, vectorized computations with NumPy.

class GMM_PySpark:
    def __init__(self, n_components=3, max_iter=100, tol=1e-4):
        self.n_components = n_components
        self.max_iter = max_iter
        self.tol = tol
        # Parameters to be learned
        self.weights_ = None
        self.means_ = None
        self.covariances_ = None

    def fit(self, rdd):
        n_samples = rdd.count()
        n_features = len(rdd.first())
        
        # --- Step 1: Initialize parameters on the driver ---
        self.weights_ = np.full(self.n_components, 1 / self.n_components)
        self.means_ = np.array(rdd.takeSample(False, self.n_components, seed=42))
        self.covariances_ = np.array([np.eye(n_features) for _ in range(self.n_components)])
        
        log_likelihood = -np.inf

        for i in range(self.max_iter):
            prev_log_likelihood = log_likelihood
            # Broadcast the current model parameters to all workers
            params_bc = sc.broadcast((self.weights_, self.means_, self.covariances_))
            
            # --- Step 2: E-Step in parallel using mapPartitions ---
            # mapPartitions is more efficient as it processes an entire iterator of points
            # at once, allowing for vectorized numpy operations.
            # It calculates responsibilities and aggregates the necessary stats for the M-step
            # *locally* on each partition.
            stats = rdd.mapPartitions(lambda iterator: self._e_step_partition(iterator, params_bc.value)) \
                       .treeReduce(self._reduce_stats) # Efficiently aggregates stats from all partitions
            
            log_likelihood = stats['log_likelihood']
            
            # --- Step 3: M-Step on the driver ---
            # The driver uses the single, aggregated 'stats' dictionary to update parameters.
            self._m_step(stats, n_samples)
            
            # --- Step 4: Check for convergence on the driver ---
            if abs(log_likelihood - prev_log_likelihood) < self.tol:
                print(f"Converged at iteration {i+1}")
                break

    def _calculate_responsibilities(self, points, weights, means, covariances):
        # This is a vectorized calculation for a batch of points
        likelihoods = np.zeros((points.shape[0], self.n_components))
        for k in range(self.n_components):
            likelihoods[:, k] = multivariate_normal.pdf(points, means[k], covariances[k], allow_singular=True)
        
        weighted_likelihoods = likelihoods * weights
        log_likelihood = np.sum(np.log(weighted_likelihoods.sum(axis=1)))
        
        # Calculate responsibilities (the core of the E-step)
        responsibilities = weighted_likelihoods / weighted_likelihoods.sum(axis=1)[:, np.newaxis]
        return responsibilities, log_likelihood

    def _e_step_partition(self, iterator, params):
        # This function runs on each worker for one partition of data
        points = np.array(list(iterator))
        if points.shape[0] == 0:
            # Handle empty partitions
            yield { 'log_likelihood': 0, 'resp_sum': np.zeros(self.n_components), 
                    'weighted_points': np.zeros_like(self.means_), 
                    'weighted_cov': np.zeros_like(self.covariances_) }
            return

        # Calculate responsibilities for all points in this partition
        responsibilities, log_likelihood = self._calculate_responsibilities(points, *params)
        
        # Aggregate statistics needed for the M-step for this partition
        resp_sum = responsibilities.sum(axis=0) # This will be Nk for this partition
        weighted_points = np.dot(responsibilities.T, points) # Sum of gamma * xn
        
        weighted_cov = np.zeros_like(self.covariances_)
        for k in range(self.n_components):
            diff = points - params[1][k] # (xn - mu_k)
            weighted_cov[k] = np.dot(responsibilities[:, k] * diff.T, diff) # Sum of gamma * (xn-mu)(xn-mu)^T
            
        # Yield a single dictionary of stats for this entire partition
        yield { 'log_likelihood': log_likelihood, 'resp_sum': resp_sum,
                'weighted_points': weighted_points, 'weighted_cov': weighted_cov }

    def _reduce_stats(self, s1, s2):
        # A simple reducer function to combine the stats dictionaries from two partitions
        return { 'log_likelihood': s1['log_likelihood'] + s2['log_likelihood'],
                 'resp_sum': s1['resp_sum'] + s2['resp_sum'],
                 'weighted_points': s1['weighted_points'] + s2['weighted_points'],
                 'weighted_cov': s1['weighted_cov'] + s2['weighted_cov'] }

    def _m_step(self, stats, n_samples):
        # Update parameters on the driver using the fully aggregated stats
        Nk = stats['resp_sum']
        self.weights_ = Nk / n_samples
        self.means_ = stats['weighted_points'] / Nk[:, np.newaxis]
        self.covariances_ = stats['weighted_cov'] / Nk[:, np.newaxis, np.newaxis]

    def predict(self, rdd):
        params_bc = sc.broadcast((self.weights_, self.means_, self.covariances_))
        
        def get_labels(iterator):
            points = np.array(list(iterator))
            if points.shape[0] == 0: return []
            resp, _ = self._calculate_responsibilities(points, *params_bc.value)
            return np.argmax(resp, axis=1) # Hard assignment for prediction

        return rdd.mapPartitions(get_labels)

Key parallel concepts in action:

mapPartitions: The star of the show. By processing an entire partition at once, we convert the data to a NumPy array and perform highly optimized matrix operations. This is far more efficient than applying a function one point at a time.
Local Aggregation: The _e_step_partition function not only calculates responsibilities but also performs the first level of aggregation. This reduces the amount of data that needs to be combined later.
treeReduce: This action combines the results from partitions more efficiently than a standard reduce by using a tree-like pattern, which is better for performance on large clusters.

GMM Result

The resulting plot shows the soft assignments GMM is known for, capable of capturing elliptical cluster shapes.

3. DBSCAN: The Challenge of Distributed Neighborhoods

DBSCAN (Density-Based Spatial Clustering of Applications with Noise) is fundamentally different. It groups together points that are closely packed together, marking as outliers points that lie alone in low-density regions.

Core Concepts:
- eps: The radius to consider when searching for neighbors.
- min_samples: The minimum number of neighbors (including the point itself) required for a point to be considered a core point.
- Core Point: A point with at least min_samples neighbors within eps.
- Border Point: A point that is not a core point but is a neighbor of a core point.
- Noise Point: A point that is neither a core nor a border point.
Algorithm: The algorithm expands clusters by connecting reachable core points. If point p is a core point, it forms a cluster with all its neighbors. If any of those neighbors are also core points, their neighbors are added to the cluster, and so on.

This expansion process is inherently sequential and non-local. A single cluster can meander across the entire dataset, meaning a point in Partition 1 might be part of the same cluster as a point in Partition 4. This makes a true, one-step parallelization nearly impossible.

DBSCAN Code Implementation

Our implementation uses a common and practical multi-stage approach. Stage 1 is parallel, while Stage 2 (which we only describe conceptually) presents the major challenge.

# First, we need a single-node implementation of DBSCAN that can be run on each partition.
class DBSCAN_single_node:
    def __init__(self, eps=0.2, min_samples=5):
        self.eps = eps; self.min_samples = min_samples
    def fit_predict(self, X):
        X = np.array(X); n_samples = X.shape[0]; self.X = X
        self.labels = np.full(n_samples, -2) # -2: unvisited, -1: noise
        cluster_id = 0
        for i in range(n_samples):
            if self.labels[i] != -2: continue
            neighbors = self._get_neighbors(i)
            if len(neighbors) < self.min_samples:
                self.labels[i] = -1 # Mark as noise
            else:
                self._expand_cluster(i, neighbors, cluster_id)
                cluster_id += 1
        return self.labels
    def _get_neighbors(self, point_idx):
        neighbors = []
        for i in range(self.X.shape[0]):
            if euclidean_distance(self.X[point_idx], self.X[i]) < self.eps:
                neighbors.append(i)
        return neighbors
    def _expand_cluster(self, core_idx, neighbors, c_id):
        self.labels[core_idx] = c_id
        i = 0
        while i < len(neighbors):
            n_idx = neighbors[i]
            if self.labels[n_idx] in [-1, -2]: # If noise or unvisited
                self.labels[n_idx] = c_id
                new_neighbors = self._get_neighbors(n_idx)
                if len(new_neighbors) >= self.min_samples:
                    # This neighbor is also a core point, add its neighbors to the queue
                    neighbors.extend([n for n in new_neighbors if n not in neighbors])
            i += 1
            
class DBSCAN_PySpark:
    def __init__(self, eps=0.2, min_samples=5):
        self.eps = eps
        self.min_samples = min_samples

    def fit_predict(self, rdd):
        # --- Stage 1: Run DBSCAN locally on each partition ---
        def run_local_dbscan(iterator):
            points = list(iterator)
            if not points: return []
            
            local_dbscan = DBSCAN_single_node(eps=self.eps, min_samples=self.min_samples)
            labels = local_dbscan.fit_predict(points)
            
            for point, label in zip(points, labels):
                yield (tuple(point), label)

        # mapPartitionsWithIndex gives us the partition ID, which we use to create
        # a globally unique ID for each local cluster (e.g., "partition_localcluster").
        partition_results_rdd = rdd.mapPartitionsWithIndex(
            lambda partition_id, iterator: [
                (point, f"{partition_id}_{label}") for point, label in run_local_dbscan(iterator) if label != -1
            ]
        )
        
        print("--- DBSCAN STAGE 1 (Local Clustering) COMPLETE ---")
        print("This implementation returns the results of the parallel local clustering.")
        
        # --- Stage 2: Merging (Conceptual & Computationally Hard) ---
        # 1. Identify all core points near the "edge" of each partition's data.
        # 2. For every pair of core points (p1, p2) from DIFFERENT partitions:
        #    - If distance(p1, p2) < eps, then their clusters must be merged.
        # 3. This creates a graph where nodes are local clusters and edges represent merges.
        # 4. Finding the connected components of this graph gives the final global clusters.
        # This requires a massive all-to-all comparison (e.g., a cartesian product of
        # edge points) and is a major research and engineering challenge in itself.

        return partition_results_rdd

Key parallel concepts in action:

mapPartitionsWithIndex: We use the “WithIndex” variant to get the ID of each partition. This is crucial for creating unique labels for the clusters found within each partition (e.g., 0_0, 0_1, 1_0). Without this, we wouldn’t know that cluster 0 from partition 0 is different from cluster 0 from partition 1.
Two-Stage Approach: The code acknowledges the algorithm’s difficulty by only implementing the first, parallelizable stage. This is a practical approach where an approximate solution is generated quickly, and a more complex (and expensive) merge step could follow if required.

DBSCAN Result

The plot shows the result of Stage 1. The characteristic “moon” shapes are found, but they are broken into multiple colors. Each color represents a distinct local cluster. A full implementation would merge the adjacent fragments of the same moon into a single cluster.

Final Conclusion

Implementing clustering algorithms in PySpark reveals that there is no one-size-fits-all parallelization strategy. The ideal approach is intimately tied to the mathematical properties of the algorithm itself.

K-Means is a perfect fit for the simple, powerful MapReduce pattern due to its independent point assignments and simple aggregation step.
GMM benefits from the more optimized mapPartitions approach, which leverages vectorized computations for its more complex, but still partition-local, E-step.
DBSCAN exposes the limits of simple parallelism. Its reliance on non-local neighborhood information requires complex, multi-stage algorithms that must explicitly handle data dependencies across partition boundaries.

Understanding these patterns is key to moving beyond single-node libraries and building truly scalable machine learning systems. The choice of how to distribute your computation is just as important as the choice of the algorithm itself.

Share on

Bluesky Facebook LinkedIn X (formerly Twitter)

Saad Lahlali

From MapReduce patterns in K-Means to the complex neighborhood problem in DBSCAN, here’s a developer’s guide to building scalable clustering from scratch.

1. K-Means: The MapReduce Archetype

K-Means Code Implementation

K-Means Result

2. GMM & The Expectation-Maximization (EM) Algorithm

GMM Code Implementation

GMM Result

3. DBSCAN: The Challenge of Distributed Neighborhoods

DBSCAN Code Implementation

DBSCAN Result

Final Conclusion

Share on