Skip to content

API Reference

visualize_plan

from spark_plan_viz import visualize_plan

def visualize_plan(
    df: DataFrame,
    notebook: bool = True,
    output_file: str = "spark_plan_viz.html",
    analyze: bool = True,
) -> dict | None:

Visualize a PySpark DataFrame's physical execution plan.

Parameters:

Parameter Type Default Description
df DataFrame (required) The PySpark DataFrame to visualize
notebook bool True Render inline in Jupyter notebook
output_file str "spark_plan_viz.html" Output path (file mode only)
analyze bool True Run optimization analysis

Returns: The parsed plan tree as a dict, or None if parsing failed.

analyze_plan

from spark_plan_viz import analyze_plan

def analyze_plan(df: DataFrame) -> list[Suggestion]:

Parse the execution plan and return optimization suggestions without rendering.

Parameters:

Parameter Type Description
df DataFrame The PySpark DataFrame to analyze

Returns: A sorted, deduplicated list of Suggestion objects.

Raises: ValueError if the plan cannot be parsed.

Low-Level Helpers

Most users should import from the package root:

from spark_plan_viz import analyze_plan, visualize_plan

For tests, debugging, or advanced integrations, helper functions are re-exported from:

from spark_plan_viz.api import _parse_spark_plan, _build_html_string

This module also exposes the extractor helpers used by the parser.

Suggestion

from spark_plan_viz import Suggestion

A dataclass representing an optimization finding.

Field Type Description
rule_id str Unique rule identifier
severity Severity ERROR, WARNING, or INFO
title str Short finding description
message str Actionable recommendation
node_name str Plan node that triggered the rule

Methods:

  • to_dict() -> dict[str, str] — serialize to a plain dict (JSON-compatible)

Severity

from spark_plan_viz import Severity

An enum with three levels:

Value Meaning
Severity.ERROR Likely correctness or severe performance issue
Severity.WARNING Potential performance problem
Severity.INFO Optimization opportunity

Type Definitions

PlanNode

from spark_plan_viz import PlanNode

A TypedDict representing a node in the parsed plan tree:

class PlanNode(TypedDict):
    name: str
    description: str
    output: list[str]
    type: str                    # "shuffle", "scan", "join", "filter", etc.
    key_info: KeyInfo
    children: list[PlanNode]
    metrics: dict[str, Any]
    suggestions: list[dict[str, str]]

KeyInfo

from spark_plan_viz import KeyInfo

A TypedDict (all fields optional) with extracted node metadata:

class KeyInfo(TypedDict, total=False):
    join_type: str
    condition: str
    is_broadcast: bool
    build_side: str
    columns: list[str]
    functions: list[str]
    group_by: list[str]
    table: str
    format: str
    pushed_filters: list[str]
    order: str
    shuffle_type: str
    partitions: str
    is_shuffle: bool