lars
05d5dee606
The merged proc_idx/proc_map plumbing wasn't run through ruff format before merging; reflow only, no logic changes.
93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
import pandas as pd
|
|
import pytest
|
|
|
|
from giant.data.loader import build_process_map_from_files, find_parquet_files
|
|
|
|
|
|
def _touch(path):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.touch()
|
|
return path
|
|
|
|
|
|
def test_find_parquet_files_single_file(tmp_path):
|
|
f = _touch(tmp_path / "shard-000.parquet")
|
|
assert find_parquet_files(f) == [f]
|
|
|
|
|
|
def test_find_parquet_files_directory_glob(tmp_path):
|
|
a = _touch(tmp_path / "shard-000.parquet")
|
|
b = _touch(tmp_path / "shard-001.parquet")
|
|
_touch(tmp_path / "not_a_parquet.root")
|
|
assert find_parquet_files(tmp_path) == sorted([a, b])
|
|
|
|
|
|
def test_find_parquet_files_empty_directory_raises(tmp_path):
|
|
with pytest.raises(FileNotFoundError):
|
|
find_parquet_files(tmp_path)
|
|
|
|
|
|
def test_manifest_resolves_relative_to_its_own_directory(tmp_path):
|
|
target = _touch(tmp_path / "processed" / "pbwo4" / "shard-000.parquet")
|
|
manifest_dir = tmp_path / "pools" / "pbwo4"
|
|
manifest_dir.mkdir(parents=True)
|
|
manifest = manifest_dir / "full.manifest"
|
|
manifest.write_text("../../processed/pbwo4/shard-000.parquet\n")
|
|
|
|
assert find_parquet_files(manifest) == [target.resolve()]
|
|
|
|
|
|
def test_manifest_skips_blank_lines_and_comments(tmp_path):
|
|
target = _touch(tmp_path / "shard-000.parquet")
|
|
manifest = tmp_path / "full.manifest"
|
|
manifest.write_text("\n# a comment\nshard-000.parquet\n\n")
|
|
|
|
assert find_parquet_files(manifest) == [target.resolve()]
|
|
|
|
|
|
def test_manifest_missing_file_raises(tmp_path):
|
|
manifest = tmp_path / "full.manifest"
|
|
manifest.write_text("does_not_exist.parquet\n")
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
find_parquet_files(manifest)
|
|
|
|
|
|
def test_manifest_with_no_entries_raises(tmp_path):
|
|
manifest = tmp_path / "full.manifest"
|
|
manifest.write_text("# only comments\n")
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
find_parquet_files(manifest)
|
|
|
|
|
|
def test_build_process_map_from_files_keeps_most_frequent(tmp_path):
|
|
"""process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only
|
|
the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other"
|
|
(last) index."""
|
|
process = ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1
|
|
path = tmp_path / "shard-000.parquet"
|
|
pd.DataFrame({"process": process}).to_parquet(path)
|
|
|
|
proc_map = build_process_map_from_files([path], n_experts=3)
|
|
|
|
assert proc_map["eIoni"] == 0
|
|
assert proc_map["phot"] == 1
|
|
assert proc_map["compt"] == 2
|
|
assert proc_map["Rayl"] == 2
|
|
assert set(proc_map.values()) <= {0, 1, 2}
|
|
|
|
|
|
def test_build_process_map_from_files_spans_multiple_files(tmp_path):
|
|
path_a = tmp_path / "a.parquet"
|
|
path_b = tmp_path / "b.parquet"
|
|
pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a)
|
|
pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b)
|
|
|
|
# phot: 1+4=5 total > eIoni: 3 > compt: 1
|
|
proc_map = build_process_map_from_files([path_a, path_b], n_experts=3)
|
|
|
|
assert proc_map["phot"] == 0
|
|
assert proc_map["eIoni"] == 1
|
|
assert proc_map["compt"] == 2
|