Hands-on lab — Day 122: Selecting and Filtering

"""Exercise 1 -- the partition invariant, the opening failure of this day.

Run: python3 01_partition_invariant.py

A filter that reports "high performers" and a filter that reports
"everyone else" should, together, account for every row. Split a column
with missing values two ways -- score > 50 and score <= 50 -- and the two
groups do NOT add up to the whole frame. Rows where score is NaN fail BOTH
comparisons and vanish from both halves, because NaN compared with < or >
or <= or >= is always False. Nobody deleted them; they are simply not in
either answer. The fix is not cleverness, it is a habit: name the missing
rows explicitly and check that all three groups together account for the
whole frame.
"""

import numpy as np
import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


scores = pd.DataFrame(
    {
        "name": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay", "Gio", "Hu"],
        "score": [72, 45, np.nan, 91, 50, np.nan, 88, 33],
    }
)
print("scores:")
print(scores)

total = len(scores)
high = scores[scores.score > 50]
low = scores[scores.score <= 50]
missing_count = int(scores.score.isna().sum())

print(f"\ntotal rows:           {total}")
print(f"high (score > 50):    {len(high)}  -> {high.name.tolist()}")
print(f"low  (score <= 50):   {len(low)}   -> {low.name.tolist()}")
print(f"missing (score NaN):  {missing_count} -> {scores.loc[scores.score.isna(), 'name'].tolist()}")

# The broken invariant: the two "obvious" halves do not sum to the total.
check("high has exactly 3 rows (Ada, Dee, Gio)", len(high) == 3)
check("low has exactly 3 rows (Bo, Eli, Hu)", len(low) == 3)
check("scores has 2 rows with a missing score (Cy, Fay)", missing_count == 2)
check(
    "the naive two-way split does NOT add up: len(high) + len(low) != total",
    len(high) + len(low) != total,
)
check(
    "the shortfall is exactly the missing-value count: total - high - low == isna().sum()",
    total - len(high) - len(low) == missing_count,
)

# The fix: name the missing rows as their own group, and the three-way
# partition -- high, low, missing -- accounts for every row exactly once.
missing_rows = scores[scores.score.isna()]
check(
    "handled explicitly, the three-way partition sums to the total: high + low + missing == total",
    len(high) + len(low) + len(missing_rows) == total,
)
check(
    "no row is double-counted: the three index sets are pairwise disjoint",
    set(high.index) & set(low.index) == set()
    and set(high.index) & set(missing_rows.index) == set()
    and set(low.index) & set(missing_rows.index) == set(),
)

# A second, equally valid fix: build the "everyone else" half as the boolean
# complement of "high" rather than a second independent comparison. Because
# NaN > 50 is False, its complement ~(score > 50) is True for a NaN row --
# so the complement bucket automatically catches the missing rows too.
low_or_missing = scores[~(scores.score > 50)]
check(
    "the complement of high (~(score > 50)) automatically catches the missing rows too",
    len(high) + len(low_or_missing) == total,
)
check(
    "the complement bucket is exactly low UNION missing",
    set(low_or_missing.index) == set(low.index) | set(missing_rows.index),
)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("01_partition_invariant.py: every assertion held.")

"""Exercise 2 -- why `and` and `or` raise on a mask, and `&`/`|`/`~` do not.

Run: python3 02_and_or_raise.py

Python's `and` and `or` are control-flow keywords: they need to convert
their operand to a single True/False using __bool__, so they can decide
which branch to take. A boolean Series has no single truth value -- it is
many booleans, one per row -- so pandas refuses to guess and raises
ValueError rather than silently picking `.any()` or `.all()` for you.
`&`, `|` and `~` are ordinary operators (bitwise-and, bitwise-or, bitwise-
not), which pandas overloads to mean elementwise boolean combination; they
never need a single truth value, so they work on a whole mask at once.
"""

import numpy as np
import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


scores = pd.DataFrame(
    {
        "name": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay", "Gio", "Hu"],
        "score": [72, 45, np.nan, 91, 50, np.nan, 88, 33],
    }
)

mask1 = scores.score > 60
mask2 = scores.name.str.len() > 2
print("mask1 (score > 60):        ", mask1.tolist())
print("mask2 (len(name) > 2):     ", mask2.tolist())

try:
    combined = mask1 and mask2
    print("mask1 and mask2 did NOT raise -- result:", combined)
    raised = False
    error_text = ""
except ValueError as exc:
    raised = True
    error_text = str(exc)
    print(f"\nmask1 and mask2 raised ValueError: {error_text}")

check("`mask1 and mask2` raises ValueError", raised)
check(
    "the error names the ambiguity: mentions 'truth value' and 'ambiguous'",
    "truth value" in error_text and "ambiguous" in error_text,
)

# The `&` form works: elementwise boolean AND, one result per row.
and_mask = mask1 & mask2
print("\nmask1 & mask2 (elementwise AND):", and_mask.tolist())
check("mask1 & mask2 does not raise and returns 8 booleans", len(and_mask) == 8)
check(
    "mask1 & mask2 matches hand computation, row by row",
    and_mask.tolist()
    == [m1 and m2 for m1, m2 in zip(mask1.tolist(), mask2.tolist())],
)

or_mask = mask1 | mask2
print("mask1 | mask2 (elementwise OR): ", or_mask.tolist())
check(
    "mask1 | mask2 matches hand computation, row by row",
    or_mask.tolist() == [m1 or m2 for m1, m2 in zip(mask1.tolist(), mask2.tolist())],
)

not_mask = ~mask1
print("~mask1 (elementwise NOT):       ", not_mask.tolist())
check(
    "~mask1 is the exact elementwise negation of mask1",
    not_mask.tolist() == [not m for m in mask1.tolist()],
)

# `or` raises for the identical reason `and` does.
try:
    scores.score.gt(60) or scores.score.lt(40)
    or_raised = False
except ValueError:
    or_raised = True
check("`or` between two masks also raises ValueError", or_raised)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("02_and_or_raise.py: every assertion held.")

"""Exercise 3 -- operator precedence: `&` binds TIGHTER than comparisons.

Run: python3 03_precedence.py

`df[df.a > 1 & df.b < 2]` does not group the way it reads. In Python, `&`
binds more tightly than `<` and `>`, and `>`/`<` chain (Python rewrites
`A > B < C` as the equivalent of `(A > B) and (B < C)`). So the expression
actually parses as `df.a > (1 & df.b) < 2`, which chains through the same
`and` this lab's exercise 2 already showed raises ValueError on a Series.
The fix is unconditional: wrap every comparison in its own parentheses
before combining them with `&`, `|` or `~`.
"""

import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


table = pd.DataFrame({"a": [0, 1, 2, 3, 4], "b": [5, 3, 1, 0, -1]})
print("table:")
print(table)

# The unparenthesised form. Written the way it "reads" -- rows where a > 1
# AND b < 2 -- but that is not what it parses as.
try:
    wrong = table[table.a > 1 & table.b < 2]
    unparenthesised_raised = False
    wrong_repr = repr(wrong)
except ValueError as exc:
    unparenthesised_raised = True
    wrong_repr = str(exc)

print(f"\ntable.a > 1 & table.b < 2 -> {'ValueError: ' + wrong_repr if unparenthesised_raised else wrong_repr}")
check(
    "the unparenthesised form raises ValueError, for the same ambiguous-truth-value reason as `and`",
    unparenthesised_raised,
)
check(
    "the error is the familiar ambiguous-truth-value message, confirming it is chained `and` in disguise",
    "ambiguous" in wrong_repr,
)

# The intended query, with each comparison parenthesised.
intended_mask = (table.a > 1) & (table.b < 2)
intended = table[intended_mask]
print("\n(table.a > 1) & (table.b < 2) mask:", intended_mask.tolist())
print("(table.a > 1) & (table.b < 2) rows:")
print(intended)

check("the parenthesised mask has exactly 3 True values", intended_mask.sum() == 3)
check(
    "the parenthesised form selects rows where a in {2,3,4} and b in {1,0,-1}",
    intended["a"].tolist() == [2, 3, 4] and intended["b"].tolist() == [1, 0, -1],
)

# A second precedence trap, and the more dangerous kind: `~` also binds
# tighter than `==`, and this one does NOT raise -- it silently computes the
# wrong thing. `~table.a` is the bitwise-NOT of the integer column itself
# (~0=-1, ~1=-2, ~2=-3, ...), computed BEFORE the == 2 comparison runs, so
# `~table.a == 2` asks "which rows have bitwise-NOT(a) equal to 2", not
# "which rows have a NOT equal to 2" -- and since no value of a in this
# table has ~a == 2, the wrong query quietly returns zero rows.
wrong_tilde = table[~table.a == 2]
print(f"\n~table.a == 2 (WRONG -- parses as (~table.a) == 2): {len(wrong_tilde)} rows")
check(
    "~table.a == 2 does NOT raise -- it silently returns the wrong (empty) result",
    len(wrong_tilde) == 0,
)
correct_tilde = table[~(table.a == 2)]
print(f"~(table.a == 2)  (correct -- excludes only a == 2): {len(correct_tilde)} rows -> {correct_tilde['a'].tolist()}")
check(
    "the parenthesised form ~(table.a == 2) correctly excludes only a == 2",
    correct_tilde["a"].tolist() == [0, 1, 3, 4],
)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("03_precedence.py: every assertion held.")

"""Exercise 4 -- a mask is a Series with an index, so filtering aligns.

Run: python3 04_mask_alignment.py

A boolean mask is not a bare array of True/False; it is a Series, with an
index of its own. When you write df[mask], pandas does not walk df and mask
in lockstep -- it looks up each of df's row labels in mask's index and uses
whatever boolean sits at that LABEL, regardless of the physical order mask
happens to be stored in. Build a mask from a reordered copy of a frame,
apply it to the original, and the result is still correct, in the
original's own row order -- because alignment is by label. Convert that
same mask to a raw NumPy array first, and the label information is gone:
applying it now walks POSITION by position, and silently returns the wrong
rows, because df's own row order does not match the order the mask array
was built in.
"""

import numpy as np
import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


scores = pd.DataFrame(
    {
        "name": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay", "Gio", "Hu"],
        "score": [72, 45, np.nan, 91, 50, np.nan, 88, 33],
    },
    index=[10, 11, 12, 13, 14, 15, 16, 17],
)
print("scores (in original row order):")
print(scores)

# The correct rows, computed directly, for comparison.
expected = scores[scores.score > 60]
print(f"\nexpected rows (score > 60), computed directly: {expected.index.tolist()}")

# Build the SAME mask from a differently-ordered copy of the frame.
reordered = scores.sort_values("score", ascending=False)
print(f"\na reordered copy's row order: {reordered.index.tolist()}")
mask_from_reordered = reordered["score"] > 60
print(f"mask, stored in the reordered copy's own order: {list(zip(mask_from_reordered.index.tolist(), mask_from_reordered.tolist()))}")

aligned_result = scores[mask_from_reordered]
print(f"\nscores[mask_from_reordered] index: {aligned_result.index.tolist()}  (label-aligned)")

check(
    "a mask built from a reordered copy, applied by LABEL, still returns the correct rows",
    aligned_result.index.tolist() == expected.index.tolist(),
)
check(
    "the aligned result is identical to computing the mask directly on scores",
    aligned_result.equals(expected),
)
check(
    "the aligned result comes back in SCORES' own row order, not the mask's storage order",
    aligned_result.index.tolist() == sorted(aligned_result.index.tolist()),
)

# Now strip the labels with .to_numpy() and apply the SAME boolean values
# positionally. This is the silent disaster: the values are identical, but
# with the index gone, pandas has nothing left to align on.
positional_array = mask_from_reordered.to_numpy()
positional_result = scores[positional_array]
print(f"\nscores[mask_from_reordered.to_numpy()] index: {positional_result.index.tolist()}  (POSITIONAL, wrong)")

check(
    "applying the same booleans positionally (via .to_numpy()) gives a DIFFERENT set of rows",
    positional_result.index.tolist() != aligned_result.index.tolist(),
)
check(
    "the positional result is wrong: it does not match scores where score > 60",
    not positional_result.equals(expected),
)
check(
    "the positional result picked rows 10, 11, 12 -- the reordered copy's first three POSITIONS, not the correct labels",
    positional_result.index.tolist() == [10, 11, 12],
)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("04_mask_alignment.py: every assertion held.")

"""Exercise 5 -- .str.contains and missing values: a trap that depends on
dtype, and pandas 3.0 changed which dtype you get by default.

Run: python3 05_str_contains_na.py

On the legacy `object` dtype, `.str.contains()` applied to a missing entry
returns `None` rather than `False`, because there is no text to search.
The resulting mask is not a clean boolean array -- its own dtype becomes
`object` -- and filtering a DataFrame with a mask that contains missing
values raises `ValueError: Cannot mask with non-boolean array containing
NA / NaN values` rather than silently doing the wrong thing. `na=False`
tells `.str.contains()` to treat a missing entry as "did not match" up
front, producing a clean boolean mask.

Pandas 3.0's new default `str` dtype (backed by PyArrow) behaves
differently, and this script measures that difference rather than
asserting the old story blindly: on a `str`-dtype column, `.str.contains()`
already returns a plain `False` for a missing entry, with no `NaN` in the
mask at all. The trap has NOT disappeared -- it still fires on `object`
dtype, which is still common (an explicit dtype="object" column, or a
column that arrived that way from elsewhere) -- but it no longer fires by
default on a plain list of Python strings under pandas 3.0.
"""

import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


raw_names = ["Alice Smith", "bob jones", None, "CAROL", "dave"]

# --- The pandas-3.0 default: str dtype -------------------------------
names_str = pd.Series(raw_names, dtype="str")
print(f"names_str.dtype: {names_str.dtype}")
mask_str = names_str.str.contains("a", case=False)
print(f"mask (str dtype):    {mask_str.tolist()}  dtype={mask_str.dtype}")

check("on pandas 3.0's default str dtype, the mask dtype is a clean bool", mask_str.dtype == bool)
check(
    "on str dtype, the missing entry (index 2) already comes back False, not NaN",
    bool(mask_str.iloc[2]) is False and not mask_str.isna().any(),
)
filtered_str = names_str[mask_str]
print(f"names_str[mask_str]: {filtered_str.tolist()}")
check(
    "filtering with the str-dtype mask works directly, no na= needed",
    filtered_str.tolist() == ["Alice Smith", "CAROL", "dave"],
)

# --- The trap: object dtype, still common in practice -----------------
names_obj = pd.Series(raw_names, dtype="object")
print(f"\nnames_obj.dtype: {names_obj.dtype}")
mask_obj = names_obj.str.contains("a", case=False)
print(f"mask (object dtype): {mask_obj.tolist()}  dtype={mask_obj.dtype}")

check("on object dtype, the missing entry's mask value is None, not False", mask_obj.iloc[2] is None)
check("on object dtype, the mask's own dtype is object, not bool", mask_obj.dtype == object)

try:
    _ = names_obj[mask_obj]
    contains_trap_raised = False
    trap_message = ""
except ValueError as exc:
    contains_trap_raised = True
    trap_message = str(exc)
print(f"names_obj[mask_obj] -> {'ValueError: ' + trap_message if contains_trap_raised else 'no error'}")

check("filtering with an object-dtype mask containing None raises ValueError", contains_trap_raised)
check(
    "the error names the real cause: masking with non-boolean / NA values",
    "non-boolean" in trap_message or "NA" in trap_message,
)

# The fix: na=False.
mask_obj_fixed = names_obj.str.contains("a", case=False, na=False)
print(f"\nmask (object dtype, na=False): {mask_obj_fixed.tolist()}  dtype={mask_obj_fixed.dtype}")
filtered_obj = names_obj[mask_obj_fixed]
print(f"names_obj[mask_obj_fixed]: {filtered_obj.tolist()}")

check("na=False produces a clean boolean mask on object dtype", mask_obj_fixed.dtype == bool)
check(
    "na=False filters correctly, matching the str-dtype result exactly",
    filtered_obj.tolist() == ["Alice Smith", "CAROL", "dave"],
)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("05_str_contains_na.py: every assertion held.")

"""Exercise 6 -- .query() versus the equivalent mask.

Run: python3 06_query_equivalence.py

`.query()` takes a string of Python-like expression syntax and evaluates
it against the DataFrame's own columns as if they were local variables,
returning exactly the same rows a hand-built boolean mask would. Its real
advantage is readability once several conditions stack up -- no repeated
`df.` prefix, no risk of the `&`/`<` precedence trap from exercise 3,
because `.query()` parses comparison and boolean-combination the way plain
Python reads. A value from outside the frame is referenced with an `@`
prefix. The honest cost: `.query()` builds and parses a small string at
every call, which is measurably slower than a mask for a single simple
condition, and it turns what used to be a static-analysis-friendly Python
expression into a string your editor cannot type-check. For one condition
on a small frame, a mask is simpler; once you are combining four or five
conditions with named thresholds, `.query()` usually reads better.
"""

import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


orders = pd.DataFrame(
    {
        "customer": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay"],
        "amount": [42.50, 108.00, 15.75, 220.10, 60.00, 9.99],
        "region": ["east", "west", "east", "west", "east", "west"],
    }
)
print("orders:")
print(orders)

# A single condition: mask and query give identical rows.
threshold = 50
mask_single = orders[orders.amount > threshold]
query_single = orders.query("amount > @threshold")
print(f"\nmask (amount > {threshold}):  {mask_single.customer.tolist()}")
print(f"query (amount > @threshold): {query_single.customer.tolist()}")

check("mask and .query() select the identical rows for a single condition", mask_single.equals(query_single))
check(
    "the @threshold syntax correctly reaches the Python variable, not a column named 'threshold'",
    query_single.customer.tolist() == ["Bo", "Dee", "Eli"],
)

# A compound condition: the case .query() is genuinely nicer for, and the
# case exercise 3's precedence trap would bite on if written unparenthesised.
mask_compound = orders[(orders.amount > threshold) & (orders.region == "east")]
query_compound = orders.query("amount > @threshold and region == 'east'")
print(f"\nmask (amount > {threshold} & region == east):  {mask_compound.customer.tolist()}")
print(f"query (amount > @threshold and region == 'east'): {query_compound.customer.tolist()}")

check(
    "mask and .query() select the identical rows for a compound condition",
    mask_compound.equals(query_compound),
)
check("the compound condition selects exactly Eli", query_compound.customer.tolist() == ["Eli"])

# .query() parses `and`/`or` in the query string the way plain Python reads
# them, with none of exercise 3's `&`-binds-tighter surprise -- because the
# string is parsed by pandas' own expression engine, not by Python's own
# operator-precedence table applied to Series objects.
query_no_parens_needed = orders.query("amount > @threshold and region == 'east'")
check(
    "inside a .query() string, 'and' works directly -- no precedence trap, unlike exercise 3's `&`",
    query_no_parens_needed.equals(query_compound),
)

# A second @variable, referencing a Python list, combined with isin.
wanted_regions = ["west"]
query_isin = orders.query("region in @wanted_regions")
mask_isin = orders[orders.region.isin(wanted_regions)]
print(f"\nquery (region in @wanted_regions): {query_isin.customer.tolist()}")
check("`in @variable` inside .query() matches .isin() exactly", query_isin.equals(mask_isin))

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("06_query_equivalence.py: every assertion held.")

"""Exercise 7 -- .isin() versus chained ==, and the empty-list trap.

Run: python3 07_isin_vs_chained.py

`series.isin(values)` and a chain of `(series == v1) | (series == v2) |
...` compute the identical boolean mask -- `.isin()` is simply the version
that scales to any number of values without writing one `==`/`|` pair per
value, and reads as "is this value one of these" rather than a wall of
`|`. The one behaviour worth knowing on purpose: `.isin([])` -- an empty
list of wanted values -- returns an all-False mask, and filtering with it
gives an EMPTY frame, not the original untouched frame. It is easy to
assume "no filter values given" means "no filter applied"; pandas disagrees,
correctly, because "is this row's value among these zero values" can only
ever be false.
"""

import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


staff = pd.DataFrame(
    {
        "name": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay", "Gio", "Hu"],
        "dept": ["eng", "sales", "eng", "hr", "sales", "eng", "hr", "sales"],
    }
)
print("staff:")
print(staff)

wanted = ["eng", "hr"]
via_isin = staff[staff.dept.isin(wanted)]
via_chained = staff[(staff.dept == "eng") | (staff.dept == "hr")]
print(f"\nisin(['eng', 'hr']):                 {via_isin.name.tolist()}")
print(f"(dept == 'eng') | (dept == 'hr'):    {via_chained.name.tolist()}")

check("isin() and the chained == form select the identical rows", via_isin.equals(via_chained))
check("both forms select exactly the 5 eng/hr staff", len(via_isin) == 5)

# Three-value case: isin scales with no extra | per value; the chained form
# needs one more == and one more | for every value added.
wanted3 = ["eng", "hr", "sales"]
via_isin3 = staff[staff.dept.isin(wanted3)]
via_chained3 = staff[(staff.dept == "eng") | (staff.dept == "hr") | (staff.dept == "sales")]
check("with a third value, isin() and the chained form still agree exactly", via_isin3.equals(via_chained3))
check("isin() with all departments listed returns the whole frame", len(via_isin3) == len(staff))

# The empty-list trap.
empty_wanted: list[str] = []
via_empty = staff[staff.dept.isin(empty_wanted)]
print(f"\nisin([]) -- an empty wanted list -- rows returned: {len(via_empty)}")

check("isin([]) returns zero rows, NOT the whole untouched frame", len(via_empty) == 0)
check("isin([]) produces an all-False mask, one entry per row", (~staff.dept.isin(empty_wanted)).all())
check(
    "isin([]) does NOT mean 'no filter' -- it means 'exclude everything', the opposite intuition",
    len(via_empty) != len(staff),
)

# The corresponding negation, ~isin, correctly means "none of these" and
# with an empty list correctly keeps everything -- worth contrasting.
via_not_in_empty = staff[~staff.dept.isin(empty_wanted)]
print(f"~isin([]) -- 'is dept NOT one of these zero values' -- rows returned: {len(via_not_in_empty)}")
check("~isin([]) DOES return every row, since nothing is excluded by an empty exclusion list", len(via_not_in_empty) == len(staff))

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("07_isin_vs_chained.py: every assertion held.")

"""Exercise 8 -- .nlargest()/.nsmallest() versus .sort_values().head().

Run: python3 08_nlargest_vs_sort_head.py

`df.nlargest(n, col)` and `df.sort_values(col, ascending=False).head(n)`
answer the same question -- the top n rows by one column -- and when there
are no ties at the cutoff they return byte-for-byte identical rows.
`.nlargest()` is the cheaper way to ask: it never sorts the whole frame,
it maintains a running set of the n largest values seen so far (an
O(n log k) approach for k rows kept, versus sort_values' full O(n log n)
sort of every row), which matters once the frame is large and n is small.
The behaviours diverge exactly at a tie sitting on the cutoff:
`.sort_values().head(n)` always returns exactly n rows, arbitrarily
keeping some tied rows and dropping others; `.nlargest(n, col, keep='all')`
can return MORE than n rows on purpose, returning every row tied at the
boundary rather than picking an arbitrary subset of them.
"""

import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


# --- No ties at the cutoff: both approaches agree exactly. -------------
scores = pd.DataFrame(
    {
        "name": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay", "Gio", "Hu"],
        "score": [72, 45, 20, 91, 50, 15, 88, 33],
    }
)
print("scores (no ties):")
print(scores)

top_nlargest = scores.nlargest(3, "score")
top_sorted_head = scores.sort_values("score", ascending=False).head(3)
print(f"\n.nlargest(3, 'score'):                       {top_nlargest.name.tolist()}")
print(f".sort_values('score', ascending=False).head(3): {top_sorted_head.name.tolist()}")

check("with no ties at the cutoff, nlargest and sort_values().head() give identical rows", top_nlargest.equals(top_sorted_head))
check("the top 3 by score are Dee (91), Gio (88), Ada (72)", top_nlargest.name.tolist() == ["Dee", "Gio", "Ada"])

bottom_nsmallest = scores.nsmallest(2, "score")
bottom_sorted_head = scores.sort_values("score", ascending=True).head(2)
check(
    "nsmallest and sort_values(ascending=True).head() also agree with no ties",
    bottom_nsmallest.equals(bottom_sorted_head),
)

# --- Ties sitting exactly on the cutoff: the two approaches diverge. ----
tied = pd.DataFrame({"name": ["A", "B", "C", "D", "E"], "score": [90, 90, 85, 80, 80]})
print("\nscores with a tie AT the cutoff (D and E both score 80):")
print(tied)

tied_nlargest_default = tied.nlargest(4, "score")
tied_sorted_head = tied.sort_values("score", ascending=False).head(4)
print(f"\n.nlargest(4, 'score') [keep='first' default]: {tied_nlargest_default.name.tolist()}")
print(f".sort_values(ascending=False).head(4):         {tied_sorted_head.name.tolist()}")

check(
    "with keep='first' (the default), nlargest(4) still matches sort_values().head(4) exactly -- both pick D over E",
    tied_nlargest_default.equals(tied_sorted_head),
)
check(".sort_values().head(n) always returns EXACTLY n rows, arbitrarily choosing among ties", len(tied_sorted_head) == 4)

tied_nlargest_all = tied.nlargest(4, "score", keep="all")
print(f"\n.nlargest(4, 'score', keep='all'):             {tied_nlargest_all.name.tolist()}")

check(
    "keep='all' returns MORE than n rows when a tie sits on the cutoff -- every tied row, not an arbitrary subset",
    len(tied_nlargest_all) == 5,
)
check(
    "keep='all' includes BOTH D and E, the tied pair sort_values().head() had to arbitrarily choose between",
    set(tied_nlargest_all.name.tolist()) == {"A", "B", "C", "D", "E"},
)
check(
    "sort_values().head(n) has no equivalent to keep='all' -- .head(4) can never return 5 rows",
    len(tied.sort_values("score", ascending=False).head(4)) == 4,
)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("08_nlargest_vs_sort_head.py: every assertion held.")

"""Exercise 9 -- .drop_duplicates() with subset and keep, and a note on
.filter(), which selects LABELS, not rows -- one of the library's more
confusingly named methods.

Run: python3 09_drop_duplicates_and_filter.py

"Duplicate" is not a fixed property of a row; it means whatever columns you
name in `subset`. The same table has 4 whole-row-unique rows, 4 rows unique
by (customer, item), and only 3 rows unique by customer alone -- three
different, all-correct answers to "how many duplicates", because they
answer three different questions. `keep='first'`/`'last'` decide which of a
duplicate group survives; the row count is identical either way, only WHICH
row differs.

`.filter()` looks like a row filter and is not one: it selects COLUMN (or
row, with axis=0) LABELS by exact name, a `like=` substring, or a `regex=`
pattern -- never by a condition on the data inside them. Passing it
row-shaped arguments does not raise; it silently matches nothing and
returns an empty selection on that axis, which is exactly the trap.
"""

import pandas as pd

checks = 0
failures = 0


def check(label, condition):
    global checks, failures
    checks += 1
    if condition:
        print(f"  ok: {label}")
    else:
        print(f"  FAIL: {label}")
        failures += 1


orders = pd.DataFrame(
    {
        "customer": ["Ada", "Bo", "Ada", "Cy", "Bo", "Ada"],
        "item": ["pen", "cup", "pen", "pen", "cup", "mug"],
        "qty": [1, 2, 1, 5, 2, 3],
    }
)
print("orders:")
print(orders)

# Whole-row duplicates: rows 0 and 2 are byte-for-byte identical.
whole_row = orders.drop_duplicates()
print(f"\ndrop_duplicates() [whole row]: kept {len(whole_row)} of {len(orders)} rows -> index {whole_row.index.tolist()}")
check("dropping whole-row duplicates keeps exactly 4 rows (row 2 is identical to row 0)", len(whole_row) == 4)
check("the surviving rows are 0, 1, 3, 5", whole_row.index.tolist() == [0, 1, 3, 5])

# Duplicate by (customer, item): rows 0/2 collide (Ada, pen) and rows 1/4
# collide (Bo, cup) -- two collisions this time, not one.
subset_ci_first = orders.drop_duplicates(subset=["customer", "item"], keep="first")
subset_ci_last = orders.drop_duplicates(subset=["customer", "item"], keep="last")
print(f"\ndrop_duplicates(subset=['customer','item'], keep='first'): index {subset_ci_first.index.tolist()}")
print(f"drop_duplicates(subset=['customer','item'], keep='last'):  index {subset_ci_last.index.tolist()}")

check("by (customer, item), 4 rows survive either way -- the COUNT does not depend on keep", len(subset_ci_first) == 4 and len(subset_ci_last) == 4)
check("keep='first' keeps the FIRST occurrence of each (customer, item) pair: rows 0, 1, 3, 5", subset_ci_first.index.tolist() == [0, 1, 3, 5])
check("keep='last' keeps the LAST occurrence instead: rows 2, 3, 4, 5", subset_ci_last.index.tolist() == [2, 3, 4, 5])
check(
    "first and last keep DIFFERENT rows for the same duplicate group -- (Ada, pen) keeps qty=1 either way here by coincidence, but the surviving ROW differs",
    subset_ci_first.index.tolist() != subset_ci_last.index.tolist(),
)

# Duplicate by customer alone: a stricter subset finds MORE duplicates,
# because it ignores what the customer actually ordered.
subset_customer = orders.drop_duplicates(subset=["customer"], keep="first")
print(f"\ndrop_duplicates(subset=['customer'], keep='first'): index {subset_customer.index.tolist()}")
check("by customer alone, only 3 rows survive -- one per distinct customer", len(subset_customer) == 3)
check("'duplicate' is not fixed: the SAME table gives 4, 4, and 3 depending on subset chosen", len(whole_row) != len(subset_customer))

# --- .filter(): selects LABELS, never a condition on the data. ---------
print("\n--- .filter() is not a row filter ---")
by_items = orders.filter(items=["customer", "qty"])
print(f"orders.filter(items=['customer', 'qty']).columns: {by_items.columns.tolist()}")
check(".filter(items=...) selects COLUMNS by exact name, unrelated to any row condition", by_items.columns.tolist() == ["customer", "qty"])
check(".filter() does not drop any rows -- same row count as the original", len(by_items) == len(orders))

by_like = orders.filter(like="qty")
print(f"orders.filter(like='qty').columns:            {by_like.columns.tolist()}")
check(".filter(like=...) matches columns by substring, still label-based", by_like.columns.tolist() == ["qty"])

# The trap: passing row-shaped intentions to .filter() does not raise --
# it silently matches nothing on the columns axis, keeping every row.
looks_like_a_row_filter = orders.filter(items=[0, 1, 2])
print(f"orders.filter(items=[0, 1, 2]) -- looks like 'keep rows 0,1,2', is NOT: columns={looks_like_a_row_filter.columns.tolist()}, rows kept={len(looks_like_a_row_filter)}")
check(
    "filter(items=[0,1,2]) does NOT keep rows 0-2 -- it looks for COLUMNS named 0, 1, 2, finds none, and keeps every row with zero columns",
    looks_like_a_row_filter.shape == (len(orders), 0),
)

print(f"\n{checks} checks, {failures} failure(s).")
if failures:
    raise SystemExit(1)
print("09_drop_duplicates_and_filter.py: every assertion held.")

lesson_id: D122
day: 122
kind: guided-build
languages: [python, bash]
setup_commands:
  - cd labs/sections/math-statistics-and-data/day-122-selecting-and-filtering
  - python3 -m venv .venv
  - .venv/bin/pip install -r requirements/requirements.txt
  - .venv/bin/python3 -c "import pandas; print(pandas.__version__)"
run_commands:
  - 'cd examples && ../.venv/bin/python3 01_partition_invariant.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 02_and_or_raise.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 03_precedence.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 04_mask_alignment.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 05_str_contains_na.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 06_query_equivalence.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 07_isin_vs_chained.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 08_nlargest_vs_sort_head.py && cd ..'
  - 'cd examples && ../.venv/bin/python3 09_drop_duplicates_and_filter.py && cd ..'
  - .venv/bin/python3 starter/check_progress.py
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, bash 3.2.57 -- bash tests/run_tests.sh -> 41 checks, 0 failure(s), exit 0. All nine reference scripts in examples/ exit 0 with every internal assertion holding, run individually and via the harness. starter/check_progress.py reports 0 of 9 exercises complete on the untouched checkout (exit 1), and the harness separately solves every blank in a scratch copy and confirms 9 of 9 complete with exit 0, without modifying the real starter/exercises.py on disk. Section 5 of the harness deliberately breaks the assertion inside 07_isin_vs_chained.py (expects 999 rows instead of the real 5), confirms the run exits non-zero with a printed FAIL line, restores the file, and confirms it passes again -- so the suite is demonstrated to be capable of failing rather than merely claimed to be. Everything was run through a real lab-local .venv created by the documented setup commands. Two honesty notes from this run. FIRST: exercise 5''s original brief framing assumed .str.contains() on a missing value always returns NaN regardless of dtype. Measured directly in this session, that is only true on the legacy object dtype; on pandas 3.0.5''s own default str dtype for a plain list of Python strings, .str.contains() on a missing entry already returns a clean False with no NaN in the mask at all, so the classic trap does not fire by default under 3.0 -- it still fires, and is demonstrated, on object dtype, which remains common. This is reported as measured rather than assumed, and the lesson corrects the brief''s framing rather than repeating it. SECOND: applying a boolean mask to a DataFrame whose index order differs from the mask''s own storage order raises a real UserWarning (Boolean Series key will be reindexed to match DataFrame index) on pandas 3.0.5 -- informational, not an error, and the aligned result is still correct -- captured verbatim in expected-output/04-mask-alignment.txt with its file path sanitized to <repo>. matplotlib, scipy and polars are not installed in this environment; polars is described from its public documentation in the lesson''s Tools section as a design contrast (no implicit index, so its filter() expression never hits the & precedence trap this lab''s exercise 3 demonstrates) and no output attributed to it is reproduced anywhere.'

# What is installed, why, and what it costs

Three 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 mask, `.loc` call, `.query()` and `.str.contains()` in this lab. Pinned exactly because two exercises depend on pandas-3.0-specific dtype behaviour — see below. |
| `pyarrow` | 25.0.1 | Apache 2.0 | The storage backend behind the pandas 3.0 default `str` dtype exercised in exercise 5. |
| `numpy` | 2.5.2 | BSD 3-Clause | `np.nan`, boolean-array semantics, and the arrays every mask is ultimately built from. |

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

## Why the versions are pinned exactly, not just floored

`exercise 5` (`.str.contains` on a column with missing values) prints a
genuinely different result depending on the column's dtype, and that dtype
default changed in pandas 3.0:

- On a **pandas-3.0 `str` dtype** column (the default this pandas version
  gives a list of Python strings), `.str.contains()` on a missing entry
  returns a plain `False`, and the resulting mask's own dtype is `bool`
  with no missing values in it at all — the classic trap does not fire.
- On the legacy **`object` dtype** column — still reachable with
  `dtype="object"`, and still what you get from many real-world sources —
  `.str.contains()` on a missing entry returns `None`, the mask's dtype is
  `object`, and filtering a DataFrame with that mask raises
  `ValueError: Cannot mask with non-boolean array containing NA / NaN
  values` rather than silently doing the wrong thing.

Both facts are captured from this exact pandas version and are stated
plainly as version-specific in `expected-output/FIELDS.md`. A different
pandas major version could print a different combination of these two
behaviours.

## 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. Section 6 of
`tests/run_tests.sh` greps every source file in `examples/` and `starter/`
to prove that nothing else does.

## What is deliberately *not* installed

**matplotlib**, **scipy** and **polars** are not installed in this
environment. The lesson's Tools section describes polars from its public
documentation as a design contrast to pandas' masks — specifically, that
`pl.col('a') > 1` composes inside one expression rather than through
Python's `&`/`|`/`~` operators, which removes the precedence trap this
lab's exercise 3 demonstrates by construction. **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 — the whole point is pandas' specific masking
and alignment behaviour, which nothing else on your system will reproduce.
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

"""Run from this directory (or `python3 starter/check_progress.py` from the
lab root): reports how many of the nine exercises in exercises.py are
complete and correct, against the same expected values the reference
examples/ scripts assert.

Exit code is 0 only when all nine are correct, matching the convention the
rest of this lab's tests use.
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

import exercises as ex

passed = 0
total = 9


def report(number, description, fn, expected_check):
    """Run one exercise function, catching the NameError an unfilled blank
    raises, and check its return value against expected_check(value)."""
    global passed
    try:
        result = fn()
    except NameError as exc:
        print(f"  {number}. {description}: NOT YET COMPLETE ({exc})")
        return
    except Exception as exc:  # noqa: BLE001 -- report any other mistake too
        print(f"  {number}. {description}: ERROR -- {exc!r}")
        return
    ok = expected_check(result)
    if ok:
        passed += 1
        print(f"  {number}. {description}: correct  (got {result!r})")
    else:
        print(f"  {number}. {description}: WRONG    (got {result!r})")


report(
    1,
    "partition invariant",
    ex.ex01_partition_invariant,
    lambda r: r == (3, 3, 8, 2),
)
report(
    2,
    "and/or raise",
    ex.ex02_and_or_raise,
    lambda r: r[0] is True and r[1] == [True, False, False, False],
)
report(3, "precedence", ex.ex03_precedence, lambda r: r == [False, False, True, True, True])
report(4, "mask alignment", ex.ex04_mask_alignment, lambda r: r == [10, 12])
report(
    5,
    "str.contains na",
    ex.ex05_str_contains_na,
    lambda r: r == ["Alice Smith", "dave"],
)
report(6, "query equivalence", ex.ex06_query_equivalence, lambda r: r == ["Bo"])
report(7, "isin empty", ex.ex07_isin_empty, lambda r: r == 0)
report(8, "nlargest ties", ex.ex08_nlargest_ties, lambda r: r == 3)
report(
    9,
    "drop_duplicates subset",
    ex.ex09_drop_duplicates_subset,
    lambda r: r == ["Ada", "Bo"],
)

print(f"\n{passed} of {total} exercises complete.")
sys.exit(0 if passed == total else 1)

"""Day 122 starter -- nine exercises, one function each.

Each function below is a working skeleton: the setup is written for you, and
exactly one line is left for you to write, marked with the sentinel name
`_FILL_THIS_IN`. Replace that name with real code. Leaving it as-is raises a
clear NameError when the function runs, which `check_progress.py` catches
and reports -- it does not crash the whole script.

Read `../examples/` for the fully worked reference AFTER you have tried each
one yourself; that is where every one of these ideas is explained in
comments. Run your progress with:

    python3 check_progress.py

from inside this `starter/` directory (or `python3 starter/check_progress.py`
from the lab root).
"""

import numpy as np
import pandas as pd


def ex01_partition_invariant():
    """On a score column with two NaN entries, high = score > 50 and
    low = score <= 50 do not sum to the total. Return
    (len(high), len(low), len(scores), missing_count)."""
    scores = pd.DataFrame(
        {
            "name": ["Ada", "Bo", "Cy", "Dee", "Eli", "Fay", "Gio", "Hu"],
            "score": [72, 45, np.nan, 91, 50, np.nan, 88, 33],
        }
    )
    high = scores[scores.score > 50]
    low = scores[scores.score <= 50]
    missing_count = _FILL_THIS_IN  # count of rows where score is NaN -- use .isna().sum()
    return len(high), len(low), len(scores), missing_count


def ex02_and_or_raise():
    """mask1 and mask2 should raise ValueError; mask1 & mask2 should not.
    Return (raised_bool, and_mask_as_list)."""
    scores = pd.Series([72, 45, np.nan, 91], name="score")
    mask1 = scores > 60
    mask2 = scores < 90
    try:
        mask1 and mask2
        raised = False
    except ValueError:
        raised = True
    and_mask = _FILL_THIS_IN  # the elementwise-AND version, using & not `and`
    return raised, and_mask.tolist()


def ex03_precedence():
    """table.a > 1 & table.b < 2 does not mean what it looks like. Return
    the CORRECT mask (as a list) for 'a > 1 AND b < 2', properly
    parenthesised."""
    table = pd.DataFrame({"a": [0, 1, 2, 3, 4], "b": [5, 3, 1, 0, -1]})
    correct_mask = _FILL_THIS_IN  # (table.a > 1) & (table.b < 2), correctly parenthesised
    return correct_mask.tolist()


def ex04_mask_alignment():
    """Build a mask from scores sorted by score descending, then apply it
    to the ORIGINAL (unsorted) scores. Return the resulting index as a
    list -- it should come back in the ORIGINAL frame's row order."""
    scores = pd.DataFrame(
        {"name": ["Ada", "Bo", "Cy", "Dee"], "score": [72, 45, 91, 33]},
        index=[10, 11, 12, 13],
    )
    reordered = scores.sort_values("score", ascending=False)
    mask = reordered["score"] > 50
    result = _FILL_THIS_IN  # apply `mask` to the ORIGINAL `scores`, not `reordered`
    return result.index.tolist()


def ex05_str_contains_na():
    """On an object-dtype Series with a missing entry, .str.contains(...)
    without na= raises when used to filter. Fix it with na=False. Return
    the filtered list of names."""
    names = pd.Series(["Alice Smith", "bob jones", None, "dave"], dtype="object")
    mask = _FILL_THIS_IN  # names.str.contains("a", case=False, na=False)
    return names[mask].tolist()


def ex06_query_equivalence():
    """Select rows where amount > threshold using .query() with an
    @variable. Return the customer list."""
    orders = pd.DataFrame(
        {"customer": ["Ada", "Bo", "Cy"], "amount": [42.5, 108.0, 15.75]}
    )
    threshold = 50
    result = _FILL_THIS_IN  # orders.query("amount > @threshold")
    return result.customer.tolist()


def ex07_isin_empty():
    """isin() with an empty list of wanted values returns zero rows, not
    the whole frame. Return the row count."""
    staff = pd.DataFrame({"dept": ["eng", "sales", "hr"]})
    empty_wanted: list[str] = []
    result = _FILL_THIS_IN  # staff[staff.dept.isin(empty_wanted)]
    return len(result)


def ex08_nlargest_ties():
    """Three rows (A, B, C) all tie for the top score, but n=2 asks for
    only the top 2. nlargest(2, 'score', keep='all') should return ALL
    THREE tied rows -- more than n. Return the row count."""
    tied = pd.DataFrame({"name": ["A", "B", "C", "D"], "score": [80, 80, 80, 60]})
    result = _FILL_THIS_IN  # tied.nlargest(2, "score", keep="all")
    return len(result)


def ex09_drop_duplicates_subset():
    """Two rows share the same customer even though they ordered
    different items. Drop duplicates by customer alone, keeping the
    first occurrence. Return the surviving customer list, in order."""
    orders = pd.DataFrame(
        {"customer": ["Ada", "Bo", "Ada"], "item": ["pen", "cup", "mug"]}
    )
    result = _FILL_THIS_IN  # orders.drop_duplicates(subset=["customer"], keep="first")
    return result.customer.tolist()

#!/usr/bin/env bash
# Tests for the Day 122 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:
#
#   * a naive two-way split of a column with missing values (score > 50,
#     score <= 50) does NOT add up to the total row count -- the shortfall
#     is exactly the count of missing rows, and a three-way partition
#     (high, low, missing) restores the invariant exactly;
#   * `mask1 and mask2` raises ValueError (ambiguous truth value of a
#     Series), while `mask1 & mask2` computes the elementwise AND with no
#     error;
#   * `df.a > 1 & df.b < 2` does not group the way it reads -- `&` binds
#     tighter than the comparisons -- and raises the same ValueError;
#     `~df.a == 2` silently computes the WRONG (not erroring) result for
#     the same reason;
#   * a boolean mask built from a reordered copy of a frame, applied to the
#     original, aligns by LABEL and returns the correct rows in the
#     original's own order; the same booleans applied positionally (via
#     .to_numpy()) return a different, wrong set of rows;
#   * .str.contains() on a missing value returns None on object dtype
#     (raising ValueError when used to filter) but a clean False on
#     pandas 3.0's default str dtype -- na=False fixes both;
#   * .query() with an @variable selects the identical rows to the
#     equivalent mask, for both a single and a compound condition;
#   * .isin() matches a chain of == / | exactly, and .isin([]) returns
#     zero rows rather than the whole frame;
#   * .nlargest()/.nsmallest() match .sort_values().head() exactly when
#     there is no tie at the cutoff, and keep='all' returns MORE than n
#     rows when there is;
#   * .drop_duplicates() with different `subset` values gives different,
#     all-correct answers to "how many duplicates", and .filter() selects
#     labels, never rows, silently matching nothing if given row-shaped
#     arguments;
#   * 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 `python3 starter/check_progress.py`, and running it 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 python: 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
}

python_bin="$(resolve_tool python3 "${PYTHON:-}")" || {
  echo "FAIL: python3 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 python3:" >&2
  echo "    PYTHON=/path/to/python3 bash tests/run_tests.sh" >&2
  exit 1
}

if ! "${python_bin}" -c "import pandas, pyarrow" >/dev/null 2>&1; then
  echo "FAIL: pandas/pyarrow are 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 122 — Filters That Add Up"
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"):
    print(f"{name:<8} {version(name)}")
print(f"platform {platform.platform()}")
print(f"exe      {sys.executable.rsplit('/', 3)[-1]}")
PY
)"
echo "${versions}" | sed 's/^/  /'

pinned_pandas="$(grep -E '^pandas==' "${lab_dir}/requirements/requirements.txt" | cut -d= -f3)"
installed_pandas="$("${python_bin}" -c "from importlib.metadata import version; print(version('pandas'))")"
check_eq "installed pandas matches requirements.txt" "${pinned_pandas}" "${installed_pandas}"

pandas_major="$("${python_bin}" -c "import pandas; print(pandas.__version__.split('.')[0])")"
check_eq "pandas is version 3 or later (this lab's captured output is 3.0.5-specific)" "3" "${pandas_major}"

# --------------------------------------------------------------------------
echo
echo "2. Every reference script runs and every assertion inside it holds"
# --------------------------------------------------------------------------

for script in 01_partition_invariant 02_and_or_raise 03_precedence \
              04_mask_alignment 05_str_contains_na 06_query_equivalence \
              07_isin_vs_chained 08_nlargest_vs_sort_head 09_drop_duplicates_and_filter; do
  out="$(cd "${lab_dir}/examples" && "${python_bin}" "${script}.py" 2>&1)"
  status=$?
  if [ "${status}" -ne 0 ]; then
    check "${script}.py exits 0" "no"
    echo "${out}" | tail -8 | sed 's/^/      /'
  else
    check "${script}.py exits 0" "yes"
  fi
  case "${out}" in
    *"${script}.py: every assertion held."*)
      check "${script}.py reports every assertion held" "yes" ;;
    *) check "${script}.py reports every assertion held" "no" ;;
  esac
done

# --------------------------------------------------------------------------
echo
echo "3. The starter checker: honest progress, both directions"
# --------------------------------------------------------------------------

starter_out="$(cd "${lab_dir}" && "${python_bin}" starter/check_progress.py 2>&1)"
starter_status=$?
echo "${starter_out}" | tail -3 | sed 's/^/  /'
case "${starter_out}" in
  *"0 of 9 exercises complete."*)
    check "an untouched starter checkout reports 0 of 9 complete" "yes" ;;
  *)
    check "an untouched starter checkout reports 0 of 9 complete" "no" ;;
esac
if [ "${starter_status}" -ne 0 ]; then
  check "the starter checker exits non-zero when incomplete" "yes"
else
  check "the starter checker exits non-zero when incomplete" "no"
fi

# Prove the checker can also report success: solve every exercise in a
# scratch copy, confirm 9 of 9 and exit 0, then discard the copy. The real
# starter/exercises.py on disk is never modified by this.
solved_dir="$(mktemp -d)"
trap 'rm -rf "${solved_dir}"' EXIT
cp "${lab_dir}/starter/exercises.py" "${solved_dir}/exercises.py"
cp "${lab_dir}/starter/check_progress.py" "${solved_dir}/check_progress.py"
"${python_bin}" - "${solved_dir}/exercises.py" <<'PY'
import sys
path = sys.argv[1]
s = open(path).read()
replacements = [
    ('missing_count = _FILL_THIS_IN  # count of rows where score is NaN -- use .isna().sum()', 'missing_count = int(scores.score.isna().sum())'),
    ('and_mask = _FILL_THIS_IN  # the elementwise-AND version, using & not `and`', 'and_mask = mask1 & mask2'),
    ('correct_mask = _FILL_THIS_IN  # (table.a > 1) & (table.b < 2), correctly parenthesised', 'correct_mask = (table.a > 1) & (table.b < 2)'),
    ('result = _FILL_THIS_IN  # apply `mask` to the ORIGINAL `scores`, not `reordered`', 'result = scores[mask]'),
    ('mask = _FILL_THIS_IN  # names.str.contains("a", case=False, na=False)', 'mask = names.str.contains("a", case=False, na=False)'),
    ('result = _FILL_THIS_IN  # orders.query("amount > @threshold")', 'result = orders.query("amount > @threshold")'),
    ('result = _FILL_THIS_IN  # staff[staff.dept.isin(empty_wanted)]', 'result = staff[staff.dept.isin(empty_wanted)]'),
    ('result = _FILL_THIS_IN  # tied.nlargest(2, "score", keep="all")', 'result = tied.nlargest(2, "score", keep="all")'),
    ('result = _FILL_THIS_IN  # orders.drop_duplicates(subset=["customer"], keep="first")', 'result = orders.drop_duplicates(subset=["customer"], keep="first")'),
]
for old, new in replacements:
    if old not in s:
        raise SystemExit(f"pattern not found, starter/exercises.py has drifted: {old!r}")
    s = s.replace(old, new)
open(path, "w").write(s)
PY
solved_out="$(cd "${solved_dir}" && "${python_bin}" check_progress.py 2>&1)"
solved_status=$?
echo "${solved_out}" | tail -3 | sed 's/^/  /'
case "${solved_out}" in
  *"9 of 9 exercises complete."*)
    check "a fully solved copy reports 9 of 9 complete" "yes" ;;
  *)
    check "a fully solved copy reports 9 of 9 complete" "no" ;;
esac
check_eq "the checker exits 0 once every exercise is correct" "0" "${solved_status}"
rm -rf "${solved_dir}"
trap - EXIT

# --------------------------------------------------------------------------
echo
echo "4. The lesson's sharpest claims, checked one value at a time"
# --------------------------------------------------------------------------

facts="$(cd "${lab_dir}/examples" && "${python_bin}" - <<'PY'
import numpy as np
import pandas as pd

# The partition invariant.
scores = pd.DataFrame({"score": [72, 45, np.nan, 91, 50, np.nan, 88, 33]})
high = scores[scores.score > 50]
low = scores[scores.score <= 50]
missing = int(scores.score.isna().sum())
print("partition_naive_sum", len(high) + len(low))
print("partition_total", len(scores))
print("partition_missing", missing)
print("partition_three_way_sum", len(high) + len(low) + missing)

# and/or raise.
try:
    (scores.score > 50) and (scores.score < 90)
    and_raised = False
except ValueError:
    and_raised = True
print("and_raised", and_raised)

# str.contains trap, by dtype.
names_str = pd.Series(["Alice", None, "dave"], dtype="str")
mask_str = names_str.str.contains("a", case=False)
print("str_dtype_mask_has_nan", bool(mask_str.isna().any()))

names_obj = pd.Series(["Alice", None, "dave"], dtype="object")
mask_obj = names_obj.str.contains("a", case=False)
try:
    names_obj[mask_obj]
    obj_filter_raised = False
except ValueError:
    obj_filter_raised = True
print("object_dtype_filter_raises", obj_filter_raised)

mask_obj_fixed = names_obj.str.contains("a", case=False, na=False)
print("na_false_fixed_count", int(mask_obj_fixed.sum()))

# isin([]) versus the whole frame.
staff = pd.DataFrame({"dept": ["eng", "sales", "hr"]})
print("isin_empty_rows", len(staff[staff.dept.isin([])]))
print("isin_empty_negated_rows", len(staff[~staff.dept.isin([])]))

# nlargest keep='all' beyond n.
tied = pd.DataFrame({"score": [80, 80, 80, 60]})
print("nlargest_keep_all_rows", len(tied.nlargest(2, "score", keep="all")))
print("sort_head_rows", len(tied.sort_values("score", ascending=False).head(2)))
PY
)"
echo "${facts}" | sed 's/^/  /'

get_fact() { printf '%s\n' "${facts}" | grep "^$1 " | cut -d' ' -f2-; }

if [ "$(get_fact partition_naive_sum)" != "$(get_fact partition_total)" ]; then
  check "the naive high+low split does NOT equal the total (6 != 8)" "yes"
else
  check "the naive high+low split does NOT equal the total (6 != 8)" "no"
fi
check_eq "the missing count exactly accounts for the shortfall" "2" "$(get_fact partition_missing)"
check_eq "the three-way partition (high+low+missing) equals the total" "$(get_fact partition_total)" "$(get_fact partition_three_way_sum)"
check_eq "\`and\` between two masks raises ValueError" "True" "$(get_fact and_raised)"
check_eq "on pandas 3.0's str dtype, .str.contains() on a missing entry is NOT NaN" "False" "$(get_fact str_dtype_mask_has_nan)"
check_eq "on object dtype, filtering with the unfixed .str.contains() mask raises ValueError" "True" "$(get_fact object_dtype_filter_raises)"
check_eq "na=False on object dtype recovers the correct 2 matches" "2" "$(get_fact na_false_fixed_count)"
check_eq "isin([]) returns zero rows, not the whole frame" "0" "$(get_fact isin_empty_rows)"
check_eq "~isin([]) (negated) returns every row instead" "3" "$(get_fact isin_empty_negated_rows)"
check_eq "nlargest(2, keep='all') returns all 3 rows tied at the cutoff, more than n" "3" "$(get_fact nlargest_keep_all_rows)"
check_eq "sort_values().head(2) is forced to exactly n=2 rows even with the same tie" "2" "$(get_fact sort_head_rows)"

# --------------------------------------------------------------------------
echo
echo "5. Prove the harness can fail, then restore it"
# --------------------------------------------------------------------------

broken_script="${lab_dir}/examples/07_isin_vs_chained.py"
cp "${broken_script}" "${broken_script}.bak"
sed -i.tmp "s/len(via_isin) == 5/len(via_isin) == 999/" "${broken_script}"
rm -f "${broken_script}.tmp"
broken_out="$(cd "${lab_dir}/examples" && "${python_bin}" 07_isin_vs_chained.py 2>&1)"
broken_status=$?
mv "${broken_script}.bak" "${broken_script}"
if [ "${broken_status}" -ne 0 ]; then
  check "a deliberately wrong assertion makes 07_isin_vs_chained.py exit non-zero" "yes"
else
  check "a deliberately wrong assertion makes 07_isin_vs_chained.py exit non-zero" "no"
fi
case "${broken_out}" in
  *FAIL:*) check "the broken run reports a FAIL line" "yes" ;;
  *) check "the broken run reports a FAIL line" "no" ;;
esac
restored_out="$(cd "${lab_dir}/examples" && "${python_bin}" 07_isin_vs_chained.py 2>&1)"
restored_status=$?
check_eq "the script is restored and exits 0 again" "0" "${restored_status}"
case "${restored_out}" in
  *"07_isin_vs_chained.py: every assertion held."*)
    check "the restored script reports every assertion held" "yes" ;;
  *) check "the restored script reports every assertion held" "no" ;;
esac

# --------------------------------------------------------------------------
echo
echo "6. Nothing left behind, and no network dependency baked into the lab"
# --------------------------------------------------------------------------

if grep -rInE 'https?://' "${lab_dir}/examples" "${lab_dir}/starter" >/dev/null 2>&1; then
  check "no URL appears in examples/ or starter/" "no"
else
  check "no URL appears in examples/ or starter/" "yes"
fi

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__' -o -name '.pytest_cache' \) -print 2>/dev/null)"
if [ -z "${stray}" ]; then
  check "no __pycache__ or .pytest_cache directories were left behind" "yes"
else
  check "no __pycache__ or .pytest_cache directories were left behind" "no"
  echo "${stray}" | sed 's/^/      /'
fi

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