Streaming cluster fit and assign

Clustering at corpus scale is not “load all vectors, call k-means once”. Framework splits the work into fit (learn centroids) and assign (label rows and write), both streaming over bounded vector batches. Peak memory stays roughly centroids + one batch + write buffer, not the full run matrix.

Layers

Layer	Owns
algorithms/cluster	Pure NumPy: nearest-centroid assign, centroid stats merge, in-memory minibatch_kmeans_v1
storage/run/vector_batch_iter.py	Stream numpy windows (vectors, row_ids[, metas]) from parquet/Qdrant via existing logical-row reads
processing/cluster	Fit/assign loops, shard fan-out, write chunks; hash-vote reduce in hash_mode_reduce.py

The step still chooses which K (or which candidate list) and what columns mean. Framework does not invent domain column names: callers pass an AssignRowSchema / DRTML bindings.

Vector batches (not feed batches)

A vector batch is a small contiguous numpy window for geometry math. It is distinct from:

• feed ProcessorBatch (domain processor rows);
• Arrow RecordBatch (parquet transport);
• embed train pair batches (SGNS updates).

iter_vector_batches projects id/vector/(optional) meta columns, skips empty or non-finite vectors, supports equality filters and an optional max_rows cap (useful when probing on a sample). Peak without that cap tracks batch_rows, not on-disk parquet size.

Fit: one shared model

Fit never returns a full label vector — only a ClusterFitResult (centroids, k, row count, inertia, backend metadata).

Serial streaming re-opens the batch factory each epoch:

1. pull early batches into an init buffer and sample starting centroids;
2. for each iteration, either
   • epoch (default for parallel): assign every batch, accumulate (sum, count) per cluster, merge into new centroids (order-independent → shard-friendly), or
   • minibatch: classic per-batch mean updates while scanning (serial order);
3. optional final pass for inertia.

Parallel fit uses the same runtime.parallelism.* plan as other modes (inprocess threads, processes, or ray). Work units are vector parquet shards. Each worker holds a broadcast copy of the current centroids, scans its shards, and returns local stats. The driver merges stats into one new centroid table. That is still one global model, not independent k-means per shard.

workers=1 or a single shard falls back to the serial path.

Sticky cache (default)

With compute.cluster.fit_sticky_cache (default on) on Ray / threads / inprocess, shards load once into worker RAM or Ray FitShardActors (open_sticky_fit_session). The same cache serves many probe K values via fit_then_assign — no parquet re-read per candidate. Stateless re-scan remains available when sticky is off or unsupported.

Shard packing uses LPT by manifest row_count (partition_vector_shards); unknown counts fall back to round-robin.

Warm-start and early-stop

• compute.cluster.warm_start (default on): ascending-K init expands the previous centroids with farthest-point selection (GEMM in algorithms/cluster), not cold random init each time.
• Epoch fit can stop early when relative inertia improvement stays below a threshold for patience epochs (InertiaEarlyStop). Domain still chooses K / winner; the stop rule is a compute knob on the fit loop.

Assign: label and publish in chunks

With fixed centroids, assign streams batches → blockwise nearest-centroid labels → logical rows → session.write in chunks (existing flush policy). Optional cluster-size rows are aggregated after counting.

Parallel assign broadcasts the same centroids; each worker writes under its own output shard id; the driver merges size counts. When several rows share a vote key (hash mode), HashClusterVoteAccumulator keeps majority votes — peak tracks unique keys, not N×D.

On a sticky session, stats-only / vote collect reuse cached matrices (GEMM + np.bincount size counts). Multiple assign payloads merge with an associative tree-merge on the caller (ThreadPool); sticky Ray gathers actor payloads then merges locally.

A stats-only path (assign_collect_stats) runs the same assign math without writing — useful inside probe prepare/evaluate.

What stays out of framework

• Choosing K and “which probe candidate wins” — domain policy (see Probe loop).
• Homegrown sklearn/torch k-means in the step — use algorithms + these loops.
• Holding the full infer matrix “because parquet looked small on disk” — see the memory rules in Parallelism.

Read next

• Probe loop
• Parallelism
• StorageSession
• Algorithm catalog
• Offline evaluation