Tutorial: K-means clustering

In this tutorial, you build a Python user-defined table function (UDTF) operator for Lakeflow Designer that runs K-means clustering with scikit-learn. UDTFs are well-suited to machine learning tasks that process entire datasets. For background on user-defined operators, see User-defined operators in Lakeflow Designer.

Overview

This tutorial steps you through creating a UDTF user-defined operator using Python. The operator performs K-Means clustering on selected columns, allowing users to:

• Choose which columns to use as features.
• Specify the number of clusters.
• Get back a table with cluster assignments for each row.

Step 1: Understand the UDTF handler pattern

A UDTF is implemented as a Python class with three key methods:

• __init__(): Called once before any rows are processed to initialize state.
• eval(row, ...): Called for each input row to accumulate data.
• terminate(): Called after all rows are processed to yield results.

This pattern allows the UDTF to:

1. Collect all data points during eval() calls
2. Train the K-Means model in terminate()
3. Yield clustered results row by row

The __init__ method runs once before any rows are processed. It sets up the state that persists across eval() calls: the configuration (ID column, feature columns, cluster count) and the accumulators for row data and features.

class SklearnKMeans:
    def __init__(self):
        # State set here persists across eval() calls for use in terminate()
        self.id_col = None
        self.feature_cols = None
        self.k = None
        self.rows = []
        self.features = []

eval() is called once per input row. It captures the configuration from the first row, then accumulates each row and its numeric feature values.

    def eval(self, row, id_column, columns, k):
        if self.id_col is None:
            self.id_col = id_column
        if self.feature_cols is None:
            self.feature_cols = columns
        if self.k is None:
            self.k = max(1, int(k))

        row_dict = row.asDict(recursive=False)
        self.rows.append(row_dict)

        feats = []
        for c in self.feature_cols:
            v = row_dict.get(c)
            if v is None:
                v = 0.0
            feats.append(float(v))
        self.features.append(feats)

terminate() runs after all rows are accumulated. It trains the K-Means model on the collected features and yields one output row per input row, pairing each row's ID with its assigned cluster.

    def terminate(self):
        import numpy as np
        from sklearn.cluster import KMeans

        if not self.rows:
            return

        X = np.asarray(self.features, dtype=float)
        n_samples = X.shape[0]
        n_clusters = min(self.k, n_samples)

        model = KMeans(
            n_clusters=n_clusters,
            n_init=10,
            random_state=42
        )
        labels = model.fit_predict(X)

        for row_dict, label in zip(self.rows, labels):
            yield str(row_dict[self.id_col]), int(label)

Step 2: Create the YAML for the operator

The YAML configuration defines how the operator appears in Lakeflow Designer. For this operator:

• Number parameter (k): Number of clusters to create
• Select widget (id_column): Dropdown populated with columns from input table
• Multi-select widget (columns): Multiple feature columns selection
• optionsSource: Automatically populates dropdowns from input table schema
• Input port: Specifies this operator accepts tabular data

schema: user-defined-operator-v0.1.0
type: uc-udtf
name: K-Means Clustering
id: kmeans
version: '1.0.0'
description: Perform K-Means clustering on selected columns
config:
  type: object
  properties:
    k:
      type: number
      title: Number of Clusters
      default: 3
      minimum: 1
      maximum: 100
      x-ui:
        widget: number
    id_column:
      type: string
      title: ID Column
      x-ui:
        widget: select
        optionsSource:
          type: inputColumns
          port: input_data
    columns:
      type: array
      items:
        type: string
      title: Feature Columns
      x-ui:
        widget: multi-select
        optionsSource:
          type: inputColumns
          port: input_data
  required:
    - k
    - id_column
    - columns
  additionalProperties: false
ports:
  input:
    - name: input_data
      title: Input Data
  output:
    - name: output
      title: Clustered Data

See User-defined operator YAML reference for a comprehensive guide to all available properties, data types, widgets, and options.

Step 3: Create the Unity Catalog function

Combine the YAML configuration and Python handler class into a single CREATE FUNCTION statement.

The example creates the function in main.example_output. Create the schema first if it doesn't exist:

CREATE SCHEMA IF NOT EXISTS main.example_output

CREATE OR REPLACE FUNCTION main.example_output.k_means(
    input_data TABLE,
    id_column STRING,
    columns ARRAY<STRING>,
    k INT
)
RETURNS TABLE (
    id STRING,
    cluster_id INT
)
LANGUAGE PYTHON
HANDLER 'SklearnKMeans'
AS $$
"""
schema: user-defined-operator-v0.1.0
type: uc-udtf
name: K-Means Clustering
id: kmeans
version: "1.0.0"
description: Perform K-Means clustering on selected columns
config:
  type: object
  properties:
    k:
      type: number
      title: Number of Clusters
      default: 3
      minimum: 1
      maximum: 100
      x-ui:
        widget: number
    id_column:
      type: string
      title: ID Column
      x-ui:
        widget: select
        optionsSource:
          type: inputColumns
          port: input_data
    columns:
      type: array
      items:
        type: string
      title: Feature Columns
      x-ui:
        widget: multi-select
        optionsSource:
          type: inputColumns
          port: input_data
  required:
    - k
    - id_column
    - columns
  additionalProperties: false
ports:
  input:
    - name: input_data
      title: Input Data
  output:
    - name: output
      title: Clustered Data
"""

class SklearnKMeans:
    def __init__(self):
        self.id_col = None
        self.feature_cols = None
        self.k = None
        self.rows = []
        self.features = []

    def eval(self, row, id_column, columns, k):
        if self.id_col is None:
            self.id_col = id_column
        if self.feature_cols is None:
            self.feature_cols = columns
        if self.k is None:
            self.k = max(1, int(k))

        row_dict = row.asDict(recursive=False)
        self.rows.append(row_dict)

        feats = []
        for c in self.feature_cols:
            v = row_dict.get(c)
            if v is None:
                v = 0.0
            feats.append(float(v))
        self.features.append(feats)

    def terminate(self):
        import numpy as np
        from sklearn.cluster import KMeans

        if not self.rows:
            return

        X = np.asarray(self.features, dtype=float)
        n_samples = X.shape[0]
        n_clusters = min(self.k, n_samples)

        model = KMeans(
            n_clusters=n_clusters,
            n_init=10,
            random_state=42
        )
        labels = model.fit_predict(X)

        for row_dict, label in zip(self.rows, labels):
            yield str(row_dict[self.id_col]), int(label)
$$

Step 4: Test with sample data

Create sample customer data for testing:

CREATE OR REPLACE TEMP VIEW customers AS
SELECT * FROM VALUES
    ('C001', 25, 35000, 20),
    ('C002', 45, 85000, 80),
    ('C003', 35, 55000, 50),
    ('C004', 50, 95000, 90),
    ('C005', 23, 30000, 15),
    ('C006', 40, 75000, 70),
    ('C007', 60, 100000, 95),
    ('C008', 30, 45000, 40)
AS t(customer_id, age, annual_income, spending_score);

Test the K-Means UDTF with three clusters:

SELECT * FROM main.example_output.k_means(
    input_data => TABLE(SELECT * FROM customers) WITH SINGLE PARTITION,
    k => 3,
    id_column => 'customer_id',
    columns => array('age', 'annual_income', 'spending_score')
)

In this case, you want to join the clustering results back with original data to see cluster assignments:

SELECT
  c.*,
  k.cluster_id
FROM customers c
INNER JOIN main.example_output.k_means(
    input_data => TABLE(SELECT * FROM customers) WITH SINGLE PARTITION,
    k => 3,
    id_column => 'customer_id',
    columns => array('age', 'annual_income', 'spending_score')
) k
ON c.customer_id = k.id
ORDER BY k.cluster_id, c.customer_id

Step 5: Register the operator

To use the operator in Lakeflow Designer, you must register it, by adding it to your .user_defined_operators.yaml file:

operators:
  - catalog: main
    schema: example_output
    functionName: k_means

Step 6: Set up permissions

Grant access to users who need to use this operator:

GRANT USE SCHEMA ON SCHEMA main.example_output TO `<user>`;
GRANT EXECUTE ON FUNCTION main.example_output.k_means TO `<user>`;

Use the operator in Lakeflow Designer

After it's registered, the operator appears in Lakeflow Designer with:

• An input port to connect your data source
• A dropdown to select which column uniquely identifies rows
• A multi-select to choose which columns to use as clustering features
• A number input for the desired number of clusters

Users can segment customers, products, or any other data into meaningful groups without writing code.

Tips for building UDTFs

1. Initialize state in __init__: Set up empty lists/variables to accumulate data
2. Accumulate in eval: Don't process yet, just collect data
3. Process in terminate: This is where the real work happens
4. Use yield to return rows: Return results one at a time from terminate
5. Handle edge cases: What if there are fewer rows than clusters?
6. Keep types explicit: UDTF returns cannot reference input types