typestar

pl_pipeline.py in Python

A lazy Polars pipeline: scan, clean, window, aggregate, join and write.

from pathlib import Path

import polars as pl

RUNS_CSV = """day,lang,chars,seconds,errors
2026-07-20,python,1820,60.0,4
2026-07-20,rust,1510,60.0,9
2026-07-21,python,1980,60.0,2
2026-07-21,sql,1240,60.0,1
2026-07-22,python,2050,60.0,3
2026-07-22,rust,1660,60.0,5
2026-07-24,css,980,45.0,7
2026-07-29,rust,1740,60.0,2
2026-07-29,,900,60.0,0
"""

TOURS = pl.DataFrame({
    "lang": ["python", "rust", "sql", "css"],
    "tours": [11, 12, 3, 3],
})


def build_plan(path: Path) -> pl.LazyFrame:
    """One plan: clean, derive, rank within language, attach tour counts."""
    return (
        pl.scan_csv(path, try_parse_dates=True)
        .filter(pl.col("lang").is_not_null() & (pl.col("chars") > 0))
        .with_columns(
            ((pl.col("chars") / 5) / (pl.col("seconds") / 60))
            .round(1).alias("wpm"),
            ((1 - pl.col("errors") / pl.col("chars")) * 100)
            .round(2).alias("accuracy"),
        )
        .with_columns(
            pl.col("wpm").mean().over("lang").round(1).alias("lang_mean"),
            pl.col("wpm").rank(descending=True).over("lang")
            .cast(pl.Int32).alias("rank_in_lang"),
            pl.when(pl.col("accuracy") >= 99.5).then(pl.lit("clean"))
            .when(pl.col("accuracy") >= 99.0).then(pl.lit("good"))
            .otherwise(pl.lit("sloppy")).alias("grade"),
        )
        .join(TOURS.lazy(), on="lang", how="left")
        .sort(["lang", "rank_in_lang"])
    )


def summarize(runs: pl.DataFrame) -> pl.DataFrame:
    return (
        runs.group_by("lang")
        .agg(
            pl.len().alias("runs"),
            pl.col("wpm").max().alias("best_wpm"),
            pl.col("wpm").mean().round(1).alias("mean_wpm"),
            pl.col("accuracy").mean().round(2).alias("mean_accuracy"),
            (pl.col("grade") == "clean").sum().alias("clean_runs"),
            pl.col("tours").first().alias("tours"),
        )
        .sort("mean_wpm", descending=True)
    )


def daily(runs: pl.DataFrame) -> pl.DataFrame:
    return (
        runs.sort("day")
        .group_by_dynamic("day", every="3d")
        .agg(pl.col("wpm").mean().round(1).alias("mean_wpm"),
             pl.len().alias("runs"))
    )


def main() -> None:
    source = Path("runs.csv")
    source.write_text(RUNS_CSV)
    plan = build_plan(source)

    print("--- plan ---")
    print(plan.explain(optimized=True).splitlines()[0])

    runs = plan.collect()
    print(f"\n{runs.height} runs kept of 9 rows read")

    print("\n--- per language ---")
    print(summarize(runs))

    print("\n--- every three days ---")
    print(daily(runs))

    best = runs.sort("wpm", descending=True).row(0, named=True)
    print(f"\nbest run: {best['lang']} at {best['wpm']} wpm "
          f"on {best['day']}")

    runs.write_parquet("runs.parquet")
    print("wrote", pl.read_parquet("runs.parquet").shape)


if __name__ == "__main__":
    main()

How it works

  1. Everything is one LazyFrame plan until the single collect at the end.
  2. The optimizer pushes the filters and projections down for you.
  3. The plan is printed first, so you can see what will actually run.

Keywords and builtins used here

The run, in numbers

Lines
97
Characters to type
2458
Tokens
833
Three-star pace
115 tpm

At the three-star pace of 115 tokens a minute, this run takes about 435 seconds.

Type this snippet

Step 2 of 2 in Encore, step 32 of 32 in Polars.

← Previous