Add tooling for a versioned geant_steps dataset layout
Introduces raw/<kind>/<gen>/<detector>/shard-NNN.root and processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet as the dataset convention, plus scripts to operate on it: migrate_geant_steps.py for the one-time move into this layout, bump_dataset_version.py to cut new gen/schema versions with a logged reason, steps_to_parquet_parallel.py to convert ROOT shards to parquet in parallel and place them correctly, and create_root_files.py to generate new ROOT shards via a minicalosim executable. The loader gains .manifest file support so pools/ (train/dev/ holdout shard lists) can be passed straight to `giant train`. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,60 @@
|
||||
import pytest
|
||||
|
||||
from giant.data.loader import 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)
|
||||
Reference in New Issue
Block a user