Hands-on lab — Day 123: Groupby and Aggregation

"""Shared fixtures. pytest finds this file by itself -- nothing imports it.

Each fixture returns a FRESH copy of its table, so one test's mutation
(there should not be any, but fixtures should not have to trust that) can
never leak into the next test.
"""

import pytest

from data import (
    build_cat_sales,
    build_large,
    build_orders,
    build_sales,
    build_weighted,
)


@pytest.fixture
def orders():
    return build_orders()


@pytest.fixture
def sales():
    return build_sales()


@pytest.fixture
def cat_sales():
    return build_cat_sales()


@pytest.fixture
def weighted():
    return build_weighted()


@pytest.fixture
def large():
    return build_large()

"""The three tables every exercise in this lab is built from.

Nothing here is randomised except the two performance-comparison tables,
which are seeded so a re-run always produces the same row counts (the
*timing* still varies machine to machine, which is why exercise 8 asserts a
ratio rather than a millisecond figure).

`orders` carries the day's opening failure on purpose: two rows have no
`region` at all, and two more have no `amount`. `sales` and `cat_sales` are
deliberately clean, so exercises 3, 4, 6 and 7 are not fighting missing data
while they teach a different point.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

# --------------------------------------------------------------------------
# `orders` -- exercises 1, 2 and 5. Twelve rows. Two have no region at all
# (order_id 5 and 8); two have no amount at all (order_id 2 and 11).
# --------------------------------------------------------------------------


def build_orders() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "order_id": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
            "region": [
                "North",
                "South",
                "North",
                "East",
                None,
                "South",
                "North",
                None,
                "East",
                "South",
                "West",
                "East",
            ],
            "rep": [
                "Ann",
                "Bo",
                "Ann",
                "Cy",
                "Bo",
                "Ann",
                "Cy",
                "Bo",
                "Ann",
                "Cy",
                "Deb",
                "Deb",
            ],
            "amount": [
                100.0,
                np.nan,
                150.0,
                300.0,
                80.0,
                200.0,
                120.0,
                90.0,
                400.0,
                175.0,
                np.nan,
                500.0,
            ],
        }
    )


# --------------------------------------------------------------------------
# `sales` -- exercises 3, 4 and 6. Twelve rows, four regions of three rows
# each, no missing values anywhere.
# --------------------------------------------------------------------------


def build_sales() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "order_id": [101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112],
            "region": [
                "North",
                "North",
                "North",
                "South",
                "South",
                "South",
                "East",
                "East",
                "East",
                "West",
                "West",
                "West",
            ],
            "rep": ["Ann", "Bo", "Ann", "Bo", "Cy", "Bo", "Cy", "Ann", "Cy", "Ann", "Bo", "Cy"],
            "amount": [
                120.0,
                80.0,
                160.0,
                200.0,
                150.0,
                250.0,
                300.0,
                420.0,
                360.0,
                60.0,
                90.0,
                45.0,
            ],
        }
    )


# --------------------------------------------------------------------------
# `cat_sales` -- exercise 7. Same rows as `sales`, but `region` and `rep`
# are declared as categoricals with categories that are never observed in
# the data ("Central" and "Deb"), so grouping by both keys can either
# manufacture every combination or only the ones actually seen.
# --------------------------------------------------------------------------


def build_cat_sales() -> pd.DataFrame:
    df = build_sales()
    df["region"] = pd.Categorical(df["region"], categories=["North", "South", "East", "West", "Central"])
    df["rep"] = pd.Categorical(df["rep"], categories=["Ann", "Bo", "Cy", "Deb"])
    return df


# --------------------------------------------------------------------------
# `weighted` -- exercise 9. Three groups of unequal size and unequal
# weights, so a weighted mean genuinely differs from a plain mean.
# --------------------------------------------------------------------------


def build_weighted() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "region": ["North", "North", "South", "South", "South", "East", "East"],
            "value": [10.0, 20.0, 5.0, 15.0, 25.0, 100.0, 50.0],
            "weight": [1.0, 3.0, 2.0, 2.0, 1.0, 1.0, 4.0],
        }
    )


# --------------------------------------------------------------------------
# `build_large` -- exercise 8. A frame big enough that the gap between a
# built-in aggregation and a Python-level `.apply` is not measurement noise.
# --------------------------------------------------------------------------


def build_large(n: int = 200_000, n_keys: int = 2000, seed: int = 42) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    return pd.DataFrame(
        {
            "key": rng.integers(0, n_keys, size=n),
            "value": rng.normal(size=n),
        }
    )

"""The worked reference suite for Day 123 -- "Groups That Reconcile".

Nine exercises, each proving one real pandas 3.0.5 behaviour by running code
and reading real values -- never by reading source. Run it:

    pytest examples

Every table these tests use comes from `data.py`, imported through the
fixtures in `conftest.py`. Read `starter/00_brief.md` for the exercise-by-
exercise explanation; this file is the answer key.
"""

import time

import numpy as np
import pandas as pd
import pytest

# --------------------------------------------------------------------------
# Exercise 1 -- the reconciliation invariant
#
# groupby drops rows whose key is missing, by default, silently. Prove the
# gap exists, prove it equals exactly the missing-key rows' total, and prove
# dropna=False makes the parts sum back to the whole.
# --------------------------------------------------------------------------


def test_1_dropna_true_undercounts_by_exactly_the_missing_rows(orders):
    grouped_total = orders.groupby("region")["amount"].sum().sum()
    overall_total = orders["amount"].sum()
    gap = overall_total - grouped_total

    assert grouped_total == 1945.0
    assert overall_total == 2115.0
    assert gap == 170.0

    missing_key_total = orders.loc[orders["region"].isna(), "amount"].sum()
    assert missing_key_total == 170.0
    assert gap == missing_key_total, "the gap must equal exactly the missing-key rows' total"

    missing_key_count = orders["region"].isna().sum()
    assert missing_key_count == 2


def test_1_dropna_false_reconciles_exactly(orders):
    grouped = orders.groupby("region", dropna=False)["amount"].sum()

    # The NaN group is real and carries the missing rows' total.
    assert grouped.loc[np.nan] == 170.0

    # Now the parts add back up to the whole, to the last cent.
    assert grouped.sum() == orders["amount"].sum() == 2115.0


# --------------------------------------------------------------------------
# Exercise 2 -- count() versus size(). size() counts ROWS; count() counts
# NON-MISSING VALUES per column. They disagree exactly where data is
# missing, and confusing them misstates every denominator downstream.
# --------------------------------------------------------------------------


def test_2_size_and_count_disagree_where_amount_is_missing(orders):
    grouped = orders.groupby("region", dropna=False)
    size = grouped.size()
    count = grouped["amount"].count()

    assert size.loc["South"] == 3
    assert count.loc["South"] == 2  # one of South's three amounts is NaN

    assert size.loc["West"] == 1
    assert count.loc["West"] == 0  # West's only amount is NaN

    diff = size - count
    assert diff.sum() == 2
    assert diff.sum() == orders["amount"].isna().sum(), (
        "the total gap between size and count must equal the missing-amount count exactly"
    )


# --------------------------------------------------------------------------
# Exercise 3 -- .agg() four ways: one function, a list, a per-column dict,
# and named aggregation, which is the readable modern form and produces
# flat column names instead of a MultiIndex.
# --------------------------------------------------------------------------


def test_3_agg_single_function(sales):
    result = sales.groupby("region")["amount"].agg("sum")
    assert result.loc["East"] == 1080.0
    assert result.loc["North"] == 360.0
    assert result.loc["South"] == 600.0
    assert result.loc["West"] == 195.0


def test_3_agg_list_of_functions(sales):
    result = sales.groupby("region")["amount"].agg(["sum", "mean", "count"])
    assert list(result.columns) == ["sum", "mean", "count"]
    assert result.loc["East", "sum"] == 1080.0
    assert result.loc["East", "mean"] == pytest.approx(360.0)
    assert result.loc["East", "count"] == 3


def test_3_agg_per_column_dict(sales):
    result = sales.groupby("region").agg({"amount": "sum", "order_id": "count"})
    assert result.loc["North", "amount"] == 360.0
    assert result.loc["North", "order_id"] == 3


def test_3_agg_named_aggregation_gives_flat_columns(sales):
    result = sales.groupby("region").agg(
        total=("amount", "sum"), avg=("amount", "mean"), n=("order_id", "count")
    )
    # Flat column names, not a MultiIndex -- that is the whole point of
    # named aggregation over the list/dict forms above.
    assert list(result.columns) == ["total", "avg", "n"]
    assert not isinstance(result.columns, pd.MultiIndex)
    assert result.loc["West", "total"] == 195.0
    assert result.loc["West", "avg"] == pytest.approx(65.0)
    assert result.loc["West", "n"] == 3


# --------------------------------------------------------------------------
# Exercise 4 -- agg reduces each group to one row; transform returns the
# input's shape, which is how a group statistic gets attached back to
# every row (here: a within-group z-score whose group mean is 0).
# --------------------------------------------------------------------------


def test_4_agg_returns_one_row_per_group(sales):
    result = sales.groupby("region")["amount"].agg("mean")
    assert result.shape == (4,)  # four regions


def test_4_transform_returns_the_input_shape(sales):
    result = sales.groupby("region")["amount"].transform("mean")
    assert result.shape == (12,)  # twelve rows, same as sales itself
    assert result.shape[0] == sales.shape[0]


def test_4_transform_attaches_a_group_mean_to_every_row(sales):
    group_mean = sales.groupby("region")["amount"].transform("mean")
    # North's three amounts are 120, 80, 160 -- mean 120.
    north_rows = sales["region"] == "North"
    assert (group_mean[north_rows] == 120.0).all()


def test_4_within_group_zscore_has_zero_mean_per_group(sales):
    group_mean = sales.groupby("region")["amount"].transform("mean")
    group_std = sales.groupby("region")["amount"].transform("std")
    zscore = (sales["amount"] - group_mean) / group_std

    per_group_zscore_mean = zscore.groupby(sales["region"]).mean()
    for region, value in per_group_zscore_mean.items():
        assert value == pytest.approx(0.0, abs=1e-9), f"{region}'s z-scores must average to 0"


# --------------------------------------------------------------------------
# Exercise 5 -- GroupBy.filter keeps or discards WHOLE GROUPS by a
# predicate. Distinct from Day 122's row-level filtering, despite the
# shared word.
# --------------------------------------------------------------------------


def test_5_filter_drops_whole_groups_below_the_threshold(orders):
    sizes = orders.groupby("region").size()
    assert sizes.to_dict() == {"East": 3, "North": 3, "South": 3, "West": 1}

    survivors = orders.groupby("region").filter(lambda g: len(g) >= 3)

    # West (size 1) is dropped whole; nothing partial survives from it.
    assert "West" not in survivors["region"].unique()
    assert set(survivors["region"].unique()) == {"East", "North", "South"}

    # The row count matches the sum of the surviving groups' own sizes.
    assert survivors.shape[0] == sizes[sizes >= 3].sum() == 9


# --------------------------------------------------------------------------
# Exercise 6 -- multi-key grouping produces a MultiIndex; as_index=False
# gives a flat frame with the same values.
# --------------------------------------------------------------------------


def test_6_multi_key_grouping_produces_a_multiindex(sales):
    result = sales.groupby(["region", "rep"])["amount"].sum()
    assert isinstance(result.index, pd.MultiIndex)
    assert result.index.names == ["region", "rep"]
    assert result.loc[("East", "Ann")] == 420.0
    assert result.loc[("West", "Cy")] == 45.0


def test_6_as_index_false_gives_the_same_values_flat(sales):
    indexed = sales.groupby(["region", "rep"])["amount"].sum()
    flat = sales.groupby(["region", "rep"], as_index=False)["amount"].sum()

    assert not isinstance(flat.index, pd.MultiIndex)
    assert list(flat.columns) == ["region", "rep", "amount"]

    east_ann = flat.loc[(flat["region"] == "East") & (flat["rep"] == "Ann"), "amount"].iloc[0]
    assert east_ann == indexed.loc[("East", "Ann")] == 420.0


# --------------------------------------------------------------------------
# Exercise 7 -- observed=. Grouping a categorical produces rows for
# unobserved combinations unless observed=True. With two categorical keys
# this explodes combinatorially.
# --------------------------------------------------------------------------


def test_7_observed_false_manufactures_unseen_combinations(cat_sales):
    # 5 region categories x 4 rep categories = 20 possible combinations,
    # only 9 of which are ever actually seen in the 12 rows of data.
    result = cat_sales.groupby(["region", "rep"], observed=False).size()
    assert len(result) == 20


def test_7_observed_true_keeps_only_combinations_actually_seen(cat_sales):
    result = cat_sales.groupby(["region", "rep"], observed=True).size()
    assert len(result) == 9
    assert result.loc[("North", "Ann")] == 2


# --------------------------------------------------------------------------
# Exercise 8 -- performance. Report the SHAPE of the gap between a
# built-in aggregation and the equivalent .apply(lambda ...), not a
# millisecond figure -- this is one machine, one day.
# --------------------------------------------------------------------------


def test_8_builtin_agg_beats_apply_by_a_wide_margin(large):
    start = time.perf_counter()
    builtin_result = large.groupby("key")["value"].agg("mean")
    builtin_seconds = time.perf_counter() - start

    start = time.perf_counter()
    apply_result = large.groupby("key")["value"].apply(lambda g: g.mean())
    apply_seconds = time.perf_counter() - start

    # Both paths must agree on the actual numbers -- speed is not the
    # only thing being asserted here.
    assert np.allclose(builtin_result.sort_index().to_numpy(), apply_result.sort_index().to_numpy())

    # The margin is asserted as a conservative ratio, never a timing.
    # This machine measured roughly 10-15x; 3x is asserted as the floor
    # so the check does not flake on a slower or busier machine.
    ratio = apply_seconds / builtin_seconds
    assert ratio >= 3.0, f"expected .apply to be at least 3x slower, measured {ratio:.1f}x"


def test_8_sort_false_does_not_change_the_values(large):
    sorted_result = large.groupby("key", sort=True)["value"].sum()
    unsorted_result = large.groupby("key", sort=False)["value"].sum()
    assert sorted_result.sort_index().equals(unsorted_result.sort_index())
    # sort=False is not guaranteed to change ORDER on every input, but it
    # must never change the VALUES -- that is the only thing asserted here.


# --------------------------------------------------------------------------
# Exercise 9 -- a weighted mean per group, computed with apply and again
# without it, asserting the two agree.
# --------------------------------------------------------------------------


def test_9_weighted_mean_via_apply(weighted):
    def weighted_mean(group: pd.DataFrame) -> float:
        return float(np.average(group["value"], weights=group["weight"]))

    via_apply = weighted.groupby("region").apply(weighted_mean, include_groups=False)

    assert via_apply.loc["North"] == pytest.approx(17.5)  # (10*1 + 20*3) / 4
    assert via_apply.loc["South"] == pytest.approx(13.0)  # (5*2+15*2+25*1) / 5
    assert via_apply.loc["East"] == pytest.approx(60.0)  # (100*1 + 50*4) / 5


def test_9_weighted_mean_without_apply_agrees(weighted):
    def weighted_mean(group: pd.DataFrame) -> float:
        return float(np.average(group["value"], weights=group["weight"]))

    via_apply = weighted.groupby("region").apply(weighted_mean, include_groups=False)

    # The vectorised route: build a value*weight column, sum both pieces
    # per group with .agg, then divide -- no apply anywhere.
    weighted_products = weighted.assign(value_weight=weighted["value"] * weighted["weight"])
    sums = weighted_products.groupby("region").agg(
        sum_value_weight=("value_weight", "sum"), sum_weight=("weight", "sum")
    )
    without_apply = sums["sum_value_weight"] / sums["sum_weight"]

    assert without_apply.sort_index().equals(via_apply.sort_index())

lesson_id: D123
day: 123
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-123-groupby-and-aggregation
  - python3 -m venv .venv
  - .venv/bin/pip install -r requirements/requirements.txt
  - .venv/bin/python3 -c "import pandas; print(pandas.__version__)"
run_commands:
  - .venv/bin/pytest examples
  - .venv/bin/pytest starter
test_commands:
  - bash tests/run_tests.sh
cleanup_commands:
  - "find . -path ./.venv -prune -o -type d -name '__pycache__' -print -exec rm -rf -- {} +"
  - rm -rf .pytest_cache
  - 'rm -rf .venv  # optional: removes the lab virtual environment'
  - 'git checkout -- starter/  # optional: reset your work'
requires_network: true
requires_api_key: false
estimated_minutes: 35
last_executed: '2026-08-19'
executed_on: 'macOS 26.5.2 (Apple Silicon, arm64), Python 3.14.0, pandas 3.0.5, pyarrow 25.0.1, numpy 2.5.2, pytest 9.1.1, bash 3.2.57 -- bash tests/run_tests.sh -> 13 checks, 0 failure(s), exit 0. pytest examples -> 20 passed. pytest starter -> 20 skipped (untouched checkout). Section 5 of the harness solves every exercise in a scratch copy (20 passed), deliberately breaks the exercise-1 gap assertion (170.0 -> 999.0), confirms the run exits non-zero with a printed FAIL, restores the file, and confirms 20 passed again -- so the suite is demonstrated to be capable of failing rather than merely claimed to be. Separately, the coordinator broke the same assertion directly in examples/test_groupby.py (not a scratch copy) and re-ran the full harness: 1 failed, 19 passed from pytest, and 6 of the harness''s own 13 checks failed with overall exit 1; restoring the line returned the harness to 13 checks, 0 failure(s), exit 0. Everything was run through a real lab-local .venv created by the documented setup commands. Two honesty notes from this run. FIRST: exercise 8''s performance ratio was measured at approximately 10-15x (built-in .agg(''mean'') versus .apply(lambda g: g.mean()) on a 200,000-row, 2,000-key frame) in prototyping before this lab was authored; the lab asserts only a conservative floor of 3.0x so the check does not flake on a slower or busier machine, and no millisecond figure is asserted anywhere. SECOND: pandas 3.0.5 has NOT changed groupby''s observed= default for categorical keys to True, despite pandas 2.1 having announced that change as coming in a future major version -- observed=False remains the default measured here, and exercise 7 states this explicitly rather than assuming the deprecation had already landed. matplotlib, scipy and polars are not installed in this environment; polars'' group_by is described from its public documentation in the lesson''s Tools section as a design contrast and no output attributed to it is reproduced anywhere.'

# What is installed, why, and what it costs

Four packages, all free and open source, installed into a lab-local
virtual environment that `rm -rf .venv` completely undoes.

| Package | Version pinned | Licence | What this lab uses it for |
| --- | --- | --- | --- |
| `pandas` | 3.0.5 | BSD 3-Clause | Every `groupby`, `.agg`, `.transform` and `.filter` call in this lab. |
| `pyarrow` | 25.0.1 | Apache 2.0 | pandas 3.0's default backend for string and several nullable dtypes; installed for parity with Days 120-122 even though this lab's tables use plain numeric and object columns throughout. |
| `numpy` | 2.5.2 | BSD 3-Clause | `np.nan`, `np.average` for exercise 9's weighted mean, and `np.random.default_rng` for exercise 8's synthetic 200,000-row table. |
| `pytest` | 9.1.1 | MIT | The test harness every exercise is written against. |

There is no paid tier of anything in this lab, no account, no key and no
signup, personally or commercially.

## The one time the network is needed

```bash
.venv/bin/pip install -r requirements/requirements.txt
```

That is the only command in the lab that opens a connection. Every script
and test after that runs completely offline.

## What is deliberately *not* installed

**matplotlib**, **scipy** and **polars** are not installed in this
environment. The lesson's Tools section describes polars' `group_by`
from its public documentation as a design contrast to pandas — no output
from polars, scipy or matplotlib is reproduced anywhere in this lab or
its lesson; every place they are mentioned says so plainly.

## If you cannot install anything at all

pandas is not in the Python standard library, and there is no reduced
path through this lab without it. If pandas genuinely cannot be
installed, read the lesson's captured output and `expected-output/`
directory instead; every number there came from a real run and is not
invented.

pandas==3.0.5
pyarrow==25.0.1
numpy==2.5.2
pytest==9.1.1

# Day 123 lab — the brief

Nine exercises, in order. Work top to bottom in `test_groupby.py`. Every
table comes from a fixture defined in `conftest.py` (`orders`, `sales`,
`cat_sales`, `weighted`, `large`) — read `data.py` once to see exactly what
each one contains before you start.

Check yourself at any point:

```bash
.venv/bin/pytest starter -v
```

On an untouched checkout that prints `20 skipped`. A **skip** means "not
attempted". Replace a `pytest.skip(...)` line with real assertions and
delete it — when every skip is gone and the suite is green, you are
finished:

```bash
.venv/bin/pytest starter -q
echo $?
```

Assert exact values everywhere except exercise 8's timing ratio, which is
inherently one machine on one day — `pytest.approx` is for that ratio and
for the z-scores in exercise 4, nowhere else.

---

## Exercise 1 — the reconciliation invariant (`orders`)

`orders` has two rows with no `region` at all. `groupby('region')` drops
them by default, silently. In `test_1_dropna_true_undercounts_by_exactly_the_missing_rows`:

- Group by `region`, sum `amount`, sum the per-group sums. That total is
  **less** than `orders['amount'].sum()`.
- Assert the exact gap, and assert the gap equals exactly the sum of
  `amount` on the rows where `region` is missing.

In `test_1_dropna_false_reconciles_exactly`: repeat with `dropna=False`.
Now there is a real `NaN` group carrying the missing rows' total, and the
grouped sum equals the overall total exactly. **This is the habit the
whole day is built on: after any groupby aggregation, check that the parts
reconcile with the whole.**

## Exercise 2 — `count()` versus `size()` (`orders`)

`orders` also has two rows with no `amount`. `size()` counts **rows**;
`count()` counts **non-missing values per column**. Group with
`dropna=False` so the region-missing rows stay in the picture too. Assert
that `size` and `count` disagree on the groups that contain a missing
`amount`, and that `(size - count).sum()` equals
`orders['amount'].isna().sum()` exactly.

## Exercise 3 — `.agg()` four ways (`sales`)

`sales` is clean: four regions, three rows each. Write four tests:

1. A single function: `.agg('sum')`.
2. A list: `.agg(['sum', 'mean', 'count'])` — assert the column names.
3. A per-column dict: `.agg({'amount': 'sum', 'order_id': 'count'})`.
4. Named aggregation: `.agg(total=('amount', 'sum'), avg=('amount', 'mean'), n=('order_id', 'count'))`
   — assert the result columns are **flat**, not a `pandas.MultiIndex`.
   This is the readable modern form.

## Exercise 4 — shapes: `agg` versus `transform` (`sales`)

- `agg` reduces each group to one row: assert the shape has one entry per
  group (4,).
- `transform` returns the **input's shape**: assert its shape matches
  `sales.shape[0]` (12), not the number of groups.
- Use `transform` to attach North's group mean (120.0) to every North row.
- Build a within-group z-score: `(amount - group_mean) / group_std`, both
  computed with `transform`. Assert each region's z-scores average to 0
  (`pytest.approx`, `abs=1e-9`).

## Exercise 5 — `GroupBy.filter` (`orders`)

`filter` keeps or drops **whole groups**, unlike Day 122's row-level
`.query()`/boolean-mask filtering — same word, different operation, worth
naming as such. `orders` grouped by `region` (default `dropna=True`) has
sizes East=3, North=3, South=3, West=1. Filter to groups of size `>= 3`.
Assert West is entirely absent from the survivors, and that the survivors'
row count equals the sum of the surviving groups' own sizes (9).

## Exercise 6 — multi-key grouping (`sales`)

Group `sales` by `['region', 'rep']`. Assert the result's index is a
`pandas.MultiIndex` with `.names == ['region', 'rep']`, and check one
value: `('East', 'Ann')` is 420.0. Repeat with `as_index=False` and assert
the same value comes back in a flat frame instead.

## Exercise 7 — `observed=` (`cat_sales`)

`cat_sales` declares `region` (5 categories) and `rep` (4 categories) as
pandas categoricals, two of which — `Central` and `Deb` — never actually
appear in the 12 rows of data. Group by both keys with `observed=False`:
assert 20 rows (5 x 4, every possible combination, most never observed).
Repeat with `observed=True`: assert 9 rows (only the combinations that are
actually present), and check `('North', 'Ann')` is 2.

## Exercise 8 — performance (`large`, 200,000 rows)

Time `large.groupby('key')['value'].agg('mean')` against
`large.groupby('key')['value'].apply(lambda g: g.mean())`. First assert
both give the same numbers (`np.allclose` after sorting both by index).
Then assert `apply_seconds / builtin_seconds >= 3.0` — a conservative
floor; this machine measured roughly 10-15x, but a slower or busier
machine should still clear 3x easily. Never assert a millisecond figure.

Second test: assert `groupby(sort=True)` and `groupby(sort=False)` give
the same **values** (once both are sorted by index for comparison) — only
the ordering of unsorted work should ever differ, never the numbers.

## Exercise 9 — a weighted mean per group (`weighted`)

Write a `weighted_mean(group)` function using
`np.average(group['value'], weights=group['weight'])`. Apply it per
`region` with `include_groups=False`. Assert North is `approx(17.5)`,
South is `approx(13.0)`, East is `approx(60.0)`.

Then compute the same weighted means **without** `apply`: build a
`value * weight` column, sum both that column and `weight` per group with
`.agg`, then divide. Assert the two routes agree exactly. This is the
honest case where `apply` is the readable choice, and you can still check
it against a faster vectorised route.

---

Prove your suite is not vacuous once you are green: re-break one assertion
on purpose (flip a comparison, change an expected number), confirm the run
exits non-zero with a printed `FAIL`, then restore it and confirm green
again.

"""Shared fixtures. pytest finds this file by itself -- nothing imports it.

Each fixture returns a FRESH copy of its table, so one test's mutation
(there should not be any, but fixtures should not have to trust that) can
never leak into the next test.
"""

import pytest

from data import (
    build_cat_sales,
    build_large,
    build_orders,
    build_sales,
    build_weighted,
)


@pytest.fixture
def orders():
    return build_orders()


@pytest.fixture
def sales():
    return build_sales()


@pytest.fixture
def cat_sales():
    return build_cat_sales()


@pytest.fixture
def weighted():
    return build_weighted()


@pytest.fixture
def large():
    return build_large()

"""The three tables every exercise in this lab is built from.

Nothing here is randomised except the two performance-comparison tables,
which are seeded so a re-run always produces the same row counts (the
*timing* still varies machine to machine, which is why exercise 8 asserts a
ratio rather than a millisecond figure).

`orders` carries the day's opening failure on purpose: two rows have no
`region` at all, and two more have no `amount`. `sales` and `cat_sales` are
deliberately clean, so exercises 3, 4, 6 and 7 are not fighting missing data
while they teach a different point.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

# --------------------------------------------------------------------------
# `orders` -- exercises 1, 2 and 5. Twelve rows. Two have no region at all
# (order_id 5 and 8); two have no amount at all (order_id 2 and 11).
# --------------------------------------------------------------------------


def build_orders() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "order_id": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
            "region": [
                "North",
                "South",
                "North",
                "East",
                None,
                "South",
                "North",
                None,
                "East",
                "South",
                "West",
                "East",
            ],
            "rep": [
                "Ann",
                "Bo",
                "Ann",
                "Cy",
                "Bo",
                "Ann",
                "Cy",
                "Bo",
                "Ann",
                "Cy",
                "Deb",
                "Deb",
            ],
            "amount": [
                100.0,
                np.nan,
                150.0,
                300.0,
                80.0,
                200.0,
                120.0,
                90.0,
                400.0,
                175.0,
                np.nan,
                500.0,
            ],
        }
    )


# --------------------------------------------------------------------------
# `sales` -- exercises 3, 4 and 6. Twelve rows, four regions of three rows
# each, no missing values anywhere.
# --------------------------------------------------------------------------


def build_sales() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "order_id": [101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112],
            "region": [
                "North",
                "North",
                "North",
                "South",
                "South",
                "South",
                "East",
                "East",
                "East",
                "West",
                "West",
                "West",
            ],
            "rep": ["Ann", "Bo", "Ann", "Bo", "Cy", "Bo", "Cy", "Ann", "Cy", "Ann", "Bo", "Cy"],
            "amount": [
                120.0,
                80.0,
                160.0,
                200.0,
                150.0,
                250.0,
                300.0,
                420.0,
                360.0,
                60.0,
                90.0,
                45.0,
            ],
        }
    )


# --------------------------------------------------------------------------
# `cat_sales` -- exercise 7. Same rows as `sales`, but `region` and `rep`
# are declared as categoricals with categories that are never observed in
# the data ("Central" and "Deb"), so grouping by both keys can either
# manufacture every combination or only the ones actually seen.
# --------------------------------------------------------------------------


def build_cat_sales() -> pd.DataFrame:
    df = build_sales()
    df["region"] = pd.Categorical(df["region"], categories=["North", "South", "East", "West", "Central"])
    df["rep"] = pd.Categorical(df["rep"], categories=["Ann", "Bo", "Cy", "Deb"])
    return df


# --------------------------------------------------------------------------
# `weighted` -- exercise 9. Three groups of unequal size and unequal
# weights, so a weighted mean genuinely differs from a plain mean.
# --------------------------------------------------------------------------


def build_weighted() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "region": ["North", "North", "South", "South", "South", "East", "East"],
            "value": [10.0, 20.0, 5.0, 15.0, 25.0, 100.0, 50.0],
            "weight": [1.0, 3.0, 2.0, 2.0, 1.0, 1.0, 4.0],
        }
    )


# --------------------------------------------------------------------------
# `build_large` -- exercise 8. A frame big enough that the gap between a
# built-in aggregation and a Python-level `.apply` is not measurement noise.
# --------------------------------------------------------------------------


def build_large(n: int = 200_000, n_keys: int = 2000, seed: int = 42) -> pd.DataFrame:
    rng = np.random.default_rng(seed)
    return pd.DataFrame(
        {
            "key": rng.integers(0, n_keys, size=n),
            "value": rng.normal(size=n),
        }
    )

"""YOUR test suite for Day 123 -- "Groups That Reconcile". Nine exercises.

Run it from the lab directory, not from here:

    pytest starter -v

Every exercise below ends in a `pytest.skip(...)` line. pytest reports a
skip as `s` in the dot line and moves on, so an unfinished suite still
exits 0. Replace each skip with real assertions -- deleting the skip line
is part of the exercise. `starter/00_brief.md` explains each exercise in
full; the fixtures you need (`orders`, `sales`, `cat_sales`, `weighted`,
`large`) come from `conftest.py` and are described there too.

Assert exact values everywhere except exercise 8's timing ratio, which is
inherently one machine on one day.
"""

import time

import numpy as np
import pandas as pd
import pytest

# --------------------------------------------------------------------------
# EXERCISE 1 -- the reconciliation invariant. groupby drops rows whose key
# is missing, by default, silently. See starter/00_brief.md exercise 1.
#
# Check with:   pytest starter -v -k test_1
# --------------------------------------------------------------------------


def test_1_dropna_true_undercounts_by_exactly_the_missing_rows(orders):
    pytest.skip(
        "exercise 1a: assert grouped_total (1945.0), overall_total (2115.0), "
        "the gap (170.0), and that the gap equals the missing-key rows' amount total exactly"
    )


def test_1_dropna_false_reconciles_exactly(orders):
    pytest.skip(
        "exercise 1b: with dropna=False, assert the NaN group's total is 170.0 "
        "and that the grouped sum equals orders['amount'].sum() exactly"
    )


# --------------------------------------------------------------------------
# EXERCISE 2 -- count() versus size(). size() counts rows; count() counts
# non-missing values per column.
#
# Check with:   pytest starter -v -k test_2
# --------------------------------------------------------------------------


def test_2_size_and_count_disagree_where_amount_is_missing(orders):
    pytest.skip(
        "exercise 2: group orders by region with dropna=False; assert size and "
        "count disagree on South and West, and that (size - count).sum() equals "
        "orders['amount'].isna().sum() exactly"
    )


# --------------------------------------------------------------------------
# EXERCISE 3 -- .agg() four ways: single function, list, per-column dict,
# and named aggregation.
#
# Check with:   pytest starter -v -k test_3
# --------------------------------------------------------------------------


def test_3_agg_single_function(sales):
    pytest.skip("exercise 3a: sales.groupby('region')['amount'].agg('sum'); assert East is 1080.0")


def test_3_agg_list_of_functions(sales):
    pytest.skip(
        "exercise 3b: .agg(['sum', 'mean', 'count']); assert the column names "
        "and East's three values"
    )


def test_3_agg_per_column_dict(sales):
    pytest.skip(
        "exercise 3c: .agg({'amount': 'sum', 'order_id': 'count'}); assert "
        "North's amount sum and order_id count"
    )


def test_3_agg_named_aggregation_gives_flat_columns(sales):
    pytest.skip(
        "exercise 3d: .agg(total=('amount','sum'), avg=('amount','mean'), "
        "n=('order_id','count')); assert the result columns are flat, NOT a "
        "MultiIndex, and check West's three values"
    )


# --------------------------------------------------------------------------
# EXERCISE 4 -- agg reduces each group to one row; transform returns the
# input's shape. Use transform to build a within-group z-score.
#
# Check with:   pytest starter -v -k test_4
# --------------------------------------------------------------------------


def test_4_agg_returns_one_row_per_group(sales):
    pytest.skip("exercise 4a: assert sales.groupby('region')['amount'].agg('mean').shape == (4,)")


def test_4_transform_returns_the_input_shape(sales):
    pytest.skip(
        "exercise 4b: assert sales.groupby('region')['amount'].transform('mean').shape "
        "matches sales.shape[0], not the number of groups"
    )


def test_4_transform_attaches_a_group_mean_to_every_row(sales):
    pytest.skip(
        "exercise 4c: assert every North row's transformed group mean equals 120.0"
    )


def test_4_within_group_zscore_has_zero_mean_per_group(sales):
    pytest.skip(
        "exercise 4d: build (amount - group_mean) / group_std using transform for "
        "both, then assert each region's z-scores average to 0 within pytest.approx"
    )


# --------------------------------------------------------------------------
# EXERCISE 5 -- GroupBy.filter keeps or drops WHOLE GROUPS by a predicate.
# Distinct from Day 122's row-level filtering.
#
# Check with:   pytest starter -v -k test_5
# --------------------------------------------------------------------------


def test_5_filter_drops_whole_groups_below_the_threshold(orders):
    pytest.skip(
        "exercise 5: filter orders.groupby('region') to groups of size >= 3; "
        "assert West is entirely absent from the survivors and that the survivors' "
        "row count equals the sum of the surviving groups' own sizes (9)"
    )


# --------------------------------------------------------------------------
# EXERCISE 6 -- multi-key grouping produces a MultiIndex; as_index=False
# gives a flat frame with the same values.
#
# Check with:   pytest starter -v -k test_6
# --------------------------------------------------------------------------


def test_6_multi_key_grouping_produces_a_multiindex(sales):
    pytest.skip(
        "exercise 6a: group sales by ['region', 'rep'], assert the result index "
        "is a pandas.MultiIndex with names ['region', 'rep'], and check "
        "('East', 'Ann') is 420.0"
    )


def test_6_as_index_false_gives_the_same_values_flat(sales):
    pytest.skip(
        "exercise 6b: repeat with as_index=False; assert the result is NOT a "
        "MultiIndex and that ('East', 'Ann')'s value still matches exercise 6a"
    )


# --------------------------------------------------------------------------
# EXERCISE 7 -- observed=. Grouping a categorical produces rows for every
# possible combination unless observed=True.
#
# Check with:   pytest starter -v -k test_7
# --------------------------------------------------------------------------


def test_7_observed_false_manufactures_unseen_combinations(cat_sales):
    pytest.skip(
        "exercise 7a: group cat_sales by ['region', 'rep'] with observed=False; "
        "assert the result has 20 rows (5 region categories x 4 rep categories)"
    )


def test_7_observed_true_keeps_only_combinations_actually_seen(cat_sales):
    pytest.skip(
        "exercise 7b: repeat with observed=True; assert the result has 9 rows "
        "and that ('North', 'Ann') is 2"
    )


# --------------------------------------------------------------------------
# EXERCISE 8 -- performance. Report the SHAPE of the gap, never a
# millisecond figure.
#
# Check with:   pytest starter -v -k test_8
# --------------------------------------------------------------------------


def test_8_builtin_agg_beats_apply_by_a_wide_margin(large):
    pytest.skip(
        "exercise 8a: time large.groupby('key')['value'].agg('mean') against "
        "large.groupby('key')['value'].apply(lambda g: g.mean()); assert both give "
        "the same numbers, then assert apply_seconds / builtin_seconds >= 3.0"
    )


def test_8_sort_false_does_not_change_the_values(large):
    pytest.skip(
        "exercise 8b: assert groupby(sort=True) and groupby(sort=False) give the "
        "same values once both are sorted for comparison"
    )


# --------------------------------------------------------------------------
# EXERCISE 9 -- a weighted mean per group, computed with apply and again
# without it, asserting the two agree.
#
# Check with:   pytest starter -v -k test_9
# --------------------------------------------------------------------------


def test_9_weighted_mean_via_apply(weighted):
    pytest.skip(
        "exercise 9a: write a weighted_mean(group) function using "
        "np.average(group['value'], weights=group['weight']); apply it per region "
        "with include_groups=False; assert North is approx(17.5), South approx(13.0), "
        "East approx(60.0)"
    )


def test_9_weighted_mean_without_apply_agrees(weighted):
    pytest.skip(
        "exercise 9b: compute the same weighted means WITHOUT apply -- build a "
        "value*weight column, sum both columns per group with .agg, then divide -- "
        "and assert the result equals exercise 9a's result exactly"
    )

#!/usr/bin/env bash
# Tests for the Day 123 lab. Run from the lab directory:
#   bash tests/run_tests.sh
#
# The harness proves the lesson's claims by running code and reading real
# values, never by reading source:
#
#   * groupby drops rows with a missing key by default, silently, and the
#     resulting gap equals exactly the missing-key rows' total; dropna=False
#     makes the parts sum back to the whole;
#   * size() and count() disagree exactly where data is missing;
#   * .agg() four ways -- single function, list, per-column dict, and named
#     aggregation -- produce the documented shapes and values, and named
#     aggregation's columns are flat, never a MultiIndex;
#   * agg returns one row per group; transform returns the input's shape,
#     and a transform-built within-group z-score averages to zero per group;
#   * GroupBy.filter keeps or drops WHOLE groups by a size predicate;
#   * multi-key grouping produces a MultiIndex, and as_index=False gives the
#     same values flat;
#   * observed=False manufactures every possible categorical combination
#     (20, on this lab's two keys); observed=True keeps only the 9 seen;
#   * a built-in aggregation beats the equivalent .apply by a wide,
#     machine-independent margin, asserted as a ratio floor, never a timing;
#   * a weighted mean computed with apply and again without it agree;
#   * the reference suite (`examples/`) passes in full;
#   * the exercise suite (`starter/`) is all-skip on an untouched checkout,
#     and the harness proves it can genuinely FAIL by solving every exercise
#     in a scratch copy, breaking one assertion on purpose, confirming a
#     non-zero exit and a printed FAIL, then restoring it;
#   * nothing is left behind on disk.
#
# Everything after the one-time install runs offline. Nothing binds a port,
# nothing writes outside the lab, nothing needs a key. Deterministic,
# non-interactive, exits 0 only if every check passes.
set -u

export PYTHONDONTWRITEBYTECODE=1

lab_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# Bytecode left by an EARLIER command is not this run's litter. The README
# documents `pytest starter -q` and `pytest examples -q` separately, and
# running either writes .pyc files that would then fail the cleanliness
# check at the end of this script -- failing the reader for following the
# instructions. Clearing them here makes that final check measure what it
# claims to: what THIS run left behind. `.venv` is untouched, because the
# packages' own bytecode is theirs, not ours.
find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -exec rm -rf {} + 2>/dev/null || true

failures=0
checks=0

check() {
  local label="$1" ok="$2"
  checks=$((checks + 1))
  if [ "${ok}" = "yes" ]; then
    echo "  ok: ${label}"
  else
    echo "  FAIL: ${label}"
    failures=$((failures + 1))
  fi
}

check_eq() {
  # check_eq <label> <expected> <actual>
  if [ "$2" = "$3" ]; then
    check "$1" "yes"
  else
    check "$1 (expected [$2], got [$3])" "no"
  fi
}

# Resolve pytest: an explicit override, then this lab's .venv, then PATH.
# Fails loudly with instructions rather than silently skipping checks.
resolve_tool() {
  local tool="$1" override="$2"
  if [ -n "${override}" ] && [ -x "${override}" ]; then echo "${override}"; return 0; fi
  if [ -x "${lab_dir}/.venv/bin/${tool}" ]; then echo "${lab_dir}/.venv/bin/${tool}"; return 0; fi
  if command -v "${tool}" >/dev/null 2>&1; then command -v "${tool}"; return 0; fi
  return 1
}

pytest_bin="$(resolve_tool pytest "${PYTEST:-}")" || {
  echo "FAIL: pytest not found." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  echo "  Or point this suite at an existing pytest:" >&2
  echo "    PYTEST=/path/to/pytest bash tests/run_tests.sh" >&2
  exit 1
}

python_bin="$(dirname "${pytest_bin}")/python3"
if [ ! -x "${python_bin}" ]; then
  python_bin="$(command -v python3 || true)"
fi
if [ -z "${python_bin}" ]; then
  echo "FAIL: python3 not found on PATH." >&2
  exit 1
fi

if ! "${python_bin}" -c "import pandas" >/dev/null 2>&1; then
  echo "FAIL: pandas is not importable from ${python_bin}." >&2
  echo "  Install the lab's dependencies with:" >&2
  echo "    python3 -m venv .venv" >&2
  echo "    .venv/bin/pip install -r requirements/requirements.txt" >&2
  exit 1
fi

echo "Day 123 — Groups That Reconcile"
echo

# --------------------------------------------------------------------------
echo "1. The tools and the versions this lab was written against"
# --------------------------------------------------------------------------

versions="$("${python_bin}" - <<'PY'
import platform
import sys
from importlib.metadata import version

print(f"python   {platform.python_version()}")
for name in ("pandas", "pyarrow", "numpy", "pytest"):
    try:
        print(f"{name:<8} {version(name)}")
    except Exception as exc:  # pragma: no cover
        print(f"{name:<8} NOT INSTALLED ({exc})")
PY
)"
echo "${versions}"
echo

pandas_version="$("${python_bin}" -c "import pandas; print(pandas.__version__)" 2>/dev/null || echo "")"
pinned_pandas="$(grep -m1 '^pandas==' "${lab_dir}/requirements/requirements.txt" | cut -d= -f3)"
check_eq "installed pandas matches requirements.txt exactly" "${pinned_pandas}" "${pandas_version}"
echo

# --------------------------------------------------------------------------
echo "2. Reference suite -- examples/ must pass in full"
# --------------------------------------------------------------------------

examples_output="$(cd "${lab_dir}" && "${pytest_bin}" examples -q 2>&1)"
examples_status=$?
echo "${examples_output}" | tail -5
check "examples/ exits 0" "$( [ ${examples_status} -eq 0 ] && echo yes || echo no )"

examples_passed_line="$(echo "${examples_output}" | grep -E '^[0-9]+ passed' || true)"
check "examples/ reports 20 passed, 0 failed" "$( echo "${examples_passed_line}" | grep -qE '^20 passed' && echo yes || echo no )"
echo

# --------------------------------------------------------------------------
echo "3. Exercise suite -- starter/ is all-skip on an untouched checkout"
# --------------------------------------------------------------------------

starter_output="$(cd "${lab_dir}" && "${pytest_bin}" starter -q 2>&1)"
starter_status=$?
echo "${starter_output}" | tail -5
check "starter/ (untouched) exits 0" "$( [ ${starter_status} -eq 0 ] && echo yes || echo no )"
check "starter/ (untouched) reports 20 skipped, 0 failed" "$( echo "${starter_output}" | grep -qE '^20 skipped' && echo yes || echo no )"
echo

# --------------------------------------------------------------------------
echo "4. Never run 'pytest examples starter' in one invocation -- same"
echo "   module name in both directories means the second collected can"
echo "   shadow the first. Documented and checked separately, above."
# --------------------------------------------------------------------------
echo

# --------------------------------------------------------------------------
echo "5. Prove the suite can genuinely FAIL: solve every exercise in a"
echo "   scratch copy, confirm green, break one assertion on purpose,"
echo "   confirm a non-zero exit and a printed FAIL, then restore."
# --------------------------------------------------------------------------

scratch_dir="$(mktemp -d "${TMPDIR:-/tmp}/d123-scratch.XXXXXX")"
cleanup_scratch() { rm -rf "${scratch_dir}"; }
trap cleanup_scratch EXIT

cp "${lab_dir}/examples/test_groupby.py" "${scratch_dir}/test_groupby.py"
cp "${lab_dir}/examples/data.py" "${scratch_dir}/data.py"
cp "${lab_dir}/examples/conftest.py" "${scratch_dir}/conftest.py"

solved_output="$("${pytest_bin}" "${scratch_dir}" -q 2>&1)"
solved_status=$?
check "scratch copy of the solved suite exits 0" "$( [ ${solved_status} -eq 0 ] && echo yes || echo no )"
check "scratch copy reports 20 passed" "$( echo "${solved_output}" | grep -qE '^20 passed' && echo yes || echo no )"

# Break test_1's exact gap assertion on purpose: 170.0 -> 999.0.
sed -i.bak 's/assert gap == 170\.0/assert gap == 999.0/' "${scratch_dir}/test_groupby.py"

broken_output="$("${pytest_bin}" "${scratch_dir}" -q 2>&1)"
broken_status=$?
check "broken scratch copy exits non-zero" "$( [ ${broken_status} -ne 0 ] && echo yes || echo no )"
check "broken scratch copy prints a FAIL/failed line" "$( echo "${broken_output}" | grep -qiE 'failed|assert' && echo yes || echo no )"

mv "${scratch_dir}/test_groupby.py.bak" "${scratch_dir}/test_groupby.py"
restored_output="$("${pytest_bin}" "${scratch_dir}" -q 2>&1)"
restored_status=$?
check "restored scratch copy exits 0 again" "$( [ ${restored_status} -eq 0 ] && echo yes || echo no )"
check "restored scratch copy reports 20 passed again" "$( echo "${restored_output}" | grep -qE '^20 passed' && echo yes || echo no )"

cleanup_scratch
trap - EXIT
echo

# --------------------------------------------------------------------------
echo "6. Nothing in examples/ or starter/ opens a network connection"
# --------------------------------------------------------------------------

url_hits="$(grep -rEl 'https?://|ftp://' "${lab_dir}/examples" "${lab_dir}/starter" 2>/dev/null || true)"
check "no URLs inside examples/ or starter/" "$( [ -z "${url_hits}" ] && echo yes || echo no )"
echo

# --------------------------------------------------------------------------
echo "7. Cleanliness -- nothing left behind by THIS run"
# --------------------------------------------------------------------------

find "${lab_dir}" -name '.venv' -prune -o -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true
find "${lab_dir}" -name '.venv' -prune -o -type d -name '.pytest_cache' -exec rm -rf {} + 2>/dev/null || true

stray="$(find "${lab_dir}" -name '.venv' -prune -o \( -type d -name '__pycache__' -print -o -type d -name '.pytest_cache' -print \) 2>/dev/null || true)"
check "no __pycache__ or .pytest_cache left behind" "$( [ -z "${stray}" ] && echo yes || echo no )"
echo

echo "-------------------------------------------------------------"
echo "${checks} checks, ${failures} failure(s)"
if [ "${failures}" -gt 0 ]; then
  exit 1
fi
exit 0