feat: replace pylance kmeans implementation with scikit-learn and expose num_init and max_iter params in api (#104)

Author: Rohit Rastogi

pyproject.toml

@@ -26,13 +26,13 @@ dependencies = [
   "numpy>=2.0.0",
   "polars>=1.20.0",
   "tiktoken>=0.9.0",
-  "pylance>=0.23.2",
   "lancedb>=0.22.0",
   "openai>=1.82.0",
   "sqlglot>=26.25.3",
   "pandas>=2.2.2",
   "cloudpickle>=3.1.1",
   "jinja2>=3.1.6",
+  "scikit-learn>=1.7.1",
 ]
 
 [project.urls]

src/fenic/_backends/local/physical_plan/transform.py

@@ -10,7 +10,7 @@
 from fenic._backends.local.lineage import OperatorLineage
 from fenic._backends.local.physical_plan.utils import apply_ingestion_coercions
 from fenic._backends.local.semantic_operators.cluster import Cluster
-from fenic.core._logical_plan.plans import CacheInfo
+from fenic.core._logical_plan.plans import CacheInfo, CentroidInfo
 from fenic.core.error import InternalError
 
 if TYPE_CHECKING:
@@ -337,15 +337,19 @@ def __init__(
         by_expr: pl.Expr,
         by_expr_name: str,
         num_clusters: int,
+        max_iter: int,
+        num_init: int,
         label_column: str,
-        centroid_info: Optional[Tuple[str, int]],
+        centroid_info: Optional[CentroidInfo],
         cache_info: Optional[CacheInfo],
         session_state: LocalSessionState,
     ):
         super().__init__([child], cache_info=cache_info, session_state=session_state)
         self.by_expr = by_expr
         self.by_expr_name = by_expr_name
         self.num_clusters = num_clusters
+        self.max_iter = max_iter
+        self.num_init = num_init
         self.label_column = label_column
         self.centroid_info = centroid_info
 
@@ -359,9 +363,11 @@ def _execute(self, child_dfs: List[pl.DataFrame]) -> pl.DataFrame:
         clustered_df = Cluster(
             child_df,
             self.by_expr_name,
-            self.num_clusters,
-            self.label_column,
-            self.centroid_info,
+            num_clusters=self.num_clusters,
+            max_iter=self.max_iter,
+            num_init=self.num_init,
+            label_column=self.label_column,
+            centroid_info=self.centroid_info,
         ).execute()
 
         # Remove the temporary column we added for clustering if it wasn't in the original

src/fenic/_backends/local/semantic_operators/cluster.py

@@ -1,14 +1,15 @@
 import logging
-from typing import Optional, Tuple
+from typing import Optional
 
 import numpy as np
 import polars as pl
 import pyarrow as pa
-from lance.util import KMeans
+from sklearn.cluster import KMeans
 
 from fenic._backends.local.semantic_operators.utils import (
     filter_invalid_embeddings_expr,
 )
+from fenic.core._logical_plan.plans import CentroidInfo
 
 logger = logging.getLogger(__name__)
 
@@ -18,21 +19,23 @@ def __init__(
         self,
         input: pl.DataFrame,
         embedding_column_name: str,
-        num_centroids: int,
+        num_clusters: int,
+        max_iter: int,
+        num_init: int,
         label_column: str,
-        centroid_info: Optional[Tuple[str, int]],
-        num_iter: int = 50,
+        centroid_info: Optional[CentroidInfo],
     ):
         self.input = input
         self.embedding_column_name = embedding_column_name
         input_height = input.height
-        if num_centroids > input_height:
+        if num_clusters > input_height:
             logger.warning(
-                f"`num_centroids` was set to {num_centroids}, but the input DataFrame only contains {input_height} rows. "
-                f"Reducing `num_centroids` to {input_height} to match the available number of rows."
+                f"`num_clusters` was set to {num_clusters}, but the input DataFrame only contains {input_height} rows. "
+                f"Reducing `num_clusters` to {input_height} to match the available number of rows."
             )
-        self.num_centroids = min(num_centroids, input_height)
-        self.num_iter = num_iter
+        self.num_clusters = min(num_clusters, input_height)
+        self.max_iter = max_iter
+        self.num_init = num_init
         self.label_column = label_column
         self.centroid_info = centroid_info
 
@@ -47,10 +50,18 @@ def execute(self) -> pl.DataFrame:
         centroids = None
         if not valid_df.is_empty():
             embeddings = np.stack(valid_df[self.embedding_column_name])
-            kmeans = KMeans(k=self.num_centroids, max_iters=self.num_iter)
-            kmeans.fit(embeddings)
-            predicted = kmeans.predict(embeddings).tolist()
-            cluster_centroids = kmeans.centroids.to_numpy(zero_copy_only=False)
+
+            # Using sklearn KMeans with k-means++ initialization (default)
+            kmeans = KMeans(
+                n_clusters=self.num_clusters,
+                max_iter=self.max_iter,
+                init='k-means++',  # This is the default, but being explicit
+                n_init=self.num_init,  # Number of times to run k-means with different centroid seeds
+                random_state=42  # For reproducibility
+            )
+
+            predicted = kmeans.fit_predict(embeddings)
+            cluster_centroids = kmeans.cluster_centers_
 
             if self.centroid_info is not None:
                 centroids = [None] * df.height
@@ -65,8 +76,8 @@ def execute(self) -> pl.DataFrame:
         if self.centroid_info is not None:
             res = res.with_columns(
                 pl.from_arrow(
-                    pa.array(centroids, type=pa.list_(pa.float32(), self.centroid_info[1]))
-                ).alias(self.centroid_info[0])
+                    pa.array(centroids, type=pa.list_(pa.float32(), self.centroid_info.num_dimensions))
+                ).alias(self.centroid_info.centroid_column)
             )
 
         return res

src/fenic/_backends/local/transpiler/plan_converter.py

@@ -264,9 +264,11 @@ def convert(
                 child_physical,
                 physical_by_expr,
                 str(logical.by_expr()),
-                logical.num_clusters(),
-                logical.label_column(),
-                logical.centroid_info(),
+                num_clusters=logical.num_clusters(),
+                max_iter=logical.max_iter(),
+                num_init=logical.num_init(),
+                label_column=logical.label_column(),
+                centroid_info=logical.centroid_info(),
                 cache_info=logical.cache_info,
                 session_state=self.session_state,
             )

src/fenic/api/dataframe/semantic_extensions.py

@@ -35,7 +35,15 @@ def __init__(self, df: DataFrame):
         """
         self._df = df
 
-    def with_cluster_labels(self, by: ColumnOrName, num_clusters: int, label_column: str = "cluster_label", centroid_column: Optional[str] = None) -> DataFrame:
+    def with_cluster_labels(
+        self,
+        by: ColumnOrName,
+        num_clusters: int,
+        max_iter: int = 300,
+        num_init: int = 1,
+        label_column: str = "cluster_label",
+        centroid_column: Optional[str] = None,
+    ) -> DataFrame:
         """Cluster rows using K-means and add cluster metadata columns.
 
         This method clusters rows based on the given embedding column or expression using K-means.
@@ -45,6 +53,8 @@ def with_cluster_labels(self, by: ColumnOrName, num_clusters: int, label_column:
         Args:
             by: Column or expression producing embeddings to cluster (e.g., `embed(col("text"))`).
             num_clusters: Number of clusters to compute (must be > 0).
+            max_iter: Maximum iterations for a single run of the k-means algorithm. The algorithm stops when it either converges or reaches this limit.
+            num_init: Number of independent runs of k-means with different centroid seeds. The best result is selected.
             label_column: Name of the output column for cluster IDs. Default is "cluster_label".
             centroid_column: If provided, adds a column with this name containing the centroid embedding
                             for each row's assigned cluster.
@@ -56,14 +66,16 @@ def with_cluster_labels(self, by: ColumnOrName, num_clusters: int, label_column:
 
         Raises:
             ValidationError: If num_clusters is not a positive integer
+            ValidationError: If max_iter is not a positive integer
+            ValidationError: If num_init is not a positive integer
             ValidationError: If label_column is not a non-empty string
             ValidationError: If centroid_column is not a non-empty string
             TypeMismatchError: If the column is not an EmbeddingType
 
         Example: Basic clustering
             ```python
             # Cluster customer feedback and add cluster metadata
-            clustered_df = df.semantic.with_cluster_labels("feedback_embeddings", 5)
+            clustered_df = df.semantic.with_cluster_labels("feedback_embeddings", num_clusters=5)
 
             # Then use regular operations to analyze clusters
             clustered_df.group_by("cluster_label").agg(count("*"), avg("rating"))
@@ -72,15 +84,23 @@ def with_cluster_labels(self, by: ColumnOrName, num_clusters: int, label_column:
         Example: Filter outliers using centroids
             ```python
             # Cluster and filter out rows far from their centroid
-            clustered_df = df.semantic.with_cluster_labels("embeddings", 3, centroid_column="cluster_centroid")
+            clustered_df = df.semantic.with_cluster_labels("embeddings", num_clusters=3, num_init=10, centroid_column="cluster_centroid")
             clean_df = clustered_df.filter(
                 embedding.compute_similarity("embeddings", "cluster_centroid", metric="cosine") > 0.7
             )
             ```
         """
         # Validate num_clusters
         if not isinstance(num_clusters, int) or num_clusters <= 0:
-            raise ValidationError("`num_clusters` must be a positive integer greater than 0.")
+            raise ValidationError("`num_clusters` must be a positive integer.")
+
+        # Validate max_iter
+        if not isinstance(max_iter, int) or max_iter <= 0:
+            raise ValidationError("`max_iter` must be a positive integer.")
+
+        # Validate num_init
+        if not isinstance(num_init, int) or num_init <= 0:
+            raise ValidationError("`num_init` must be a positive integer.")
 
         # Validate clustering target
         if not isinstance(by, ColumnOrName):
@@ -106,7 +126,13 @@ def with_cluster_labels(self, by: ColumnOrName, num_clusters: int, label_column:
 
         return self._df._from_logical_plan(
             SemanticCluster(
-                self._df._logical_plan, by_expr, num_clusters, label_column, centroid_column
+                self._df._logical_plan,
+                by_expr,
+                num_clusters=num_clusters,
+                max_iter=max_iter,
+                num_init=num_init,
+                label_column=label_column,
+                centroid_column=centroid_column,
             )
         )
 

src/fenic/core/_logical_plan/plans/__init__.py

@@ -18,6 +18,7 @@
 )
 from fenic.core._logical_plan.plans.transform import (
     SQL,
+    CentroidInfo,
     DropDuplicates,
     Explode,
     Filter,
@@ -52,4 +53,5 @@
     "Sort",
     "Union",
     "Unnest",
+    "CentroidInfo",
 ]

src/fenic/core/_logical_plan/plans/transform.py

@@ -2,6 +2,7 @@
 
 import logging
 import re
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple
 
 import duckdb
@@ -415,21 +416,30 @@ def with_children(self, children: List[LogicalPlan]) -> LogicalPlan:
         result.set_cache_info(self.cache_info)
         return result
 
+@dataclass
+class CentroidInfo:
+    centroid_column: str
+    num_dimensions: int
+
 class SemanticCluster(LogicalPlan):
     def __init__(
         self,
         input: LogicalPlan,
         by_expr: LogicalExpr,
         num_clusters: int,
+        max_iter: int,
+        num_init: int,
         label_column: str,
         centroid_column: Optional[str],
     ):
         self._input = input
         self._by_expr = by_expr
         self._num_clusters = num_clusters
+        self._max_iter = max_iter
+        self._num_init = num_init
         self._label_column = label_column
         self._centroid_column = centroid_column
-        self._centroid_info: Optional[Tuple[str, int]] = None
+        self._centroid_info: Optional[CentroidInfo] = None
         super().__init__(self._input.session_state)
 
     def children(self) -> List[LogicalPlan]:
@@ -446,7 +456,7 @@ def _build_schema(self) -> Schema:
         new_fields = [ColumnField(self._label_column, IntegerType)]
         if self._centroid_column:
             new_fields.append(ColumnField(self._centroid_column, by_expr_type))
-            self._centroid_info = (self._centroid_column, by_expr_type.dimensions)
+            self._centroid_info = CentroidInfo(self._centroid_column, by_expr_type.dimensions)
 
         return Schema(column_fields=self._input.schema().column_fields + new_fields)
 
@@ -456,7 +466,13 @@ def _repr(self) -> str:
     def num_clusters(self) -> int:
         return self._num_clusters
 
-    def centroid_info(self) -> Optional[Tuple[str, int]]:
+    def max_iter(self) -> int:
+        return self._max_iter
+
+    def num_init(self) -> int:
+        return self._num_init
+
+    def centroid_info(self) -> Optional[CentroidInfo]:
         return self._centroid_info
 
     def by_expr(self) -> LogicalExpr:
@@ -469,7 +485,13 @@ def with_children(self, children: List[LogicalPlan]) -> LogicalPlan:
         if len(children) != 1:
             raise ValueError("SemanticCluster must have exactly one child")
         result = SemanticCluster(
-            children[0], self._by_expr, self._num_clusters, self._label_column, self._centroid_column
+            children[0],
+            self._by_expr,
+            num_clusters=self._num_clusters,
+            max_iter=self._max_iter,
+            num_init=self._num_init,
+            label_column=self._label_column,
+            centroid_column=self._centroid_column,
         )
         result.set_cache_info(self.cache_info)
         return result