Performance analysis and charts for https://plotly.com/blog/chart-smarter-not-harder-universal-dataframe-support/

narwhals_performance_comparison.py

# Script for generating figures following performance_analysis.py
# Authors: @ndrezn @emilykl

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go

def read_and_clean_df(file_path):

    df = pd.read_csv(file_path)

    # In the environment column, replace "venv/final-5.24.1" with "v5.24"
    # and "venv/final-6.0.0rc0" with "v6.0"
    df['Environment'] = df['Environment'].replace(
        {
            "venv/final-5.24.1": "v5.24",
            "venv/final-6.0.0rc0": "v6.0"
        }
    )

    # Add "px." to the beginning of each value in the "Chart type" column
    df['Chart type'] = "px." + df['Chart type']

    # Aggregate by "Chart type", "Dataframe type", and "Environment" columns
    # For each column, calculate the mean, min, and max of the "Time (s)" column
    df_grouped = df.groupby(["Chart type", "Dataframe type", "Environment"]).agg(
        {
            "Time (s)": ["mean", "min", "max"]
        }
    )
    # Flatten everything so it's just a normal table
    df_grouped = df_grouped.reset_index()
    # Flatten the column names as well
    df_grouped.columns = ["Chart type", "Dataframe type", "Environment", "Mean Time (s)", "Min Time (s)", "Max Time (s)"]
    df_grouped["error_y"] = df_grouped["Max Time (s)"] - df_grouped["Mean Time (s)"]
    df_grouped["error_y_minus"] = -(df_grouped["Min Time (s)"] - df_grouped["Mean Time (s)"])
    
    return df_grouped

def make_heatmap(df_grouped):

    pivot_data = df_grouped.pivot_table(
        index=["Chart type", "Dataframe type"], columns="Environment", values="Mean Time (s)"
    ).reset_index()
    pivot_data["Improvement (%)"] = (
        (pivot_data["v5.24"] - pivot_data["v6.0"]) / pivot_data["v5.24"]
    ) * 100
    pivot_data["Improvement Factor"] = (
        pivot_data["v5.24"] / pivot_data["v6.0"]
    )

    # heatmap_data_percentage = pivot_data.pivot(
    #     index="Chart type", columns="Dataframe type", values="Improvement (%)"
    # )
    heatmap_data_factor = pivot_data.pivot(
        index="Chart type", columns="Dataframe type", values="Improvement Factor"
    )

    # Filter to a limited set of chart types (column "Chart type")
    chart_types = ["px.scatter", "px.line", "px.bar", "px.box", "px.histogram"]
    heatmap_data_factor = heatmap_data_factor[heatmap_data_factor.index.isin(chart_types)]

    fig = go.Figure(
        data=go.Heatmap(
            z=heatmap_data_factor.values,
            x=heatmap_data_factor.columns,
            y=heatmap_data_factor.index,
            colorscale = ["#CFCDEB", "#675AFF"],
            # Hide colorbar
            showscale=False,
            text=[[f"{val:.1f}x" for val in row] for row in heatmap_data_factor.values],
            hoverinfo="text",
        )
    )
    fig.update_layout(
        title=dict(
            text="Performance improvement of Plotly v6.0 over v5.24 for large datasets",
            x=0.5,
            subtitle=dict(
                text="Reduction factor in time required to generate a plot object for a 1-million-row dataframe",
                font_color="#888888",
                font_style="italic",
            ),
        ),
        xaxis_title="Dataframe type",
        yaxis_title="Chart type",
    )
    fig.update_traces(
        text=[[f"{val:.1f}x" for val in row] for row in heatmap_data_factor.values],
        texttemplate="%{text}",
    )
    return fig


def make_bar_chart(df_grouped):
    df_for_bar = df_grouped.copy()

    chart_types = ["px.scatter", "px.line", "px.bar", "px.box", "px.histogram"]
    df_for_bar = df_for_bar[df_for_bar["Chart type"].isin(chart_types)]

    fig = px.bar(
        df_for_bar,
        x="Chart type",
        y="Mean Time (s)",
        color="Environment",
        facet_col="Dataframe type",
        barmode="group",
        title="Time to generate Plotly Express charts for 1 million rows",
        color_discrete_map={"v5.24": "#FFA15A", "v6.0": "#675AFF"},
    )

    # Set titles of subplots
    fig.for_each_annotation(lambda a: a.update(text=a.text.split("=")[-1]))

    # Remove x-axis titles
    fig.update_xaxes(title_text="")

    # Set y-axis title, but ONLY for first subplot
    fig.update_yaxes(title_text="Time (seconds)", row=1, col=1)

    # Set legend title
    fig.update_layout(legend_title="Plotly version")

    # Center plot title
    fig.update_layout(title_x=0.5)

    return fig

if __name__ == '__main__':
    
    df_grouped = read_and_clean_df("performance_results_final_ALL_CHARTS.csv")

    fig_heatmap = make_heatmap(df_grouped)
    fig_bar = make_bar_chart(df_grouped)

    fig_heatmap.show()
    fig_bar.show()

performance_analysis.py

# Script for testing time to generate various Ploty Express charts
# Authors: @ndrezn @emilykl

# Note: The script is intended to be run in various virtual environments with
# different sets of dependencies installed; hence the reference to environments
# in the usage instructions. "Environment" can be any string describing the testing scenario

import sys
import time
from functools import wraps
import os

import polars as pl
import pandas as pd
import numpy as np
import plotly.express as px


try:
    import orjson
except ImportError:
    orjson = None

np.random.seed(42)

REPEAT = 21
IGNORE_FIRST = True
TO_JSON = True
NUM_ROWS=1_000_000


# Decorator to time a function multiple times and return the average time
def timeit(repeat=1, ignore_first=False):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            total_times = []
            for i, _ in enumerate(range(repeat)):
                start_time = time.time()  # Record start time
                _ = func(*args, **kwargs)
                end_time = time.time()  # Record end time
                if ignore_first and i == 0: # Ignore the first run if ignore_first is True
                    continue
                total_times.append(end_time - start_time)  # Keep a list of all total times
            return total_times
        return wrapper
    return decorator

# Generate a large dataset as a dict, with additional columns for color and facet
def generate_large_dataset_as_dict(num_rows=1_000_000):
    return {
        "x": np.random.uniform(0, 100, num_rows),
        "y": np.random.uniform(0, 100, num_rows),
        "z": np.random.uniform(0, 100, num_rows),
        "category": np.random.choice(["A", "B", "C"], num_rows),  # Original category
        "colorby": np.random.choice(["Group 1", "Group 2"], num_rows),  # Color by group
        "facetby": np.random.choice(
            ["Region 1", "Region 2"], num_rows
        ),  # Facet by region
    }


# Make a Polars DataFrame from a dict
def make_polars(dataset_as_dict):
    return pl.DataFrame(dataset_as_dict)

# Make a Pandas DataFrame from a dict
def make_pandas(dataset_as_dict):
    return pd.DataFrame(dataset_as_dict)

# Make a PyArrow Table from a dict
def make_pyarrow(dataset_as_dict):
    return pl.DataFrame(dataset_as_dict).to_arrow()

dataset_as_dict = generate_large_dataset_as_dict(num_rows=NUM_ROWS)
pandas_df = make_pandas(dataset_as_dict)
polars_df = make_polars(dataset_as_dict)
pyarrow_table = make_pyarrow(dataset_as_dict)

# Scatter plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_scatter(df, to_json=True):
    fig = px.scatter(
        df,
        x="x",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Scatter Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Scatter 3d plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_scatter_3d(df, to_json=True):
    fig = px.scatter_3d(
        df,
        x="x",
        y="y",
        z="z",
        color="colorby",  # Color by the "colorby" column
        # facet_col="facetby",  # Facet by the "facetby" column 
        # facet is NOT supported for scatter3D plots
        title="Scatter 3D Plot with Color",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Scatter polar plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_scatter_polar(df, to_json=True):
    fig = px.scatter_polar(
        df,
        r="x",
        theta="y",
        color="colorby",  # Color by the "colorby" column
        # facet_col="facetby",  # Facet by the "facetby" column
        # facet is NOT supported for scatter polar plots
        title="Scatter Polar Plot with Color",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Scatter ternary plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_scatter_ternary(df, to_json=True):
    fig = px.scatter_ternary(
        df,
        a="x",
        b="y",
        c="z",
        color="colorby",  # Color by the "colorby" column
        # facet_col="facetby",  # Facet by the "facetby" column
        # facet is NOT supported for scatter ternary plots
        title="Scatter Ternary Plot with Color",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Line plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_line(df, to_json=True):
    fig = px.line(
        df,
        x="x",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Line Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Line 3d plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_line_3d(df, to_json=True):
    fig = px.line_3d(
        df,
        x="x",
        y="y",
        z="z",
        color="colorby",  # Color by the "colorby" column
        # facet_col="facetby",  # Facet by the "facetby" column
        # facet is NOT supported for line 3D plots
        title="Line 3D Plot with Color",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Area plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_area(df, to_json=True):
    fig = px.area(
        df,
        x="x",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Area Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")


# Bar plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_bar(df, to_json=True):
    fig = px.bar(
        df,
        x="category",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Bar Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Bar polar plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_bar_polar(df, to_json=True):
    fig = px.bar_polar(
        df,
        r="y",
        theta="category",
        color="colorby",  # Color by the "colorby" column
        # facet_col="facetby",  # Facet by the "facetby" column
        # facet is NOT supported for bar polar plots
        title="Bar Polar Plot with Color",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Violin plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_violin(df, to_json=True):
    fig = px.violin(
        df,
        x="category",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Violin Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Box plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_box(df, to_json=True):
    fig = px.box(
        df,
        x="category",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Box Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# ECDF plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_ecdf(df, to_json=True):
    fig = px.ecdf(
        df,
        x="x",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="ECDF Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Strip plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_strip(df, to_json=True):
    fig = px.strip(
        df,
        x="category",
        y="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Strip Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Histogram plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_histogram(df, to_json=True):
    fig = px.histogram(
        df,
        x="x",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Histogram Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Pie chart
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_pie(df, to_json=True):
    fig = px.pie(
        df,
        names="category",
        values="y",
        color="colorby",  # Color by the "colorby" column
        facet_col="facetby",  # Facet by the "facetby" column
        title="Pie Chart with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")

# Funnel plot
@timeit(repeat=REPEAT, ignore_first=IGNORE_FIRST)
def figure_generation_funnel(df, to_json=True):
    fig = px.funnel(
        df,
        x="y",
        color="category",
        facet_col="facetby",  # Facet by the "facetby" column
        title="Funnel Plot with Color and Facet",
    )
    if to_json:
        if orjson:
            fig.to_json(engine="orjson")
        else:
            fig.to_json(engine="json")


# Function to test all charts and gather the times for Polars, Pandas, and PyArrow
def test_all_charts():

    # List all functions to test
    # Comment out to skip
    # Note: This list does not include ALL Plotly Express chart types for various reasons
    # The full list can be found here: https://plotly.com/python-api-reference/plotly.express.html
    
    names_funcs = [
        ("scatter", figure_generation_scatter),
        ("scatter_3d", figure_generation_scatter_3d),
        ("scatter_polar", figure_generation_scatter_polar),
        ("scatter_ternary", figure_generation_scatter_ternary),
        ("line", figure_generation_line),
        ("line_3d", figure_generation_line_3d),
        ("area", figure_generation_area),
        ("bar", figure_generation_bar),
        ("bar_polar", figure_generation_bar_polar),
        ("violin", figure_generation_violin),
        ("box", figure_generation_box),
        ("ecdf", figure_generation_ecdf),
        ("strip", figure_generation_strip),
        ("histogram", figure_generation_histogram),
        ("pie", figure_generation_pie),
        # ("funnel", figure_generation_funnel),
        # There is a BUG with the above funnel plot code which shows up
        # ONLY on the Narwhals branch -- need to investigate further
        # Skipping for now
    ]

    results = []

    # Polars performance
    print("Testing Polars...")
    for name, func in names_funcs:
        results.extend((name, "Polars", n) for n in func(polars_df, to_json=TO_JSON))

    # Pandas performance
    print("Testing Pandas...")
    for name, func in names_funcs:
        results.extend((name, "Pandas", n) for n in func(pandas_df, to_json=TO_JSON))

    # PyArrow performance
    print("Testing PyArrow...")
    for name, func in names_funcs:
        results.extend((name, "PyArrow", n) for n in func(pyarrow_table, to_json=TO_JSON))

    return results


# Run the performance tests and append results to an existing CSV
def run_and_save_results(csv_filename, environment):
    env_results = test_all_charts()

    # Create DataFrame with separate columns for Chart Type and Dataframe Type
    df = pl.DataFrame(
        {
            "Chart type": [result[0] for result in env_results],
            "Dataframe type": [result[1] for result in env_results],
            "Time (s)": [result[2] for result in env_results],
            "Environment": [environment] * len(env_results),
        }
    )

    # Check if file exists, if not, write CSV, otherwise append
    if not os.path.exists(csv_filename):
        df.write_csv(csv_filename)
        print(f"Results saved to {csv_filename}")
    else:
        # Read existing CSV and concatenate new data
        existing_df = pl.read_csv(csv_filename)
        combined_df = pl.concat([existing_df, df], how="vertical")
        combined_df.write_csv(csv_filename)
        print(f"Results appended to {csv_filename}")


# Entry point for running the tests
if __name__ == "__main__":

    if len(sys.argv) != 3:
        print("Usage: python performance_test.py <csv_filename> <environment-name>")
    else:
        csv_filename = sys.argv[1]
        environment = sys.argv[2]
        run_and_save_results(csv_filename, environment)

performance_results_final_ALL_CHARTS.csv