diff --git a/uc_salary_pipeline.py b/uc_salary_pipeline.py index ef2fa2f..6863450 100644 --- a/uc_salary_pipeline.py +++ b/uc_salary_pipeline.py @@ -40,10 +40,18 @@ Optionally, an aggregate CSV (--aggregate-output) with one row per category (broken down further by --group-by year/campus): [year] [campus] category, category_label, employee_count, total_gross_pay, mean_gross_pay, median_gross_pay + +Optionally, a top-earners CSV (--top-earners-output) with the --top-n +(default 10) highest gross_pay employees per category (also broken down +further by --group-by, e.g. per year): + category, category_label, rank, year, campus, title, + gross_pay, regular_pay, overtime_pay, other_pay, first_name, last_name """ import argparse import csv +import heapq +import itertools import json import logging import re @@ -103,14 +111,26 @@ def parse_uc_wages_record(record: dict) -> dict | None: } +def _group_sort_key(key: tuple) -> tuple: + """Sort a (group_dims..., category) key by its group values, then by + category in classifier.CATEGORIES order (not alphabetically).""" + *prefix, category = key + return (*prefix, classifier.CATEGORIES.index(category)) + + def process_file(path: Path, campus: str, year: int, writer: csv.writer, senior_threshold: float, support_threshold: float, - agg: dict | None = None, group_dims: tuple = ()) -> int: + agg: dict | None = None, group_dims: tuple = (), + top_earners: dict | None = None, top_n: int = 10, + tie_breaker: itertools.count | None = None) -> int: """Classify every record in one UC-wages JSON file, writing a CSV row per employee. Returns the number of rows written. - If `agg` is given, also accumulates each record's gross_pay under an + If `agg` is given, also accumulates each record's gross_pay under a (group_dims..., category) key for later aggregate reporting. + + If `top_earners` is given, also maintains a size-`top_n` min-heap of + the highest gross_pay employees per (group_dims..., category) key. """ with open(path, encoding="utf-8") as fh: data = json.load(fh) @@ -140,6 +160,16 @@ def process_file(path: Path, campus: str, year: int, writer: csv.writer, count += 1 if agg is not None: agg.setdefault(group_prefix + (category,), []).append(parsed["gross_pay"]) + if top_earners is not None: + entry = (parsed["gross_pay"], next(tie_breaker), + (year, campus, parsed["title"], parsed["gross_pay"], + parsed["regular_pay"], parsed["overtime_pay"], + parsed["other_pay"], parsed["first_name"], parsed["last_name"])) + heap = top_earners.setdefault(group_prefix + (category,), []) + if len(heap) < top_n: + heapq.heappush(heap, entry) + else: + heapq.heappushpop(heap, entry) return count @@ -151,14 +181,10 @@ def write_aggregate_csv(agg: dict, group_dims: tuple, outpath: Path) -> None: "employee_count", "total_gross_pay", "mean_gross_pay", "median_gross_pay", ] - def sort_key(key): - *prefix, category = key - return (*prefix, classifier.CATEGORIES.index(category)) - with open(outpath, "w", newline="", encoding="utf-8") as fh: writer = csv.writer(fh) writer.writerow(header) - for key in sorted(agg, key=sort_key): + for key in sorted(agg, key=_group_sort_key): *prefix, category = key values = agg[key] total = sum(values) @@ -170,6 +196,29 @@ def write_aggregate_csv(agg: dict, group_dims: tuple, outpath: Path) -> None: log.info("Aggregate CSV -> %s", outpath) +def write_top_earners_csv(top_earners: dict, outpath: Path) -> None: + """Write the top-N earners per (group_dims..., category) key, ranked + highest gross_pay first within each group.""" + header = ["category", "category_label", "rank", "year", "campus", "title", + "gross_pay", "regular_pay", "overtime_pay", "other_pay", + "first_name", "last_name"] + + with open(outpath, "w", newline="", encoding="utf-8") as fh: + writer = csv.writer(fh) + writer.writerow(header) + for key in sorted(top_earners, key=_group_sort_key): + category = key[-1] + ranked = sorted(top_earners[key], key=lambda e: e[0], reverse=True) + for rank, (gross_pay, _, row) in enumerate(ranked, start=1): + year, campus, title, gross, regular, overtime, other, first, last = row + writer.writerow([ + category, classifier.CATEGORY_LABELS[category], rank, + year, campus, title, round(gross, 2), round(regular, 2), + round(overtime, 2), round(other, 2), first, last, + ]) + log.info("Top earners CSV -> %s", outpath) + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @@ -202,10 +251,18 @@ def main(): parser.add_argument("--group-by", nargs="*", choices=["year", "campus"], default=["year"], help="Additional breakdown dimensions for " - "--aggregate-output, alongside category " - "(default: year). Pass both for a year x campus " - "x category breakdown, or nothing to collapse " - "everything to one row per category.") + "--aggregate-output and --top-earners-output, " + "alongside category (default: year). Pass both " + "for a year x campus x category breakdown, or " + "nothing to collapse everything to one group " + "per category.") + parser.add_argument("--top-earners-output", + help="If given, also write a CSV of the --top-n " + "highest gross_pay employees per category (see " + "--group-by).") + parser.add_argument("--top-n", type=int, default=10, + help="Number of top earners per group for " + "--top-earners-output (default: 10).") parser.add_argument("--skip-download", action="store_true", help="Classify whatever is already cached in " "--data-dir without hitting the network.") @@ -228,6 +285,8 @@ def main(): group_dims = tuple(dim for dim in ("year", "campus") if dim in args.group_by) agg = {} if args.aggregate_output else None + top_earners = {} if args.top_earners_output else None + tie_breaker = itertools.count() if top_earners is not None else None output_path = Path(args.output) total_rows = 0 @@ -243,7 +302,9 @@ def main(): continue count = process_file(path, campus, year, writer, args.senior_threshold, args.support_threshold, - agg=agg, group_dims=group_dims) + agg=agg, group_dims=group_dims, + top_earners=top_earners, top_n=args.top_n, + tie_breaker=tie_breaker) total_rows += count log.info("%s %s: classified %d records", campus, year, count) @@ -256,6 +317,9 @@ def main(): if agg is not None: write_aggregate_csv(agg, group_dims, Path(args.aggregate_output)) + if top_earners is not None: + write_top_earners_csv(top_earners, Path(args.top_earners_output)) + if __name__ == "__main__": main()