Stop copy-pasting notebook code: package it as a Python wheel and ship it to Databricks

Here's how it usually goes. The same clean_columns() function gets copied into 40 notebooks. Someone fixes a bug in three of them, and now nobody knows which version a job runs.

A wheel (.whl) fixes this. It's a versioned, installable Python package that you test once, publish once and install everywhere.

#When a wheel is the right tool

#Step 1: project layout

sales_utils/
├── pyproject.toml
├── README.md
├── src/
│   └── sales_utils/
│       ├── __init__.py
│       ├── cleaning.py
│       └── jobs.py
└── tests/
    └── test_cleaning.py

The src/ layout stops tests from importing the local folder by accident. You always test the installed package.

# src/sales_utils/cleaning.py
from pyspark.sql import DataFrame, functions as F

def clean_columns(df: DataFrame) -> DataFrame:
    """Lower-case, snake_case column names and trim string columns."""
    for c in df.columns:
        df = df.withColumnRenamed(c, c.strip().lower().replace(" ", "_"))
    for c, t in df.dtypes:
        if t == "string":
            df = df.withColumn(c, F.trim(F.col(c)))
    return df

# src/sales_utils/jobs.py
import argparse
from pyspark.sql import SparkSession
from sales_utils.cleaning import clean_columns

def run_daily():
    p = argparse.ArgumentParser()
    p.add_argument("--source", required=True)
    p.add_argument("--target", required=True)
    args = p.parse_args()
    spark = SparkSession.builder.getOrCreate()
    clean_columns(spark.read.table(args.source)).write.mode("overwrite").saveAsTable(args.target)

#Step 2: pyproject.toml

[build-system]
requires = ["setuptools>=68", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "sales_utils"
version = "1.2.0"
description = "Shared cleaning and job helpers for the sales lakehouse"
requires-python = ">=3.10"
dependencies = []            # pyspark is provided by the Databricks runtime; don't pin it here

[project.optional-dependencies]
dev = ["pytest", "pyspark==3.5.*", "databricks-connect"]

[project.scripts]
run_daily = "sales_utils.jobs:run_daily"     # entry point used by the wheel job task

[tool.setuptools.packages.find]
where = ["src"]

Warning

Don't list pyspark, delta-spark or other runtime-provided libraries as dependencies. Installing a different version on the cluster can break the runtime. Put them under dev for local testing only.

#Step 3: test and build locally

python -m venv .venv && source .venv/bin/activate      # Windows: .venv\Scripts\activate
pip install -e ".[dev]" build
pytest -q
python -m build --wheel
ls dist/
# sales_utils-1.2.0-py3-none-any.whl

py3-none-any means pure Python and installable on any Databricks runtime with a matching Python. Packages with compiled extensions need a wheel built for the runtime's Linux platform and Python version.

#Step 4: put the wheel where Databricks can install it

#Option A: Unity Catalog volume (simplest)

databricks fs cp dist/sales_utils-1.2.0-py3-none-any.whl dbfs:/Volumes/shared/libs/wheels/ --overwrite

GRANT READ VOLUME ON VOLUME shared.libs.wheels TO `data-engineers`;

#Option B: a private package feed (Azure Artifacts, JFrog, Nexus)

pip install twine
twine upload --repository-url https://pkgs.dev.azure.com/contoso/_packaging/data-libs/pypi/upload/ dist/*

Clusters then install with --index-url (or --extra-index-url) pointing at the feed. Keep the feed credential in a secret scope, not in notebook text. A feed is better when many workspaces or teams consume the same libraries.

#Step 5: install and use it

# notebook-scoped (classic or serverless notebooks)
%pip install /Volumes/shared/libs/wheels/sales_utils-1.2.0-py3-none-any.whl
dbutils.library.restartPython()

from sales_utils.cleaning import clean_columns
display(clean_columns(spark.read.table("main.bronze.orders")))

Other ways to install it:

• Cluster library: Compute → Libraries → Install new → Volumes. Every notebook on the cluster gets it.
• Serverless jobs and notebooks: add the wheel path to the environment dependencies.
• Allow-list: on shared-access clusters, admins must allow-list the volume path for libraries.

#Step 6: run it as a wheel task

A Python wheel task calls the entry point directly, with no notebook:

{
  "task_key": "daily_clean",
  "python_wheel_task": {
    "package_name": "sales_utils",
    "entry_point": "run_daily",
    "parameters": ["--source", "main.bronze.orders", "--target", "main.silver.orders"]
  },
  "libraries": [{ "whl": "/Volumes/shared/libs/wheels/sales_utils-1.2.0-py3-none-any.whl" }]
}

#Step 7: automate with Databricks Asset Bundles

A bundle builds the wheel, uploads it and deploys the job in one command:

# databricks.yml
bundle:
  name: sales_utils

artifacts:
  default:
    type: whl
    build: python -m build --wheel
    path: .

resources:
  jobs:
    daily_clean:
      name: sales-daily-clean
      tasks:
        - task_key: daily_clean
          python_wheel_task:
            package_name: sales_utils
            entry_point: run_daily
            parameters: ["--source", "main.bronze.orders", "--target", "main.silver.orders"]
          libraries:
            - whl: ./dist/*.whl
          environment_key: default
      environments:
        - environment_key: default
          spec:
            client: "2"
            dependencies:
              - ./dist/*.whl

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production
    workspace:
      host: https://adb-1234567890123456.7.azuredatabricks.net

databricks bundle validate
databricks bundle deploy -t dev
databricks bundle run daily_clean -t dev
databricks bundle deploy -t prod          # from CI, as a service principal

databricks bundle init default-python generates a working template with this structure. Start from it rather than writing YAML by hand. Serverless environment fields change between CLI versions, so check the generated template for your version.

#Versioning and release rules

• Use semantic versioning: patch for fixes, minor for new functions, major for breaking changes.
• Never overwrite a published version. Jobs pinned to 1.2.0 must keep getting the same code.
• Pin versions in production jobs. Use "latest" only in dev.
• Tag the Git commit (v1.2.0) and build from the tag in CI.
• Keep a CHANGELOG.md so consumers know what changed.

#Troubleshooting

Tip

Test the business logic with plain unit tests and small DataFrames, using local PySpark or Databricks Connect. Keep entry points thin: parse arguments, call tested functions, and write the output.