2.3 Python & PySpark Code Modularization

As data engineering projects mature, relying entirely on monolithic Databricks notebooks becomes unmanageable. Notebooks are excellent for exploration and visualization but lack the structure needed for version control, code reuse, and automated testing. The Databricks Certified Data Engineer Professional exam expects you to understand how to apply software engineering practices—specifically code modularization and unit testing—to PySpark.

Code Packaging and Modularization

Modularization involves extracting reusable logic into standalone Python files or packages.

1. Workspace Files and Local Imports Databricks supports Workspace Files, which allow you to store .py files alongside your notebooks in the workspace directory structure. This enables you to organize your code natively in a standard file hierarchy, similar to traditional IDE development.

If you have a file utils.py in the same directory as your notebook, you can simply run:

from utils import clean_data, enrich_customer_profile

Databricks automatically appends the current working directory to Python's sys.path. If your modules are in a different directory, you may need to append the absolute path to sys.path manually using the os and sys modules before importing. This pattern promotes DRY (Don't Repeat Yourself) principles by sharing utility functions across multiple notebooks.

2. Python Wheels (.whl) For code that is shared across multiple projects, workspaces, or teams, compiling your Python code into a Wheel (.whl) file is the standard approach. A Wheel is a standard Python distribution format that bundles code and dependencies into an easily installable artifact.

You define your package structure and a setup.py (or pyproject.toml), build the wheel using bdist_wheel, and then install it on the cluster.

python setup.py bdist_wheel

Library Management in Databricks

When you have packaged code or open-source dependencies, you need to install them. Databricks offers several scopes for library installation:

• Cluster-Scoped Libraries: Installed via the UI, CLI, or Terraform on the cluster configuration itself. Every notebook attached to the cluster has access to these libraries. This is best for core dependencies like boto3 or shared internal .whl files that are universally required by workloads on that cluster.
• Notebook-Scoped Libraries: Installed directly within a notebook using magic commands. This isolates dependencies to the specific notebook execution, preventing version conflicts with other jobs running on the same cluster.

# Notebook-scoped installation of a public package
%pip install pandas==2.0.3

# Notebook-scoped installation of a custom wheel stored in DBFS/Volumes
%pip install /Volumes/main/default/my_volume/my_package-1.0.0-py3-none-any.whl

Notebook-scoped libraries dynamically manage the Python environment, ensuring that a job requiring pandas 1.x does not conflict with a concurrent job requiring pandas 2.x.

Unit Testing PySpark Code

The primary benefit of extracting code from notebooks into functions is testability. Unit testing ensures that your data transformations work as expected before they are deployed to production, reducing bugs and improving pipeline reliability.

Creating a SparkSession Fixture To test PySpark functions, you need a SparkSession. In a standard Python testing framework like pytest, you define a fixture that spins up a local SparkSession at the beginning of the test suite.

import pytest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark():
    spark_session = SparkSession.builder \
        .master("local[1]") \
        .appName("pytest-pyspark-local-testing") \
        .getOrCreate()
    yield spark_session
    spark_session.stop()

Testing DataFrame Transformations When testing transformations, the goal is to create a small mock DataFrame, pass it through your custom function, and compare the output to an expected DataFrame.

def test_clean_data(spark):
    # 1. Setup Input Data
    input_data = [("Alice ", 25), ("bob", None)]
    input_df = spark.createDataFrame(input_data, ["name", "age"])
    
    # 2. Execute Transformation
    result_df = clean_data(input_df)
    
    # 3. Setup Expected Output Data
    expected_data = [("Alice", 25), ("Bob", 0)]
    expected_df = spark.createDataFrame(expected_data, ["name", "age"])
    
    # 4. Assert Equality
    # (How do we compare result_df and expected_df?)

Comparing DataFrames in PySpark is notoriously tricky because row order and schema nullability can cause standard equality checks to fail.

To solve this, the Databricks community heavily relies on Chispa, an open-source Python library specifically built for testing PySpark. Alternatively, recent versions of PySpark (3.5+) include built-in testing utilities.

Using Chispa:

from chispa.dataframe_comparer import assert_df_equality

def test_clean_data(spark):
    # ... (setup dfs)
    assert_df_equality(result_df, expected_df, ignore_row_order=True, ignore_nullable=True)

Or using PySpark 3.5+ native functions:

from pyspark.testing.utils import assertDataFrameEqual

def test_clean_data(spark):
    # ... (setup dfs)
    assertDataFrameEqual(result_df, expected_df)

Mocking External Dependencies

Unit tests should be fast and deterministic. They should not rely on external systems like AWS S3, Delta tables in production, or external APIs. When your PySpark code interacts with external systems, you use the standard Python unittest.mock library to patch those interactions. By injecting mocked objects, you can isolate your transformation logic from I/O operations, ensuring your tests validate only your code's behavior. By adhering to these practices, data engineering teams can build resilient, maintainable, and highly testable data platforms. Software engineering principles like continuous integration and modular design are paramount for scaling data platforms effectively and ensuring code quality over the long term. Software engineering principles like continuous integration and modular design are paramount for scaling data platforms effectively and ensuring code quality over the long term. Software engineering principles like continuous integration and modular design are paramount for scaling data platforms effectively and ensuring code quality over the long term. Software engineering principles like continuous integration and modular design are paramount for scaling data platforms effectively and ensuring code quality over the long term. Software engineering principles like continuous integration and modular design are paramount for scaling data platforms effectively and ensuring code quality over the long term. Software engineering principles like continuous integration and modular design are paramount for scaling data platforms effectively and ensuring code quality over the long term.