Cleaning and Shaping Data for a Causal Question

2.1 Polars expressions: select, filter, with_columns (60 min)

Concept. Almost everything in Polars is built from expressions: pl.col("unit_price") > 200, pl.col("quantity") * pl.col("unit_price"). You pass expressions to four verbs:

Examples. Create day02/expressions.py and run it after each step:

# day02/expressions.py
import polars as pl

orders = pl.read_csv("data/raw/orders.csv", try_parse_dates=True)

# 1. select: keep some columns
print(orders.select("order_id", "customer_id", "unit_price").head(3))

# 2. filter: keep some rows
big = orders.filter(pl.col("unit_price") > 200)
print("Lines over $200:", big.height)

# 3. with_columns: add a new column built from others
orders = orders.with_columns(
    revenue=pl.col("quantity") * pl.col("unit_price"),
)
print(orders.select("quantity", "unit_price", "revenue").head(3))

# 4. Several conditions at once
discounted_big = orders.filter(
    (pl.col("discount_pct") > 0) & (pl.col("unit_price") > 100)
)
print("Discounted lines over $100:", discounted_big.height)

# 5. Conditional column: when / then / otherwise
orders = orders.with_columns(
    period=pl.when(pl.col("order_date") >= pl.date(2026, 3, 2))
             .then(pl.lit("post"))
             .otherwise(pl.lit("pre"))
)
print(orders.group_by("period").len().sort("period"))

Lines over $200: 2044
...
Discounted lines over $100: 1554
...
│ post   ┆ 9020 │
│ pre    ┆ 9522 │

Notes:

• Combine conditions with & (and) and | (or), each condition in brackets.
• pl.when(...).then(...).otherwise(...) is the Polars version of if/else for a whole column, like gen ... if ... / replace in Stata.
• pl.date(2026, 3, 2) is a date literal. The email date splits the sample into pre and post.

Try it. Starting from the raw orders, keep only completed orders, add revenue and a Boolean column discounted (discount above 0), keep only order_id, revenue, discounted, and print the share of discounted lines.

import polars as pl
orders = pl.read_csv("data/raw/orders.csv", try_parse_dates=True)
out = (orders
       .filter(pl.col("status") == "completed")
       .with_columns(
           revenue=pl.col("quantity") * pl.col("unit_price"),
           discounted=pl.col("discount_pct") > 0,
       )
       .select("order_id", "revenue", "discounted"))
print(out.head(3))
print("Share discounted:", round(out["discounted"].mean(), 3))

2.2 From raw to clean: duplicates, text, cancellations (50 min)

Concept. Yesterday's dictionary found three problems. Each one would bias a causal estimate if we ignored it:

We write cleaning as functions in a script, never as manual edits to files, so anyone can rerun it. And we finish with assertions: checks that stop the script loudly if a rule is broken.

Create day02/clean.py:

# day02/clean.py
"""Clean the raw NorthMart files and save them as Parquet in data/clean/."""
from pathlib import Path

import polars as pl

RAW = Path("data/raw")
CLEAN = Path("data/clean")


def clean_customers(df: pl.DataFrame) -> pl.DataFrame:
    return df.with_columns(
        region=pl.col("region").str.strip_chars().str.to_uppercase().fill_null("UNKNOWN")
    )


def clean_orders(df: pl.DataFrame) -> pl.DataFrame:
    return (
        df.unique(maintain_order=True)                 # drop exact duplicate rows
          .filter(pl.col("status") == "completed")     # cancelled orders never happened
          .with_columns(revenue=(pl.col("quantity") * pl.col("unit_price")).round(2))
    )


def main() -> None:
    CLEAN.mkdir(parents=True, exist_ok=True)

    customers_raw = pl.read_csv(RAW / "customers.csv", try_parse_dates=True)
    orders_raw = pl.read_csv(RAW / "orders.csv", try_parse_dates=True)
    products = pl.read_csv(RAW / "products.csv")
    campaign = pl.read_csv(RAW / "email_campaign.csv", try_parse_dates=True)

    customers = clean_customers(customers_raw)
    orders = clean_orders(orders_raw)

    print("Regions before:", sorted(customers_raw["region"].drop_nulls().unique().to_list()))
    print("Regions after: ", sorted(customers["region"].unique().to_list()))
    print(f"Order lines: {orders_raw.height:,} raw -> {orders.height:,} clean")

    # Checks: stop loudly if something is wrong
    assert orders["order_id"].is_unique().all(), "order_id should be unique after cleaning"
    assert customers["customer_id"].is_unique().all(), "customer_id should be unique"
    assert orders["revenue"].min() > 0, "revenue should be positive"

    customers.write_parquet(CLEAN / "customers.parquet")
    orders.write_parquet(CLEAN / "orders.parquet")
    products.write_parquet(CLEAN / "products.parquet")
    campaign.write_parquet(CLEAN / "email_campaign.parquet")
    print("Saved 4 Parquet files to", CLEAN)


if __name__ == "__main__":
    main()

python day02/clean.py

Regions before: [' AB ', ' ATL ', ' BC ', ' ON ', ' QC ', 'AB', 'ATL', 'BC', 'ON', 'QC', 'ab', 'atl', 'bc', 'on', 'qc']
Regions after:  ['AB', 'ATL', 'BC', 'ON', 'QC', 'UNKNOWN']
Order lines: 18,542 raw -> 17,839 clean
Saved 4 Parquet files to data/clean

Why these choices:

• .unique() removes exact duplicates. Removing rows by order_id alone could hide a real data error, so we only drop rows that are identical in every column.
• Missing region → "UNKNOWN", not dropped. Dropping customers with missing region would silently change the sample. A visible category keeps them and lets us check later whether missingness relates to treatment.
• Parquet stores types (dates stay dates) and is much faster and smaller than CSV. From now on we read data/clean/*.parquet.

Try it. Add one more assertion: every order's customer_id must exist in customers. (Hint: is_in.)

assert orders["customer_id"].is_in(customers["customer_id"].implode()).all(), \
    "orders contain unknown customers"

2.3 Group-by and joins: one row per customer (55 min)

Concept. Causal questions have a unit. Ours is the customer, but orders are recorded per line. We collapse lines to customers with group_by().agg() and bring in other tables with join.

Two joining rules protect the analysis:

1. Start from the unit table (customers), then left-join everything onto it. Starting from orders silently drops the 271 customers who never bought.
2. Treatment lives in its own table. email_campaign lists only recipients, so after a left join non-recipients have nulls. We turn that into treated = 0.

Create day02/summaries.py:

# day02/summaries.py
import polars as pl

orders = pl.read_parquet("data/clean/orders.parquet")
customers = pl.read_parquet("data/clean/customers.parquet")
campaign = pl.read_parquet("data/clean/email_campaign.parquet")

# 1. One number for the whole table
print(f"Total revenue: ${orders['revenue'].sum():,.0f}")

# 2. group_by + agg: one row per region... but region lives in customers, so join first
orders_c = orders.join(customers, on="customer_id", how="left")
by_region = (orders_c.group_by("region")
             .agg(n_orders=pl.len(), revenue=pl.col("revenue").sum().round(0).cast(pl.Int64))
             .sort("revenue", descending=True))
print(by_region)

# 3. Customer-level outcome table
per_customer = orders.group_by("customer_id").agg(
    n_orders=pl.col("order_id").n_unique(),
    revenue=pl.col("revenue").sum(),
    first_order=pl.col("order_date").min(),
)
print(per_customer.head())
print("Customers with at least one order:", per_customer.height, "of", customers.height)

# 4. The treatment lives in its own table: left join + fill the gaps with 0
customers_t = (customers
               .join(campaign.select("customer_id", "sent_date"), on="customer_id", how="left")
               .with_columns(treated=pl.col("sent_date").is_not_null().cast(pl.Int8)))
print(customers_t.group_by("treated").len().sort("treated"))

# 5. Keep customers who never ordered! Start from customers, not orders.
analysis = (customers_t
            .join(per_customer, on="customer_id", how="left")
            .with_columns(pl.col("n_orders", "revenue").fill_null(0)))
print("Rows in analysis table:", analysis.height)

Total revenue: $2,810,574
│ ON      ┆ 6608     ┆ 1013648 │
...
Customers with at least one order: 3729 of 4000
│ 0       ┆ 2721 │
│ 1       ┆ 1279 │
Rows in analysis table: 4000

The final table has 4,000 rows: every customer, buyers or not.

Try it. Join orders to products and compute revenue and the average discount by product category, sorted by revenue.

import polars as pl
orders = pl.read_parquet("data/clean/orders.parquet")
products = pl.read_parquet("data/clean/products.parquet")
print(orders.join(products, on="product_id", how="left")
      .group_by("category")
      .agg(revenue=pl.col("revenue").sum().round(0).cast(pl.Int64),
           avg_discount=pl.col("discount_pct").mean().round(1))
      .sort("revenue", descending=True))

│ Electronics ┆ 1678960 ┆ 6.1          │
│ Home        ┆ 435289  ┆ 5.9          │
...

2.4 Lazy queries for bigger data (30 min)

Concept. read_parquet loads the whole file at once. scan_parquet instead builds a plan, optimises it, and reads only what the plan needs. For our 18,000 lines it makes no difference; for the 200 million lines of a real retailer, it is the difference between seconds and a crash.

Create day02/lazy_demo.py:

# day02/lazy_demo.py
import polars as pl

# scan_parquet builds a plan; nothing is read yet
plan = (
    pl.scan_parquet("data/clean/orders.parquet")
      .filter(pl.col("order_date") >= pl.date(2026, 3, 2))
      .group_by("customer_id")
      .agg(post_revenue=pl.col("revenue").sum())
)
print(plan.explain())       # see the optimised plan
result = plan.collect()     # now Polars does the work
print(result.sort("post_revenue", descending=True).head(5))

AGGREGATE[maintain_order: false]
  [col("revenue").sum().alias("post_revenue")] BY [col("customer_id")]
  FROM
  simple π 2/2 ["customer_id", "revenue"]
    Parquet SCAN [data/clean/orders.parquet]
    PROJECT 3/9 COLUMNS
    SELECTION: col("order_date") >= 2026-03-02

Read the plan from the bottom up: Polars reads only 3 of 9 columns and applies the date filter while reading. You wrote the steps in a natural order; Polars rearranged them to be efficient.

Habit for the rest of the course: scan_* + .collect() in pipelines, read_* for quick exploration.

2.5 Day 2 mini-project: the naive comparison (60 min)

Goal. Answer management's question the way most dashboards would: compare emailed and non-emailed customers after 2 March. Then check whether the two groups were comparable before the email.

Create day02/naive_comparison.py:

# day02/naive_comparison.py
"""Mini-project: how different do emailed and non-emailed customers look, before and after?"""
import polars as pl

SENT = pl.date(2026, 3, 2)

customers = pl.read_parquet("data/clean/customers.parquet")
orders = pl.read_parquet("data/clean/orders.parquet")
campaign = pl.read_parquet("data/clean/email_campaign.parquet")

# 1. Orders per customer, before and after the email date
per_customer = orders.group_by("customer_id").agg(
    orders_pre=(pl.col("order_date") < SENT).sum(),
    orders_post=(pl.col("order_date") >= SENT).sum(),
)

# 2. One row per customer, treatment flag, zeros for non-buyers
analysis = (
    customers
    .join(campaign.select("customer_id", treated=pl.lit(1)), on="customer_id", how="left")
    .join(per_customer, on="customer_id", how="left")
    .with_columns(
        pl.col("treated").fill_null(0),
        pl.col("orders_pre", "orders_post").fill_null(0),
    )
)

# 3. Compare group means
summary = (analysis.group_by("treated")
           .agg(n=pl.len(),
                orders_pre=pl.col("orders_pre").mean(),
                orders_post=pl.col("orders_post").mean(),
                loyalty_share=pl.col("is_loyalty").mean())
           .sort("treated"))
print(summary)

t = summary.filter(pl.col("treated") == 1).row(0, named=True)
c = summary.filter(pl.col("treated") == 0).row(0, named=True)

naive = t["orders_post"] / c["orders_post"] - 1
pre_gap = t["orders_pre"] / c["orders_pre"] - 1
growth_t = t["orders_post"] / t["orders_pre"] - 1
growth_c = c["orders_post"] / c["orders_pre"] - 1

print(f"\nNaive 'effect' (post-period gap):    {naive:+.0%}")
print(f"Gap BEFORE the email was even sent:  {pre_gap:+.0%}")
print(f"Growth, emailed customers:           {growth_t:+.1%}")
print(f"Growth, other customers:             {growth_c:+.1%}")
print(f"Difference in growth:                {growth_t - growth_c:+.1%}")
print("\nThe simulation's TRUE effect is +15% more orders.")

python day02/naive_comparison.py

│ 0       ┆ 2721 ┆ 1.850055   ┆ 1.620728    ┆ 0.208012      │
│ 1       ┆ 1279 ┆ 3.232213   ┆ 3.331509    ┆ 0.491009      │

Naive 'effect' (post-period gap):    +106%
Gap BEFORE the email was even sent:  +75%
Growth, emailed customers:           +3.1%
Growth, other customers:             -12.4%
Difference in growth:                +15.5%

The simulation's TRUE effect is +15% more orders.

Interpretation. This is the most important output of Module 1:

1. The naive dashboard says the email doubled orders (+106%). The truth is +15%.
2. Emailed customers already ordered 75% more before the email existed, and were more than twice as likely to be loyalty members. The groups were never comparable. That is selection bias, caused by targeting.
3. Comparing each group with itself over time (growth) and then comparing the growth rates gets close to the truth. That idea is difference-in-differences, and Module 5 makes it rigorous (with standard errors and a check of its key assumption).

None of steps 2–3 is possible without clean pre-period data for both groups. That is why causal work starts with data engineering.

Try it. Repeat the comparison using revenue instead of order counts, split by is_loyalty. Within each loyalty group, are treated and control customers comparable before the email?

import polars as pl
SENT = pl.date(2026, 3, 2)
customers = pl.read_parquet("data/clean/customers.parquet")
orders = pl.read_parquet("data/clean/orders.parquet")
campaign = pl.read_parquet("data/clean/email_campaign.parquet")
per_customer = orders.group_by("customer_id").agg(
    revenue_pre=pl.col("revenue").filter(pl.col("order_date") < SENT).sum(),
    revenue_post=pl.col("revenue").filter(pl.col("order_date") >= SENT).sum(),
)
analysis = (customers
            .join(campaign.select("customer_id", treated=pl.lit(1)), on="customer_id", how="left")
            .join(per_customer, on="customer_id", how="left")
            .with_columns(pl.col("treated").fill_null(0),
                          pl.col("revenue_pre", "revenue_post").fill_null(0)))
print(analysis.group_by("treated", "is_loyalty")
      .agg(n=pl.len(), revenue_pre=pl.col("revenue_pre").mean().round(0),
           revenue_post=pl.col("revenue_post").mean().round(0))
      .sort("is_loyalty", "treated"))

│ 0       ┆ false      ┆ 2155 ┆ 265.0       ┆ 247.0        │
│ 1       ┆ false      ┆ 651  ┆ 430.0       ┆ 480.0        │
│ 0       ┆ true       ┆ 566  ┆ 374.0       ┆ 315.0        │
│ 1       ┆ true       ┆ 628  ┆ 585.0       ┆ 569.0        │

2.6 Write it down and save with Git (15 min)

Add a short section to docs/data_dictionary.md (or a new docs/day02_notes.md) answering, in three sentences: What did the naive comparison say, why is it wrong, and what data would we need to do better? Writing the reasoning down is part of every capstone in this course.

git add .
git commit -m "Day 2: cleaning pipeline, customer table, naive comparison"
git push

Day 2 checklist

• I can use select, filter, with_columns and when/then/otherwise
• I can clean duplicates, messy text and cancelled rows in a rerunnable script with assertions
• I start joins from the unit table so units with no events are kept as zeros
• I can explain the difference between read_parquet and scan_parquet
• I can explain why the naive +106% is wrong, using the pre-period gap