add top-earners-output: top-N gross pay employees per category

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 "
+                             "--aggregate-output and --top-earners-output, "
-                             "(default: year). Pass both for a year x campus "
+                             "alongside category (default: year). Pass both "
-                             "x category breakdown, or nothing to collapse "
+                             "for a year x campus x category breakdown, or "
-                             "everything to one row per category.")
+                             "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()