Polars profiling

lf_profile.py

import polars as pl

lf = pl.LazyFrame({"a": [*range(1_000_000)], "b": [*range(1_000_000)]})

perm = lf.with_columns(pl.col("a").shuffle()).with_columns((pl.col("a") * pl.col("b")).alias("c"))

agg = (perm.sort(by="a").group_by_dynamic("a", every="10000i", period="50000i")
                       .agg([pl.col("b").sum().alias("b_sum"),
                             pl.col("c").std().alias("c_std"),
                             pl.col("c").quantile(0.95).alias("c_p95")
]))

result, profile_df = agg.profile()

# Add duration and percentage calculations
timings = profile_df.with_columns([
    (pl.col("end") - pl.col("start")).alias("duration")
]).with_columns([
    (pl.col("duration") / pl.col("duration").sum() * 100).alias("percent_total")
])

print(timings)

# shape: (5, 5)
# ┌──────────────────┬───────┬─────────┬──────────┬───────────────┐
# │ node             ┆ start ┆ end     ┆ duration ┆ percent_total │
# │ ---              ┆ ---   ┆ ---     ┆ ---      ┆ ---           │
# │ str              ┆ u64   ┆ u64     ┆ u64      ┆ f64           │
# ╞══════════════════╪═══════╪═════════╪══════════╪═══════════════╡
# │ optimization     ┆ 0     ┆ 94      ┆ 94       ┆ 0.008489      │
# │ with_column(a)   ┆ 94    ┆ 8160    ┆ 8066     ┆ 0.728446      │
# │ with_column(c)   ┆ 8170  ┆ 10283   ┆ 2113     ┆ 0.190826      │
# │ sort(a)          ┆ 10287 ┆ 34732   ┆ 24445    ┆ 2.207644      │
# │ group_by_dynami) ┆ 34737 ┆ 1107308 ┆ 1072571  ┆ 96.864595     │
# └──────────────────┴───────┴─────────┴──────────┴───────────────┘