Apache Spark Framework for Clustering Algorithms in Distributed Mode

Apache Spark is a distributed computing engine that has libraries for building data pipelines using programming APIs, SQL API, and APIs for carrying out tasks that are part of the machine learning life cycle, such as feature engineering, model training, inference, and evaluation. The advantage of SparkML is that it can leverage computing power across the nodes, VMs, or containers to execute computationally and memory-intensive model training. 

However, SparkML does not support every machine learning algorithm out of the box. Several advanced or specialized algorithms, particularly in the area of unsupervised learning, are not currently implemented in SparkML, and DBSCAN is one such unsupervised clustering algorithm. These algorithms are highly valuable for tasks where cluster boundaries are non-linear or when the number of clusters is not known in advance.

Because SparkML does not natively support these models, practitioners are often forced to train them using standalone implementations outside of the Spark framework, typically using single-node libraries such as those in Scikit-learn. This approach is not efficient because clustering multiple subsets of data sequentially using single-threaded processing can become a performance bottleneck. In contrast, if we could parallelize the execution of such algorithms across partitions or nodes in a Spark cluster by writing custom distributed training workflows, we could drastically improve training throughput and reduce time-to-insight for large-scale clustering problems.

Solution

Training clustering machine learning (ML) algorithms that are not natively supported by SparkML, such as DBSCAN or OPTICS, can still be achieved efficiently by leveraging the parallel processing capabilities of Apache Spark. Although SparkML does not provide built-in support for these algorithms, we can design a distributed training pipeline using the following techniques:

1. Custom partitioner: A custom partitioner can be used to logically divide the dataset into partitions based on domain-specific keys. For instance, if the data involves relationships between meters and transformers, we can partition the dataset such that each partition contains data related to a specific transformer. This ensures that each partition is independent and self-contained for localized clustering.
2. mapPartitions transformation: Spark’s mapPartitions transformation enables us to apply arbitrary processing logic to each partition of the dataset. This is especially useful when we want to run third-party or custom ML code, such as clustering using Scikit-learn, on each data partition.
3. Embedding clustering logic in partitions: The clustering algorithm’s training logic (for example, DBSCAN implemented in Python) can be embedded directly inside the mapPartitions function. This allows Spark to process each partition in parallel, executing the clustering logic independently on each partition without requiring native SparkML support.
4. Collecting and aggregating results: Once each partition completes its clustering task, the results can be aggregated and converted back into a Spark DataFrame. This enables further downstream processing, visualization, or actions to be performed on the clustered results in the familiar Spark ecosystem.

Let’s illustrate this approach with a practical use case.

Example Use Case: Meter-Transformer Voltage Clustering

Imagine we have a dataset that records voltage readings emitted by electric meters, each of which is associated with a transformer. The goal is to determine whether a given meter is truly associated with the transformer it is mapped to in the system. Typically, meters connected to the same transformer exhibit voltage signals that lie within a similar range due to the shared electrical characteristics of the transformer.

This insight allows us to frame the problem as a clustering task. By grouping meters based on the similarity of their voltage signals using an unsupervised clustering algorithm (such as DBSCAN), we can identify clusters of meters likely connected to the same transformer. Meters that fall outside the expected cluster (i.e., outliers) may have been misassigned and should be flagged for manual review. Technicians can verify these outlier meters and verify which transformer they are connected to, and the corresponding database is updated.

By applying a custom partitioner to group meters by transformer, and running the clustering logic within mapPartitions, multiple transformers’ data can be processed in parallel, leveraging Spark’s distributed computing and reducing overall processing time and effective resource utilization.

Custom Partitioner

A custom partitioner in Spark is a user-defined function that determines how keys in an RDD are assigned to partitions. When you use partitionBy(num_partitions, partitioner_function), Spark applies your function to each key to decide which partition it should go to. This ensures that all records with the same key (e.g., transformerId) are grouped in the same partition, enabling data locality for downstream operations like grouping, aggregation, or parallel processing on logically related data.

from pyspark import SparkContext, SparkConf
from pyspark.sql import SparkSession
from pyspark.rdd import portable_hash
from pyspark.sql.functions import col

# Initialize Spark session
spark = SparkSession.builder.appName("TransformerPartitioning").getOrCreate()

# Sample DataFrame (replace with your actual loading logic)
df = spark.read.option("header", "true").csv("your_data.csv")

# Step 1: Extract relevant key-value RDD
rdd = df.rdd.map(lambda row: (row['transformerid'], row))

# Step 2: Get number of unique transformer IDs
num_partitions = df.select("transformerid").distinct().count()

# Step 3: Define a custom partitioner function
class TransformerPartitioner:
    def __init__(self, num_partitions):
        self.num_partitions = num_partitions

    def __call__(self, key):
        # Use a hash to spread transformer IDs across partitions
        return portable_hash(key) % self.num_partitions

# Step 4: Partition the RDD using custom partitioner
partitioned_rdd = rdd.partitionBy(num_partitions, TransformerPartitioner(num_partitions))

# Step 5: Convert back to DataFrame if needed
partitioned_df = partitioned_rdd.map(lambda x: x[1]).toDF()

# Optional: Check partition count
print("Number of partitions:", partitioned_df.rdd.getNumPartitions())

Map Partitions

As our dataset contains fields such as meterId, transformerId, and voltage signals (v1, v2). The DataFrame is first repartitioned by transformerId using a custom partitioner, so that all records belonging to the same transformer ID are grouped into a unique partition. The mapPartitions function is then used to apply a custom function that collects all rows in a partition, converts them into a local Pandas DataFrame, and applies scikit-learn's DBSCAN clustering on the voltage signal features. 

from pyspark.sql import SparkSession
from pyspark.sql.functions import col
from sklearn.cluster import DBSCAN
import pandas as pd
import numpy as np

# Initialize Spark session
spark = SparkSession.builder.appName("DBSCANMapPartitions").getOrCreate()

# Sample schema: meterId, transformerId, voltage_signal_1, voltage_signal_2
data = [
    (1, 'TX1', 1.0, 0.9),
    (2, 'TX1', 1.1, 0.95),
    (3, 'TX2', 5.0, 4.9),
    (4, 'TX2', 5.2, 5.1),
    (5, 'TX3', 9.0, 9.1),
    (6, 'TX3', 9.2, 8.9),
]

df = spark.createDataFrame(data, ["meterId", "transformerId", "v1", "v2"])

# Repartition by transformerId so each transformer gets its own partition
df = df.repartition("transformerId")

# Convert DataFrame to RDD
rdd = df.rdd.mapPartitions(lambda iterator: train_dbscan(iterator))

# Define mapPartitions logic
def train_dbscan(partition):
    import pandas as pd
    from sklearn.cluster import DBSCAN

    # Convert iterator to list of rows
    rows = list(partition)
    if not rows:
        return iter([])

    # Extract features from rows (e.g., v1 and v2 columns)
    df_local = pd.DataFrame(rows, columns=["meterId", "transformerId", "v1", "v2"])
    features = df_local[["v1", "v2"]].values

    # Train DBSCAN
    dbscan = DBSCAN(eps=0.3, min_samples=2)
    labels = dbscan.fit_predict(features)

    # Return original rows with cluster label
    df_local["cluster"] = labels
    return df_local.itertuples(index=False, name=None)

# Convert back to DataFrame with schema
result_df = rdd.toDF(["meterId", "transformerId", "v1", "v2", "cluster"])
result_df.show()

The resulting cluster labels are appended to the data, and the rows with their assigned cluster are returned. Finally, the processed RDD is converted back to a Spark DataFrame for further analysis or storage. This approach enables the distributed and parallel execution of DBSCAN on logically grouped subsets of data.