Parallelize profiling by Parquet shard

tests/test_cli.py

@@ -3,9 +3,17 @@
 from __future__ import annotations
 
 import json
+from argparse import Namespace
 from pathlib import Path
 
-from swe_data_processing.cli import _run_streaming_stage, build_parser
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from swe_data_processing.cli import (
+    _run_parallel_parquet_profile,
+    _run_streaming_stage,
+    build_parser,
+)
 
 
 def test_streaming_stage_writes_each_record(tmp_path: Path) -> None:
@@ -36,3 +44,34 @@ def test_cli_has_only_deterministic_pipeline_commands() -> None:
     help_text = parser.format_help()
     assert "classify" not in help_text
     assert "repair" not in help_text
+
+
+def test_parallel_parquet_profile_uses_atomic_shard_parts(tmp_path: Path) -> None:
+    dataset = tmp_path / "dataset"
+    dataset.mkdir()
+    for index in range(2):
+        pq.write_table(
+            pa.Table.from_pylist(
+                [{"trajectory_id": f"sample-{index}", "resolved": 1}]
+            ),
+            dataset / f"part-{index}.parquet",
+        )
+    output = tmp_path / "metrics.jsonl"
+    errors = tmp_path / "errors.jsonl"
+    status = _run_parallel_parquet_profile(
+        Namespace(
+            input=dataset,
+            output=output,
+            errors=errors,
+            tokenizer_json=None,
+            workers=2,
+            resume=False,
+        )
+    )
+
+    assert status == 0
+    assert [json.loads(line)["sample_id"] for line in output.read_text().splitlines()] == [
+        "sample-0",
+        "sample-1",
+    ]
+    assert errors.read_text() == ""