perf: stream CUR parquet normalization in batches

Author: edmundmiller

bin/benchmark_report_normalize.py

@@ -222,21 +222,27 @@ def _to_tags_dict(value: Any) -> dict[str, Any]:
     return {}
 
 
-def _normalize_cost_rows(costs_parquet: Path) -> list[dict[str, Any]]:
+def _iter_parquet_rows(costs_parquet: Path):
     try:
         import pyarrow.parquet as pq
     except ImportError as exc:
         raise RuntimeError("pyarrow is required to normalize CUR parquet") from exc
 
-    table = pq.read_table(costs_parquet)
-    rows = table.to_pylist()
-    cols = set(table.column_names)
+    parquet_file = pq.ParquetFile(costs_parquet)
+    cols = set(parquet_file.schema_arrow.names)
+
+    for batch in parquet_file.iter_batches():
+        for row in batch.to_pylist():
+            yield cols, row
 
-    is_map = "resource_tags" in cols and "resource_tags_user_unique_run_id" not in cols
 
+def _normalize_cost_rows(costs_parquet: Path) -> list[dict[str, Any]]:
     grouped: dict[tuple[str, str, str], dict[str, float | str]] = {}
+    is_map: bool | None = None
 
-    for row in rows:
+    for cols, row in _iter_parquet_rows(costs_parquet):
+        if is_map is None:
+            is_map = "resource_tags" in cols and "resource_tags_user_unique_run_id" not in cols
         if is_map:
             tags = _to_tags_dict(row.get("resource_tags"))
             run_id = tags.get("user_unique_run_id") or tags.get("user_nf_unique_run_id")

modules/local/normalize_benchmark_jsonl/tests/test_normalize.py

@@ -1,6 +1,9 @@
 import json
 
-from benchmark_report_normalize import extract_runs, extract_tasks, load_run_data, normalize_jsonl
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from benchmark_report_normalize import _normalize_cost_rows, extract_runs, extract_tasks, load_run_data, normalize_jsonl
 
 
 def test_cached_count_extracted(make_run, flat_task):
@@ -38,6 +41,33 @@ def test_normalize_writes_jsonl_bundle(tmp_path, make_run, flat_task, write_run_
     assert task["process_short"] == "PROCESS_A"
 
 
+def test_normalize_cost_rows_reads_parquet_in_batches(tmp_path):
+    parquet_path = tmp_path / "costs.parquet"
+    table = pa.table(
+        {
+            "resource_tags_user_unique_run_id": ["run1", "run1"],
+            "resource_tags_user_pipeline_process": ["PROC_A", "PROC_A"],
+            "resource_tags_user_task_hash": ["abcdef12", "abcdef12"],
+            "split_line_item_split_cost": [1.25, 2.75],
+            "split_line_item_unused_cost": [0.25, 0.75],
+        }
+    )
+    pq.write_table(table, parquet_path, row_group_size=1)
+
+    rows = _normalize_cost_rows(parquet_path)
+
+    assert rows == [
+        {
+            "run_id": "run1",
+            "process": "PROC_A",
+            "hash": "abcdef12",
+            "cost": 5.0,
+            "used_cost": 4.0,
+            "unused_cost": 1.0,
+        }
+    ]
+
+
 def test_load_run_data(tmp_path, make_run, write_run_json):
     data_dir = tmp_path / "data"
     write_run_json(data_dir, [make_run(), make_run(run_id="run2")])