KMeans Clustering in Apache DataFusion

Let's assume we have a table with data in the following format:

1. id column, let's say in64 but it can be any kind of unique ID
2. feature column, a dense vector of the size d filled by f32 values

Our goal is to find K clusters. See Wikipedia for more details.

Long story short, we need to:

1. Initialise K "centers": vectors of the same size d
2. For each point in the data, we need to find the nearest center based on distance between features and center
3. We need to recompute the vector of each center as an average of all the points for that this center is the nearest

Continue this until convergence.

Let's imagine we have an aggregate function (UDAF) that is paremetrized by the "centers" and for each point in data compute the distance to all the centers, choose the closest one and update this cluster sum of coridnates and count of points in it. Then we can express one iteration of KMeans in the form like this:

  SELECT r.sums, r.counts FROM
  (SELECT kmeans_iteratio(features, centers) AS r FROM dataset GROUP BY 1);

Or in DataFusion DataFrame API:

  let expr = kmeans_step_expr(col("features"), k, d, centers.clone());
  let batches = features.clone().aggregate(vec![], vec![expr])?.collect().await?;
  let sums = batches[0]
               .column(0)
               .as_any()
               .downcast_ref::<FixedSizeListArray>()
               .ok_or_else(|| {
                    datafusion::common::DataFusionError::Execution(
                        "k-means step result must be a FixedSizeList".to_string(),
                    )
                })?;
  let cnts = batches[0]
               .columns(1)
               .as_any()
               .downcast_ref::<Float64Array>()
               .expect("k_means_step returns Float64");
  // for each value we should update centers by dividing
  // each value from sums by the value from cnts <...>

NOTE: As one may already see, if the dataset is a table or parquet file, the peak memory usage of one iteration is O(|K*d|) that does not depend on amount of data in the table. DataFusion will read batches one by one, compute distances, update state and free the memory for the next batch. The only thing that is store in mmeory is the state: k vectors for sums and k double values for counts. This is out-of-core implementation by design opposite to one from scikit-learn that requires to have all the data in memory.

That was the first time in my live I wrote manual SIMD kernels and I can say it was fun! I used a wide crate for it. The L2 distance kernel is like this:

  pub(crate) fn l2_distance(x: &[f32], c: &[f32], d: usize) -> f32 {
    let mut acc = f32x8::splat(0.0);
    let mut t = 0;
    while t + 8 <= d {
        let xv = f32x8::from(&x[t..t + 8]);
        let cv = f32x8::from(&c[t..t + 8]);
        acc = (xv - cv).mul_add(xv - cv, acc);
        t += 8;
    }

    let mut dist = acc.reduce_add();

    while t < d {
        let xi = x[t];
        let ci = c[t];
        dist += (xi - ci) * (xi - ci);
        t += 1;
    }

    dist
  }

A "classical" approach: 8 lines, SIMD body and a non SIMD tail for last elements. Because for KMeans we need to find the nearest center, scale does not matter and we can compare squared L2 because square root is monotonic function.

For expressions like df.aggregate(vec![], vec![expr]) DataFusion falls back to the Accumulator so it does not make any sense to implement GroupsAccumulator.

What we should keep inside:

• current centers
• dimension (size of "features")
• k
• metric (if we want to support not only L2)
• sums
• counts

For performance reason I put both sums and counts into one long vector: first k*d are sums, last k are counts.

  #[derive(Debug)]
  pub(crate) struct KMeansStepAccumulator {
      k: usize,
      d: usize,
      centers: Vec<f32>,
      state: Vec<f64>,
      metric: DistanceMetric,
  }

Update batch is trivial: scan rows, for each row call the SIMD-kernel.

  fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
    debug_assert_eq!(values[0].null_count(), 0usize);
    let v = as_f32_list_like(&values[0], "k_means_step", "first")?;

    // no nulls are assumed in feature (embeddings)
    for i in 0..v.len() {
      let vv = v.value(i);
      let (cluster, _) = nearest_centers(vv, &self.centers, self.k, self.d, self.metric);
      self.state[self.k * self.d + cluster] += 1.0f64;

      for t in 0..self.d {
        self.state[cluster * self.d + t] += vv[t] as f64;
      }
    }

    Ok(())
  }

For merging two batches it is enough to serialize only the state field and pass it as an arrow's List.

  fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
    let s = as_f64_list_like(&states[0], "k_means_step", "state")?;

    let n = self.k * self.d + self.k;
    // non-nulls semantic by contract
    for i in 0..s.len() {
      let ss = s.value(i);
      for t in 0..n {
        self.state[t] += ss[t];
      }
    }

    Ok(())
  }

Most of other code is a pure DataFusion machinery.

It is long. More than 250 lines of code. I would better leave the link to the source. Long story short, it is almost direct port of the Scala code from the Spark MLLib.

One of the the main problem of KMeans algorithm is to choose the K. scikit-learn documentation recommends to fit multiple K and choose the right one by metrics like silhouette coefficients.

Working on my implementation I realized that the most expensive part of the algorithm is scan and aggregation while the number crunching in kernels is just around 30% of all the work. As well I noticed that the expensive scan and aggregate can be shared along multiple K using the same UDAF in a way like this:

  SELECT r1.sums, r1.counts, r2.sums, r2.counts, ... FROM
  (SELECT
    kmeans_iteratio(features, centers1, k1) AS r1,
    kmeans_iteratio(features, centers2, k2) AS r2,
    ...
    FROM dataset GROUP BY 1);

This feature is already in my code.

As any out-of-core implementation, this one will be always slower than the in-memory one like in scikit-learn. At the same time the overall performance is not so bad. I tested it as a part of Power Iteration Clustering on graph with 8M nodes that means KMeans was called on a dataset with 8M vectors of size 19. Wall time was few seconds with a hard limit on DataFusion memory pool 4GB.