File size: 2,642 Bytes
201dc77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """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 # noqa: F401
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*")) # no partial output
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 # stale ingest never served
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 # idempotent
|