| """ingest_parquet / local row readers. The pyarrow-dependent cases skip when |
| pyarrow is absent (it is not an align dependency); the error path always runs.""" |
| from __future__ import annotations |
|
|
| import json |
|
|
| import pytest |
|
|
| from align.data import ingest_parquet as ip |
| from align.data import sources |
|
|
|
|
| def _has_pyarrow() -> bool: |
| try: |
| import pyarrow |
| return True |
| except ImportError: |
| return False |
|
|
|
|
| def test_jsonl_rows_need_no_pyarrow(tmp_path): |
| p = tmp_path / "x.jsonl" |
| p.write_text('{"a":1}\n\n{"a":2}\n') |
| assert list(sources.iter_local_rows(p)) == [{"a": 1}, {"a": 2}] |
|
|
|
|
| @pytest.mark.skipif(_has_pyarrow(), reason="tests the pyarrow-absent error path") |
| def test_missing_pyarrow_points_at_hub(tmp_path): |
| p = tmp_path / "x.parquet" |
| p.write_bytes(b"PAR1") |
| with pytest.raises(ip.IngestUnavailable, match="align.data.hub.iter_rows"): |
| list(sources.iter_local_rows(p)) |
| with pytest.raises(ImportError): |
| ip.ingest_parquet(p, out=tmp_path / "o.jsonl") |
| assert not list(tmp_path.glob("*.tmp*")) |
|
|
|
|
| def test_prior_ingest_is_used_without_pyarrow(tmp_path, monkeypatch): |
| """An already-ingested JSONL serves a parquet read: the training venv never |
| needs pyarrow once a one-time ingest was done elsewhere.""" |
| p = tmp_path / "x.parquet" |
| p.write_bytes(b"PAR1") |
| monkeypatch.setattr(ip, "INGEST_DIR", str(tmp_path / "ing")) |
| out = ip.ingested_path(p, out_dir=tmp_path / "ing") |
| out.parent.mkdir() |
| out.write_text(json.dumps({"q": 1}) + "\n") |
| assert list(ip.iter_parquet_rows(p)) == [{"q": 1}] |
|
|
|
|
| def test_ingested_path_changes_with_file(tmp_path): |
| p = tmp_path / "x.parquet" |
| p.write_bytes(b"a") |
| a = ip.ingested_path(p, out_dir=tmp_path) |
| p.write_bytes(b"bb") |
| assert ip.ingested_path(p, out_dir=tmp_path) != a |
|
|
|
|
| def test_json_default_bytes_marked_not_dropped(): |
| assert ip._json_default(b"\x00\x01") == {"__bytes_b64__": "AAE="} |
| with pytest.raises(TypeError): |
| ip._json_default(object()) |
|
|
|
|
| @pytest.mark.skipif(not _has_pyarrow(), reason="pyarrow not installed") |
| def test_roundtrip_with_pyarrow(tmp_path): |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| rows = [{"q": f"p{i}", "v": float(i)} for i in range(10)] |
| p = tmp_path / "d.parquet" |
| pq.write_table(pa.Table.from_pylist(rows), str(p)) |
| assert list(ip.iter_parquet_rows(p, prefer_ingested=False)) == rows |
| out = ip.ingest_parquet(p, out=tmp_path / "d.jsonl") |
| assert list(sources.iter_local_rows(out)) == rows |
| assert ip.ingest_parquet(p, out=out) == out |
|
|
