Capstone Project 2: PySpark Code Reviewer Agent

Data engineers frequently review pull requests containing PySpark pipelines. Common mistakes—such as forgetting to broadcast small dimension tables, creating Cartesian products, or inducing severe partition data skew—can cost thousands of dollars in cloud compute bills.

In this project, you will build an Autonomous PySpark Code Reviewer Agent that analyzes code diffs for data engineering antipatterns.

1. Antipatterns Checked by the Agent

flowchart TD
    Code["PySpark Pull Request Code"] --> Scanner["Code Reviewer Agent Brain"]

    Scanner --> Check1["1. Missing Broadcast Join<br/>(Joining huge fact table with tiny lookup table without broadcast)"]
    Scanner --> Check2["2. Cartesian Product Risk<br/>(Join condition missing or using non-equi join)"]
    Scanner --> Check3["3. Excessive repartition(1)<br/>(Bottlenecking all data onto a single executor)"]
    Scanner --> Check4["4. Uncached Repeated Actions<br/>(Evaluating the same expensive DataFrame 3+ times)"]

    Scanner --> Report["Structured Code Review Audit Report"]

    style Code fill:#334155,color:#fff
    style Scanner fill:#2563eb,color:#fff
    style Report fill:#059669,color:#fff

2. Complete Runnable Implementation

import re
from typing import List, Dict

class PySparkReviewAgent:
    def __init__(self):
        self.rules = [
            {
                "id": "RULE_01_BROADCAST",
                "pattern": r"\.join\(\s*df_([a-zA-Z0-9_]*dim[a-zA-Z0-9_]*),",
                "severity": "WARNING",
                "message": "Potential unbroadcasted lookup table. If this dimension table is under 100MB, wrap with broadcast(df_dim) to avoid an expensive shuffle."
            },
            {
                "id": "RULE_02_REPARTITION_ONE",
                "pattern": r"\.repartition\(1\)",
                "severity": "CRITICAL",
                "message": "Using .repartition(1) forces all cluster data onto a single core, causing executor OutOfMemory errors. Use .coalesce() or reasonable partition counts."
            },
            {
                "id": "RULE_03_COUNT_CHECK",
                "pattern": r"if\s+.*\.count\(\)\s*>\s*0:",
                "severity": "PERFORMANCE",
                "message": "Calling .count() > 0 evaluates the entire DataFrame just to check for existence. Use .isEmpty() or .take(1) instead."
            }
        ]

    def review_code(self, code_snippet: str) -> List[Dict]:
        findings = []
        lines = code_snippet.splitlines()

        for rule in self.rules:
            for idx, line in enumerate(lines, 1):
                if re.search(rule["pattern"], line, re.IGNORECASE):
                    findings.append({
                        "rule_id": rule["id"],
                        "line_number": idx,
                        "line_content": line.strip(),
                        "severity": rule["severity"],
                        "recommendation": rule["message"]
                    })

        return findings

if __name__ == "__main__":
    sample_pyspark_pr = """
def process_nightly_transactions(df_huge_fact, df_dim_status):
    # Check if data exists
    if df_huge_fact.count() > 0:
        # Join fact with dimension
        df_joined = df_huge_fact.join(df_dim_status, df_huge_fact.status_id == df_dim_status.id)

        # Write to single file
        df_joined.repartition(1).write.parquet("s3://bucket/output/")
    """

    agent = PySparkReviewAgent()
    results = agent.review_code(sample_pyspark_pr)

    print("=== PySpark Code Review Audit Report ===")
    for f in results:
        print(f"
[{f['severity']}] Line {f['line_number']}: {f['rule_id']}")
        print(f"Code: {f['line_content']}")
        print(f"Fix : {f['recommendation']}")

Key Takeaways

• Automated AI code review agents protect production data platforms by catching antipatterns before merge.
• Rule-based heuristics combined with LLM reasoning deliver high-speed, verifiable PR comments in CI/CD pipelines.