The first time k-means broke for me, it did not get slow. It refused to start. A billion 8-dimensional points sat in object storage, and sklearn.cluster.KMeans never reached its first iteration: it tried to materialize the array, hit the memory ceiling, and died. The algorithm was fine. The machine was the problem.

That is the stake. The interesting failure is not “my clustering takes too long,” it is “the loop cannot run at all, because the data does not fit.” Hold that distinction, because everything that follows is about it.

So a warning up front: this post is not a benchmark and it has no speedup numbers, because those would be specific to one machine and one dataset and would not transfer to yours. What transfers is the reasoning, namely when you should distribute k-means and, once you do, exactly how the data moves.

I will carry one example the whole way: a billion 8-dimensional points in object storage, far too large for your laptop. That is the dataset, and every section below is what happens to it.

What distributing actually changes

Standard k-means is a short loop:

  1. Initialize: pick k cluster centroids.
  2. Assign: put each point with its nearest centroid.
  3. Update: move each centroid to the mean of its assigned points.
  4. Check: test the stop condition (centroids stopped moving, or you hit the iteration cap).
  5. Repeat steps 2 and 3 until then.

On one machine that is a few lines, and for moderate data it is genuinely fast:

import numpy as np
from sklearn.cluster import KMeans

X = np.random.randn(1000, 2)
kmeans = KMeans(n_clusters=3, random_state=42)
labels = kmeans.fit_predict(X)

The cost of that loop is $O(n \times k \times d \times i)$, where n is the number of points, k is the number of clusters, d is the dimensionality, and i is the number of iterations. The term that hurts is n, because every one of those n points gets touched k times on every iteration. When n is a billion and the array will not fit in memory, the loop does not slow down, it stops being runnable. That is the wall.

Now look at where the work lives, because one fact decides everything else: the assignment step is the expensive one, and each point is assigned independently of every other point. To decide which centroid a point belongs to, you need that point and the current k centroids, and nothing else. That independence is the entire reason distributing k-means works.

So you pin the data where it already lives. You split the billion points across workers, each worker holding a partition that never moves, and broadcast the k centroids to every worker. Each worker assigns its own partition locally, which is embarrassingly parallel, and instead of shipping points anywhere, each worker emits a tiny summary per cluster: the sum of its points and a count. A tree-reduce combines those partial sums and counts into k updated centroids, those centroids broadcast back out, and the next iteration begins.

The thing that crosses the wire each iteration is k centroids, not n points. For our billion points in 8 dimensions, that is a handful of tiny vectors moving instead of the entire dataset. The data stays pinned. Only the summaries travel.

One distributed k-means iteration Three worker partitions stay pinned in place. Each assigns its own points locally and emits per-cluster partial sums and counts. A tree-reduce combines them into k updated centroids, which are broadcast back to every worker for the next iteration. Only the k centroids cross the wire. data stays pinned partition 1 assign locally partition 2 assign locally partition 3 sums + counts tree-reduce combine partials k centroids updated broadcast k centroids back to every worker
One iteration of distributed k-means. Partitions never move. Each worker assigns its own points and emits per-cluster sums and counts, a tree-reduce folds them into k updated centroids, and those k centroids broadcast back out to all three workers. Tiny summaries flow up, a tiny model flows down.

The real bottleneck is initialization

Once you see the iteration that way, the loop is cheap to distribute and initialization becomes where the cost moves.

Classic k-means++ picks seeds one at a time. Each new centroid is chosen with probability proportional to its squared distance from the seeds chosen so far, so you have to scan the whole dataset before you can pick the next seed, and that is k sequential passes over the data. On one machine, fine. On a cluster, each of those passes is a full distributed scan, effectively k full distributed passes stacked back to back, all of them before the actual clustering loop begins. For a large k over a billion points, you can spend longer seeding than clustering.

The fix is **k-means   ** (k-means parallel). Instead of drawing one seed per pass, it oversamples: each pass draws several candidate centers at once, building up a candidate pool in O(log k) passes rather than k sequential passes. Then it reduces that pool down to k final seeds with one local reduction, no further distributed scans. Same goal as k-means++, good well-spread seeds, but the number of passes over the distributed data drops from k to O(log k).
k-means++ versus k-means parallel seeding On the left, k-means++ runs k sequential full passes over distributed data, one seed per pass, drawn as a tall stack of repeated scans. On the right, k-means parallel runs a few oversampling passes labeled O of log k that gather a candidate pool, then one local reduction down to k seeds. k-means++ k sequential passes full scan, +1 seed full scan, +1 seed full scan, +1 seed ... k times k-means|| O(log k) oversampling passes pass: oversample candidates pass: oversample candidates pass: oversample candidates candidate pool local reduce to k seeds no distributed scan
Why k-means|| wins at scale. k-means++ needs k full distributed passes, one seed each. k-means|| oversamples in O(log k) passes into a candidate pool, then reduces locally to k seeds with no further scan.
This is why Spark’s MLlib defaults to k-means   , and it is the one distributed-specific lever you actually need to know: 
from pyspark.ml.clustering import KMeans as SparkKMeans

# k-means|| (k-means parallel) is the scalable seeding, not k-means++.
# Both aim for good seeds, but k-means|| runs in O(log k) passes
# rather than k sequential passes.
kmeans = SparkKMeans(k=5, initMode="k-means||", initSteps=2)
kmeans.setTol(1e-4)

Two engines, one algorithm

PySpark MLlib and Dask-ML both implement the data-movement model from the first section, and they differ in which world they live in, not in what k-means does.

Reach for Spark when the data already lives in the Spark or JVM ecosystem and is huge. You load it, pack the columns into a single features vector with VectorAssembler, and fit, and Spark handles the partitioning and the reduce:

from pyspark.ml.clustering import KMeans as SparkKMeans
from pyspark.ml.feature import VectorAssembler

df = spark.read.option("header", "true").csv(data_path)
features = VectorAssembler(inputCols=numeric_cols, outputCol="features")
df_vec = features.transform(df)

model = SparkKMeans(k=5, maxIter=50, featuresCol="features").fit(df_vec)
predictions = model.transform(df_vec)

Reach for Dask when the workflow is already NumPy and Pandas native. The API mirrors scikit-learn, so the jump is small: read into a Dask dataframe, hand it a Dask array, fit.

import dask.dataframe as dd
from dask_ml.cluster import KMeans as DaskKMeans

df = dd.read_csv(data_path)
X = df[numeric_cols].to_dask_array(lengths=True)

kmeans = DaskKMeans(n_clusters=5, max_iter=100, init_max_iter=3)
kmeans.fit(X)
labels = kmeans.predict(X)

The original version of this post carried a rule of thumb that Spark earns its keep on very large datasets, on the order of >1TB, but treat that as a heuristic, not a measured threshold. The real decision is ecosystem: if your pipeline is JVM and the data is already in Spark, stay in Spark, and if your pipeline is Python and the data is already in arrays, stay in Dask. The friction of crossing that boundary usually outweighs any per-iteration difference.

When the data never stops: streaming

The two engines above assume the data, however large, eventually lands somewhere and you fit over all of it. There is a third regime: data that never fully lands, a stream you see one batch at a time. For that you want partial_fit, which updates the centroids from a single mini-batch and then forgets it.

Here is a sharp edge worth knowing, because the obvious move does not work. dask_ml.cluster.KMeans does not expose partial_fit, so call it and you get an AttributeError. The streaming path is to wrap scikit-learn’s MiniBatchKMeans, which does support partial_fit, in dask_ml.wrappers.Incremental:

from sklearn.cluster import MiniBatchKMeans
from dask_ml.wrappers import Incremental

def incremental_kmeans_dask(data_stream, k=3):
    base = MiniBatchKMeans(n_clusters=k, random_state=42)
    kmeans = Incremental(base)
    for batch in data_stream:
        # We drive the stream ourselves, feeding explicit mini-batches.
        # Incremental delegates each call to partial_fit on the wrapped
        # estimator (here it is not mapping over Dask-array blocks).
        kmeans.partial_fit(batch)
    return kmeans

Each batch nudges the centroids a little, then is discarded, and you never hold the whole stream. That is the entire point. This is the regime where even “fits across the cluster” stops being true, because there is no fixed dataset to fit, only an arrival rate.

Batch fit versus streaming partial_fit On the top, batch fit ingests the whole dataset at once into a set of centroids. On the bottom, an unbounded stream of mini-batches flows one at a time into a single evolving set of centroids, each batch nudging the centroids in place before being discarded. batch fit whole dataset k centroids streaming partial_fit b1 b2 b3 b4 ... k centroids nudged update in place each batch updates, then is discarded
Two regimes. Batch fit pulls the whole dataset into one fit. Streaming partial_fit lets an unbounded sequence of mini-batches each nudge the same evolving centroids in place, holding only one batch at a time.

Choosing k without eyeballing

One practical note survives the move to a cluster: you still have to pick k. The elbow method (sweep k, plot the within-cluster sum of squares, look for the bend) usually means squinting at a chart, but you can automate the squint. For each k, record the inertia, then pick the k whose point on the curve sits farthest, in perpendicular distance, from the straight line joining the first and last points. A no-eyeball elbow detector:

import numpy as np

def find_elbow_point(k_range, inertias):
    points = np.array(list(zip(k_range, inertias)), dtype=float)
    line_vec = points[-1] - points[0]
    line_vec_norm = line_vec / np.linalg.norm(line_vec)
    vec_from_first = points - points[0]
    scalar_proj = vec_from_first @ line_vec_norm
    proj_points = np.outer(scalar_proj, line_vec_norm) + points[0]
    distances = np.linalg.norm(points - proj_points, axis=1)
    return int(k_range[int(np.argmax(distances))])

One storage habit matters far more on a cluster than on a laptop: store features as float32 rather than float64 when you can. Halving the per-point footprint is often the difference between a partition fitting in a worker’s memory and spilling, which ties straight back to the stake at the top, because the whole reason you reached for a cluster was that the data stopped fitting. Standardizing and dropping outliers first (z-score, keep z_scores < 3) still helps the clustering itself, the same way it does on one machine.

The pattern, and the limit

Strip away the engines and distributed k-means is one idea: figure out which step is embarrassingly parallel, run that step where the data already sits, and ship as little as possible per round. Here the assignment step parallelizes per partition, the update step is a tree-reduce of partial sums, and the only thing crossing the wire each iteration is the k centroids.

Decision: distribute k-means or not A flow that starts by asking whether the data fits in one machine's RAM. If yes, use single-node MiniBatchKMeans and the cluster almost always loses. If no, distribute, then branch on Spark for JVM or huge data, Dask for NumPy or Pandas native workflows, or Incremental for streaming data that never fully lands. does it fit in one machine's RAM? yes single-node MiniBatchKMeans the cluster almost always loses no: distribute which engine? Spark JVM / huge Dask NumPy native stream Incremental
The decision that matters. The first question is whether the data fits in one machine. If it does, a single-node MiniBatchKMeans almost always beats a cluster. Only when it does not do you distribute, and then the choice is ecosystem: Spark, Dask, or Incremental for streams.

That same shape, parallelize the per-record step, reduce small summaries, broadcast a small model, recurs across a lot of distributed machine learning. Gradient descent over sharded data does it, distributed gradient boosting does it, and federated training does it. K-means is just the version where the model is small enough that you can watch the trick happen.

And none of it is free. You pay shuffle cost on every reduce and coordination cost across workers, and that overhead is real, so for anything that fits on one machine, a single-node MiniBatchKMeans will beat a cluster, because the cluster spends its first seconds just agreeing on who holds what. Distribution is not a default. It is what you reach for when the data has already won the argument against your RAM. Data stays pinned, and the centroids are the only thing that crosses the wire.

You May Also Enjoy