Add UDTF example with improved documentation

lennartkats-db · lennartkats-db · commit ca9d17368265 · 2025-12-19T14:08:47.000+01:00
- Add k-means clustering UDTF example for Unity Catalog
- Focus documentation on UDTF pattern and SQL accessibility
- Include Python implementation and SQL usage examples
- Add CI/CD integration instructions
diff --git a/knowledge_base/udtf-operator/.gitignore b/knowledge_base/udtf-operator/.gitignore
@@ -0,0 +1,10 @@
+.databricks/
+build/
+dist/
+__pycache__/
+*.egg-info
+.venv/
+scratch/**
+!scratch/README.md
+**/explorations/**
+**/!explorations/README.md
diff --git a/knowledge_base/udtf-operator/.vscode/__builtins__.pyi b/knowledge_base/udtf-operator/.vscode/__builtins__.pyi
@@ -0,0 +1,3 @@
+# Typings for Pylance in Visual Studio Code
+# see https://github.com/microsoft/pyright/blob/main/docs/builtins.md
+from databricks.sdk.runtime import *
diff --git a/knowledge_base/udtf-operator/.vscode/extensions.json b/knowledge_base/udtf-operator/.vscode/extensions.json
@@ -0,0 +1,7 @@
+{
+    "recommendations": [
+        "databricks.databricks",
+        "redhat.vscode-yaml",
+        "ms-python.black-formatter"
+    ]
+}
diff --git a/knowledge_base/udtf-operator/.vscode/settings.json b/knowledge_base/udtf-operator/.vscode/settings.json
@@ -0,0 +1,39 @@
+{
+    "jupyter.interactiveWindow.cellMarker.codeRegex": "^# COMMAND ----------|^# Databricks notebook source|^(#\\s*%%|#\\s*\\<codecell\\>|#\\s*In\\[\\d*?\\]|#\\s*In\\[ \\])",
+    "jupyter.interactiveWindow.cellMarker.default": "# COMMAND ----------",
+    "python.testing.pytestArgs": [
+        "."
+    ],
+    "files.exclude": {
+        "**/*.egg-info": true,
+        "**/__pycache__": true,
+        ".pytest_cache": true,
+        "dist": true,
+    },
+    "files.associations": {
+        "**/.gitkeep": "markdown"
+    },
+
+    // Pylance settings (VS Code)
+    // Set typeCheckingMode to "basic" to enable type checking!
+    "python.analysis.typeCheckingMode": "off",
+    "python.analysis.extraPaths": ["src", "lib", "resources"],
+    "python.analysis.diagnosticMode": "workspace",
+    "python.analysis.stubPath": ".vscode",
+
+    // Pyright settings (Cursor)
+    // Set typeCheckingMode to "basic" to enable type checking!
+    "cursorpyright.analysis.typeCheckingMode": "off",
+    "cursorpyright.analysis.extraPaths": ["src", "lib", "resources"],
+    "cursorpyright.analysis.diagnosticMode": "workspace",
+    "cursorpyright.analysis.stubPath": ".vscode",
+
+    // General Python settings
+    "python.defaultInterpreterPath": "./.venv/bin/python",
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true,
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter",
+        "editor.formatOnSave": true,
+    },
+}
diff --git a/knowledge_base/udtf-operator/README.md b/knowledge_base/udtf-operator/README.md
@@ -0,0 +1,103 @@
+# User-Defined Table Functions in Unity Catalog
+
+This project demonstrates how to create and register a Python User-Defined Table Function (UDTF) in Unity Catalog using Databricks Asset Bundles. Once registered, the UDTF becomes available to analysts and other users across your Databricks workspace, callable directly from SQL queries.
+
+**Learn more:** [Introducing Python UDTFs in Unity Catalog](https://www.databricks.com/blog/introducing-python-user-defined-table-functions-udtfs-unity-catalog)
+
+## Concrete example: Definition and Usage
+
+This project includes a k-means clustering algorithm as a UDTF.
+
+### Python Implementation
+
+The UDTF is defined in [`src/kmeans_udtf.py`](src/kmeans_udtf.py) as follows:
+
+```python
+class SklearnKMeans:
+    def __init__(self, id_column: str, columns: list, k: int):
+        self.id_column = id_column
+        self.columns = columns
+        self.k = k
+        self.data = []
+
+    def eval(self, row: Row):
+        # Process each input row
+        self.data.append(row)
+
+    def terminate(self):
+        # Perform computation and yield results
+        # ... clustering logic ...
+        for record in results:
+            yield (record.id, record.cluster)
+```
+
+### SQL Usage
+
+Once registered, any analyst or SQL user in your workspace can call the UDTF from SQL queries:
+
+```sql
+SELECT * FROM main.your_schema.k_means(
+    input_data => TABLE(SELECT * FROM my_data),
+    id_column => 'id',
+    columns => array('feature1', 'feature2', 'feature3'),
+    k => 3
+)
+```
+
+The UDTF integrates seamlessly with:
+- SQL queries in notebooks
+- Databricks SQL dashboards
+- Any tool that connects to your Databricks workspace via SQL
+
+See [`src/sample_notebook.ipynb`](src/sample_notebook.ipynb) for complete examples.
+
+## Getting Started With This Project
+
+### Prerequisites
+
+* Databricks workspace with Unity Catalog enabled
+* Databricks CLI installed and configured
+* Python with `uv` package manager
+
+### Setup and Testing
+
+1. Install dependencies:
+   ```bash
+   uv sync --dev
+   ```
+
+2. Run tests (registers and executes the UDTF):
+   ```bash
+   uv run pytest
+   ```
+
+### Deployment
+
+Deploy to dev:
+```bash
+databricks bundle deploy --target dev
+databricks bundle run register_udtf_job --target dev
+```
+
+Deploy to production:
+```bash
+databricks bundle deploy --target prod
+databricks bundle run register_udtf_job --target prod
+```
+
+The UDTF will be registered at `main.your_username.k_means` (dev) or `main.prod.k_means` (prod).
+
+## Advanced Topics
+
+**CI/CD Integration:**
+- Set up CI/CD for Databricks Asset Bundles following the [CI/CD documentation](https://docs.databricks.com/dev-tools/bundles/ci-cd.html)
+- To automatically register the UDTF on deployment, add `databricks bundle run -t prod register_udtf_job` to your deployment script after `databricks bundle deploy -t prod` (alternatively, the job in `resources/udtf_job.yml` can use a schedule for registration)
+
+**Serverless compute vs. clusters:** The job uses serverless compute by default. Customize catalog/schema in `databricks.yml` or via job parameters.
+
+## Learn More
+
+- [Introducing Python UDTFs in Unity Catalog](https://www.databricks.com/blog/introducing-python-user-defined-table-functions-udtfs-unity-catalog) - Blog post covering UDTF concepts and use cases
+- [Python UDTFs Documentation](https://docs.databricks.com/udf/udtf-unity-catalog.html) - Official documentation
+- [Databricks Asset Bundles](https://docs.databricks.com/dev-tools/bundles/index.html) - CI/CD and deployment framework
+- [Unity Catalog Functions](https://docs.databricks.com/udf/unity-catalog.html) - Governance and sharing
diff --git a/knowledge_base/udtf-operator/databricks.yml b/knowledge_base/udtf-operator/databricks.yml
@@ -0,0 +1,42 @@
+# This is a Databricks asset bundle definition for udtf_operator.
+# See https://docs.databricks.com/dev-tools/bundles/index.html for documentation.
+bundle:
+  name: udtf_operator
+  uuid: 67b02248-6ca2-4230-8ce5-fc7f5e37c7a5
+
+include:
+  - resources/*.yml
+  - resources/*/*.yml
+
+# Variable declarations. These variables are assigned in the dev/prod targets below.
+variables:
+  catalog:
+    description: The catalog to use
+  schema:
+    description: The schema to use
+
+targets:
+  dev:
+    # The default target uses 'mode: development' to create a development copy.
+    # - Deployed resources get prefixed with '[dev my_user_name]'
+    # - Any job schedules and triggers are paused by default.
+    # See also https://docs.databricks.com/dev-tools/bundles/deployment-modes.html.
+    mode: development
+    default: true
+    workspace:
+      host: https://company.databricks.com
+    variables:
+      catalog: main
+      schema: ${workspace.current_user.short_name}
+  prod:
+    mode: production
+    workspace:
+      host: https://company.databricks.com
+      # We explicitly deploy to /Workspace/Users/user@company.com to make sure we only have a single copy.
+      root_path: /Workspace/Users/user@company.com/.bundle/${bundle.name}/${bundle.target}
+    variables:
+      catalog: main
+      schema: prod
+    permissions:
+      - user_name: user@company.com
+        level: CAN_MANAGE
diff --git a/knowledge_base/udtf-operator/fixtures/.gitkeep b/knowledge_base/udtf-operator/fixtures/.gitkeep
@@ -0,0 +1,9 @@
+# Test fixtures directory
+
+Add JSON or CSV files here. In tests, use them with `load_fixture()`:
+
+```
+def test_using_fixture(load_fixture):
+    data = load_fixture("my_data.json")
+    assert len(data) >= 1
+```
diff --git a/knowledge_base/udtf-operator/fixtures/titanic_sample.csv b/knowledge_base/udtf-operator/fixtures/titanic_sample.csv
@@ -0,0 +1,9 @@
+PassengerId,Age,Pclass,Survived
+1,22.0,3,0
+2,38.0,1,1
+3,26.0,3,1
+4,35.0,1,1
+5,35.0,3,0
+6,54.0,1,0
+7,2.0,3,0
+8,27.0,1,0
diff --git a/knowledge_base/udtf-operator/pyproject.toml b/knowledge_base/udtf-operator/pyproject.toml
@@ -0,0 +1,30 @@
+[project]
+name = "udtf_operator"
+version = "0.0.1"
+authors = [{ name = "lennart.kats@databricks.com" }]
+requires-python = ">=3.10,<3.13"
+dependencies = [
+    "scikit-learn>=1.3.0",
+    "numpy>=1.24.0",
+]
+
+[dependency-groups]
+dev = [
+    "pytest",
+    "databricks-dlt",
+    "databricks-connect>=15.4,<15.5",
+    "ipykernel",
+]
+
+[project.scripts]
+main = "udtf_operator.main:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src"]
+
+[tool.black]
+line-length = 125
diff --git a/knowledge_base/udtf-operator/resources/udtf_job.yml b/knowledge_base/udtf-operator/resources/udtf_job.yml
@@ -0,0 +1,26 @@
+resources:
+  jobs:
+    register_udtf_job:
+      name: register_udtf_job
+      schedule:
+        quartz_cron_expression: '0 0 8 * * ?' # daily at 8am
+        timezone_id: "UTC"
+      parameters:
+        - name: catalog
+          default: ${var.catalog}
+        - name: schema
+          default: ${var.schema}
+      tasks:
+        - task_key: create_udtf
+          spark_python_task:
+            python_file: ../src/main.py
+            parameters:
+              - "--catalog"
+              - "{{job.parameters.catalog}}"
+              - "--schema"
+              - "{{job.parameters.schema}}"
+          environment_key: default
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "4"
diff --git a/knowledge_base/udtf-operator/src/kmeans_udtf.py b/knowledge_base/udtf-operator/src/kmeans_udtf.py
@@ -0,0 +1,68 @@
+"""K-Means UDTF for Unity Catalog"""
+
+import inspect
+from pyspark.sql import SparkSession
+
+
+class SklearnKMeans:
+    """K-means clustering UDTF handler"""
+
+    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):
+        """Process each row"""
+        if self.id_col is None:
+            self.id_col = id_column
+            self.feature_cols = columns
+            self.k = max(1, int(k))
+
+        row_dict = row.asDict(recursive=False)
+        self.rows.append(row_dict)
+
+        feats = [float(row_dict.get(c) or 0.0) for c in self.feature_cols]
+        self.features.append(feats)
+
+    def terminate(self):
+        """Emit cluster assignments"""
+        if not self.rows:
+            return
+
+        import numpy as np
+        from sklearn.cluster import KMeans
+
+        X = np.asarray(self.features, dtype=float)
+        n_clusters = min(self.k, X.shape[0])
+
+        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)
+
+
+def register(catalog: str, schema: str, name: str = "k_means"):
+    """Register k_means UDTF in Unity Catalog"""
+    spark = SparkSession.builder.getOrCreate()
+    source = inspect.getsource(SklearnKMeans)
+
+    spark.sql(f"""
+        CREATE OR REPLACE FUNCTION {catalog}.{schema}.{name}(
+            input_data TABLE,
+            id_column STRING,
+            columns ARRAY<STRING>,
+            k INT
+        )
+        RETURNS TABLE (id STRING, cluster_id INT)
+        LANGUAGE PYTHON
+        HANDLER 'SklearnKMeans'
+        AS $$
+{source}
+        $$
+    """)
+
+    return f"{catalog}.{schema}.{name}"
diff --git a/knowledge_base/udtf-operator/src/main.py b/knowledge_base/udtf-operator/src/main.py
@@ -0,0 +1,14 @@
+"""Main entry point for UDTF registration"""
+
+import argparse
+import kmeans_udtf
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--catalog", required=True)
+    parser.add_argument("--schema", required=True)
+    args = parser.parse_args()
+
+    fn = kmeans_udtf.register(args.catalog, args.schema)
+    print(f"✅ Registered {fn}")
diff --git a/knowledge_base/udtf-operator/src/sample_notebook.ipynb b/knowledge_base/udtf-operator/src/sample_notebook.ipynb
diff --git a/knowledge_base/udtf-operator/tests/conftest.py b/knowledge_base/udtf-operator/tests/conftest.py
diff --git a/knowledge_base/udtf-operator/tests/test_kmeans_udtf.py b/knowledge_base/udtf-operator/tests/test_kmeans_udtf.py
diff --git a/knowledge_base/udtf-operator/uv.lock b/knowledge_base/udtf-operator/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# Typings for Pylance in Visual Studio Code`
	`2`	`+# see https://github.com/microsoft/pyright/blob/main/docs/builtins.md`
	`3`	`+from databricks.sdk.runtime import *`