pyspark-expert

name	pyspark-expert
description	Write production-quality PySpark code, unit tests, and project layouts for Spark 3.5+ with Delta Lake. Use this skill whenever the user asks to write, review, refactor, or debug PySpark code, create pytest-based unit tests for Spark transformations, define schemas, design ETL modules, or discuss DataFrame API best practices. Also trigger when the user mentions PySpark, Spark DataFrame, SparkSession, pyspark.sql.functions, Delta Lake tables, or asks for help structuring a Spark project. Even if the request seems simple (e.g., "write a groupBy"), consult this skill because it enforces code style, anti-pattern avoidance, and idiomatic Spark 3.5+ usage that Claude would otherwise miss.
allowed-tools	Read, Write, Edit, Bash, Grep
license	GPL v3.0
model	sonnet
author	Leandro Kellermann de Oliveira <[e-mail]>

PySpark Expert

This skill produces production-quality PySpark code targeting Spark 3.5+ with Delta Lake. It enforces a consistent project layout, idiomatic DataFrame API usage, DDL-based schema definitions, and pytest patterns with separated test data files.

Before writing any PySpark code, read references/PATTERNS.md for detailed examples, anti-patterns, and the full project layout specification.

---

Core Principles

1. Functional composition over mutation chains — prefer df.select(...) and df.transform(fn) over long .withColumn() chains.
2. Schemas are contracts — define every schema as a DDL string constant, never inline StructType builders. The same schema constants are used in production code and in test data fixtures.
3. Built-ins over UDFs — exhaust pyspark.sql.functions and pyspark.sql.Column methods before reaching for a UDF. When a UDF is genuinely needed, explain why no built-in alternative exists.
4. Explicit over implicit — always alias derived columns, qualify ambiguous column references in joins, and specify column lists in select() rather than select("*").
5. Test with real Spark — no mocking of DataFrame internals. Use a shared SparkSession pytest fixture and assertDataFrameEqual which validates both schema and data in one assertion.

---

Project Layout (enforced)

When generating a PySpark project or module, always use this structure:

project_name/
├── src/
│   ├── __init__.py
│   ├── schemas.py             # DDL schema constants (single source of truth)
│   ├── transformations.py     # Pure DataFrame → DataFrame functions
│   ├── readers.py             # I/O: read from sources
│   ├── writers.py             # I/O: write to sinks
│   └── tasks/                 # Entry points that compose read → transform → write
│       ├── __init__.py
│       └── process_orders.py
├── tests/
│   └── unit_tests/
│       ├── conftest.py                # Shared SparkSession fixture
│       ├── transformations_data.py    # Mock inputs & expected outputs as fixtures
│       └── transformations_test.py    # Flat pytest functions
├── pyproject.toml
└── README.md

The separation matters because:

• schemas.py is the single source of truth — imported by both production code and test data fixtures. Never redefine schemas in tests.
• transformations.py contains pure functions (DataFrame) → DataFrame with no I/O, which makes them trivially testable.
• readers.py / writers.py isolate I/O so tests never touch real storage.
• tasks/ are thin orchestrators; they hold no business logic.
• _data.py / _test.py pairs keep mock data separate from test logic.

For larger projects, schemas.py and transformations.py can become packages (schemas/, transformations/) with submodules. The test file naming convention still applies: one <module>_data.py and one <module>_test.py per source module.

---

Schema Definition Rules

Always define schemas as DDL string constants in src/schemas.py:

# src/schemas.py

ORDERS_SCHEMA = """
    order_id BIGINT,
    customer_id BIGINT,
    order_date DATE,
    amount DECIMAL(12,2),
    status STRING,
    items ARRAY<STRUCT<
        product_id: BIGINT,
        quantity: INT,
        unit_price: DECIMAL(10,2)
    >>
"""

CUSTOMER_SCHEMA = """
    customer_id BIGINT,
    first_name STRING,
    last_name STRING,
    gender STRING,
    country STRING,
    continent STRING
"""

Use them everywhere consistently — in production code, in test data fixtures, and in CREATE TABLE statements:

# In production code
df = spark.read.schema(ORDERS_SCHEMA).parquet(path)

# In test data fixtures
df = spark.createDataFrame(data, ORDERS_SCHEMA)

Why DDL strings instead of StructType chains:

• Readable at a glance — no nested constructor noise.
• Diffable in code reviews.
• Directly usable in spark.createDataFrame(), CREATE TABLE statements, and schema validation.

---

Code Style Rules

Read references/PATTERNS.md for the full anti-pattern catalog with examples. Here is the summary:

Prefer select() + transform() over .withColumn() chains

# ✗ BAD — each withColumn triggers a new projection in the logical plan
df = df.withColumn("full_name", F.concat("first_name", F.lit(" "), "last_name"))
df = df.withColumn("age_group", F.when(F.col("age") < 18, "minor").otherwise("adult"))
df = df.withColumn("created_at", F.current_timestamp())

# ✓ GOOD — single select, one projection
df = df.select(
    "*",
    F.concat("first_name", F.lit(" "), "last_name").alias("full_name"),
    F.when(F.col("age") < 18, "minor").otherwise("adult").alias("age_group"),
    F.current_timestamp().alias("created_at"),
)

Use DataFrame.transform() for pipeline composition

def enrich_with_age_group(df: DataFrame) -> DataFrame:
    return df.select(
        "*",
        F.when(F.col("age") < 18, "minor").otherwise("adult").alias("age_group"),
    )

def add_full_name(df: DataFrame) -> DataFrame:
    return df.select(
        "*",
        F.concat_ws(" ", "first_name", "last_name").alias("full_name"),
    )

result = (
    raw_df
    .transform(add_full_name)
    .transform(enrich_with_age_group)
)

Never use collect() unless you genuinely need all data on the driver

Always qualify join columns to avoid ambiguity

# ✗ BAD
joined = orders.join(customers, "customer_id")

# ✓ GOOD — explicit aliases prevent ambiguity downstream
joined = (
    orders.alias("o")
    .join(customers.alias("c"), F.col("o.customer_id") == F.col("c.customer_id"))
    .select("o.*", "c.name", "c.email")
)

Import style

from pyspark.sql import DataFrame, SparkSession
import pyspark.sql.functions as F
import pyspark.sql.types as T

Always use the F. and T. namespace. Never do from pyspark.sql.functions import *.

---

TDD workflow (enforced)

When developing a new transformation, use a Python virtual environment and always follow this order:

1. Define the mock input and expected output fixtures in _data.py.
2. Write the test function in _test.py — it will fail because the transformation doesn't exist yet.
3. Implement the transformation in src/transformations.py until the test passes.
4. Add edge cases (nulls, empty DataFrames, constraint violations) as new fixture + test pairs.

Use @pytest.mark.skip(reason="TBD") to stub tests you plan to implement later or if you failed to implement more than 3 times. In case of failure, write a TODO.md file in the project root folder.

---

Spark 3.5+ Features to Prefer

When writing code for Spark 3.5+, use these modern APIs:

Instead of	Use (3.5+)
Manual try/except on casts	F.try_to_number(), F.try_cast()
.withColumn() chains	.select() + .transform()
collect_list without ordering	collect_list + .orderBy() in window
Multiple when/otherwise for ranges	F.greatest(), F.least() when applicable
Manual null-safe comparisons	Column.eqNullSafe() (<=>)
Row-level Python UDFs	pandas_udf for transforms, foreachPartition for side effects

---

Minimal documentation:

Always consider the following README.md content when working on a project:

• About: what is the project about.
• Getting Started: commands to install, configure, run unit tests with pytest, and how is expected to run.
• IMPROVEMENTS.md: file summarizing main improvements of the new refactored version in comparison with the original one.

Quick Reference: When to Read references/PATTERNS.md

Always read references/PATTERNS.md when:

• Writing a new transformation function
• Reviewing or refactoring existing PySpark code
• Creating test fixtures with complex types (structs, arrays, maps)
• Deciding between UDF vs built-in approach
• Structuring a new PySpark project from scratch

Quick Reference: When to Read references/PYTEST_PATTERNS.md

Always read references/PYTEST_PATTERNS.md when:

• Writing a new transformation function, unless explicitly told not to do it.
• Refactoring existing PySpark code.
• Creating text fixtures with complex types.
• Structuring a new PySpark project from scratch.