Add update-manifest and create-manifest subcommands to bump_dataset_version
update-manifest rewrites the schemaN component in existing manifest files to a specified or auto-detected highest schema, verifying all target files exist before writing. create-manifest builds a new manifest from explicit parquet file paths, supporting --pool/--type (full|holdout|dev) to derive the output path from root, and enforcing holdout isolation by checking for cross-manifest overlap whenever a holdout manifest is involved. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,8 +1,15 @@
|
||||
import os
|
||||
import pytest
|
||||
from scripts import bump_dataset_version
|
||||
|
||||
plan_bump_gen = bump_dataset_version.plan_bump_gen
|
||||
plan_bump_schema = bump_dataset_version.plan_bump_schema
|
||||
apply_bump = bump_dataset_version.apply_bump
|
||||
plan_update_manifest = bump_dataset_version.plan_update_manifest
|
||||
apply_update_manifest = bump_dataset_version.apply_update_manifest
|
||||
plan_create_manifest = bump_dataset_version.plan_create_manifest
|
||||
apply_create_manifest = bump_dataset_version.apply_create_manifest
|
||||
check_holdout_overlap = bump_dataset_version.check_holdout_overlap
|
||||
|
||||
|
||||
def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
|
||||
@@ -84,3 +91,219 @@ def test_apply_bump_appends_without_clobbering_existing_log(tmp_path):
|
||||
text = (tmp_path / "VERSIONS.md").read_text()
|
||||
assert "existing entry" in text
|
||||
assert "reason B" in text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# update-manifest
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_parquet(path):
|
||||
"""Create a zero-byte stand-in for a parquet file."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.touch()
|
||||
|
||||
|
||||
def test_update_manifest_bumps_to_specified_schema(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
old_rel = "../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet"
|
||||
manifest.write_text(old_rel + "\n")
|
||||
|
||||
lines, missing = plan_update_manifest(manifest, "schema2")
|
||||
assert missing == []
|
||||
# exactly one data line was updated
|
||||
changed = [(old, new) for old, new in lines if new is not None]
|
||||
assert len(changed) == 1
|
||||
assert "schema2" in changed[0][1]
|
||||
assert "schema1" not in changed[0][1]
|
||||
|
||||
|
||||
def test_update_manifest_auto_detects_highest_schema(tmp_path):
|
||||
for schema in ("schema1", "schema2", "schema3"):
|
||||
d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
|
||||
d.mkdir(parents=True)
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
parquet.touch()
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
|
||||
|
||||
lines, missing = plan_update_manifest(manifest, None)
|
||||
assert missing == []
|
||||
changed = [(old, new) for old, new in lines if new is not None]
|
||||
assert "schema3" in changed[0][1]
|
||||
|
||||
|
||||
def test_update_manifest_reports_missing_targets(tmp_path):
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
|
||||
|
||||
# schema2 dir exists but the parquet file does not
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
|
||||
|
||||
lines, missing = plan_update_manifest(manifest, "schema2")
|
||||
assert len(missing) == 1
|
||||
assert "schema2" in str(missing[0])
|
||||
|
||||
|
||||
def test_update_manifest_skips_already_at_target(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/steps/gen1/schema2/pbwo4/shard-000.parquet\n")
|
||||
|
||||
lines, missing = plan_update_manifest(manifest, "schema2")
|
||||
assert missing == []
|
||||
# line is unchanged — new is None
|
||||
assert all(new is None for _, new in lines)
|
||||
|
||||
|
||||
def test_update_manifest_preserves_comments_and_blanks(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
content = "# header\n\n../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n"
|
||||
manifest.write_text(content)
|
||||
|
||||
lines, _ = plan_update_manifest(manifest, "schema2")
|
||||
assert lines[0] == ("# header", None)
|
||||
assert lines[1] == ("", None)
|
||||
assert lines[2][1] is not None # the data line was updated
|
||||
|
||||
|
||||
def test_apply_update_manifest_writes_file(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
manifest_dir.mkdir(parents=True)
|
||||
manifest = manifest_dir / "full.manifest"
|
||||
manifest.write_text("../../processed/steps/gen1/schema1/pbwo4/shard-000.parquet\n")
|
||||
|
||||
plan, _ = plan_update_manifest(manifest, "schema2")
|
||||
apply_update_manifest(manifest, plan)
|
||||
|
||||
written = manifest.read_text()
|
||||
assert "schema2" in written
|
||||
assert "schema1" not in written
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# create-manifest
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_create_manifest_writes_relative_paths(tmp_path):
|
||||
pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
|
||||
_make_parquet(pq1)
|
||||
_make_parquet(pq2)
|
||||
|
||||
output = tmp_path / "pools" / "pbwo4" / "train.manifest"
|
||||
lines, missing, resolved = plan_create_manifest(output, [pq1, pq2])
|
||||
|
||||
assert missing == []
|
||||
assert len(lines) == 2
|
||||
assert all("schema2" in l for l in lines)
|
||||
assert all(not l.startswith("/") for l in lines)
|
||||
assert resolved == [pq1.resolve(), pq2.resolve()]
|
||||
|
||||
apply_create_manifest(output, lines)
|
||||
assert output.exists()
|
||||
written = output.read_text().strip().splitlines()
|
||||
assert len(written) == 2
|
||||
|
||||
|
||||
def test_create_manifest_reports_missing_files(tmp_path):
|
||||
ghost = tmp_path / "processed" / "gen1" / "schema2" / "shard-000.parquet"
|
||||
output = tmp_path / "pools" / "full.manifest"
|
||||
lines, missing, _ = plan_create_manifest(output, [ghost])
|
||||
assert len(missing) == 1
|
||||
assert missing[0] == ghost.resolve()
|
||||
|
||||
|
||||
def test_create_manifest_creates_parent_dirs(tmp_path):
|
||||
pq = tmp_path / "a.parquet"
|
||||
pq.touch()
|
||||
output = tmp_path / "deep" / "nested" / "pool.manifest"
|
||||
lines, _, _ = plan_create_manifest(output, [pq])
|
||||
apply_create_manifest(output, lines)
|
||||
assert output.exists()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# check_holdout_overlap
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_no_overlap_check_when_no_holdout_involved(tmp_path):
|
||||
pool_dir = tmp_path / "pools" / "pbwo4"
|
||||
pool_dir.mkdir(parents=True)
|
||||
pq = tmp_path / "a.parquet"
|
||||
pq.touch()
|
||||
# writing full.manifest, no holdout.manifest exists
|
||||
output = pool_dir / "full.manifest"
|
||||
overlaps = check_holdout_overlap(output, [pq.resolve()])
|
||||
assert overlaps == []
|
||||
|
||||
|
||||
def test_overlap_detected_when_creating_holdout(tmp_path):
|
||||
pool_dir = tmp_path / "pools" / "pbwo4"
|
||||
pool_dir.mkdir(parents=True)
|
||||
pq = tmp_path / "a.parquet"
|
||||
pq.touch()
|
||||
|
||||
# full.manifest already lists the same file
|
||||
full = pool_dir / "full.manifest"
|
||||
full.write_text(os.path.relpath(pq.resolve(), start=pool_dir) + "\n")
|
||||
|
||||
output = pool_dir / "holdout.manifest"
|
||||
overlaps = check_holdout_overlap(output, [pq.resolve()])
|
||||
assert len(overlaps) == 1
|
||||
assert overlaps[0][0] == "full.manifest"
|
||||
assert overlaps[0][1] == pq.resolve()
|
||||
|
||||
|
||||
def test_overlap_detected_when_holdout_already_exists(tmp_path):
|
||||
pool_dir = tmp_path / "pools" / "pbwo4"
|
||||
pool_dir.mkdir(parents=True)
|
||||
pq = tmp_path / "a.parquet"
|
||||
pq.touch()
|
||||
|
||||
# holdout.manifest already lists the file
|
||||
holdout = pool_dir / "holdout.manifest"
|
||||
holdout.write_text(os.path.relpath(pq.resolve(), start=pool_dir) + "\n")
|
||||
|
||||
# now creating full.manifest with the same file
|
||||
output = pool_dir / "full.manifest"
|
||||
overlaps = check_holdout_overlap(output, [pq.resolve()])
|
||||
assert len(overlaps) == 1
|
||||
assert overlaps[0][0] == "holdout.manifest"
|
||||
|
||||
|
||||
def test_no_overlap_when_files_are_disjoint(tmp_path):
|
||||
pool_dir = tmp_path / "pools" / "pbwo4"
|
||||
pool_dir.mkdir(parents=True)
|
||||
pq_holdout = tmp_path / "holdout.parquet"
|
||||
pq_full = tmp_path / "full.parquet"
|
||||
pq_holdout.touch()
|
||||
pq_full.touch()
|
||||
|
||||
holdout = pool_dir / "holdout.manifest"
|
||||
holdout.write_text(os.path.relpath(pq_holdout.resolve(), start=pool_dir) + "\n")
|
||||
|
||||
output = pool_dir / "full.manifest"
|
||||
overlaps = check_holdout_overlap(output, [pq_full.resolve()])
|
||||
assert overlaps == []
|
||||
|
||||
Reference in New Issue
Block a user