lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Add file counts and reference tracking to dwarf status

Browse Source 
Shows per-directory file counts throughout the tree, plus a referenced
count for raw/ (matched against any same-named parquet under processed/)
and each schemaN dir (matched against pools/*.manifest entries).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-02 11:09:47 +02:00
parent 7761cd21b6
commit 57e044e347
1 changed files with 143 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/bump_dataset_version.py
+143 -16
View File
@@ -157,6 +157,81 @@ def _du(path: Path) -> int:
return total
def _count_files(path: Path) -> int:
"""Total number of regular files under *path*, recursively (0 if missing)."""
if not path.is_dir():
return 0
total = 0
for _dirpath, _dirnames, filenames in os.walk(path):
total += len(filenames)
return total
# Must match giant.data.loader.MANIFEST_SUFFIX.
MANIFEST_SUFFIX = ".manifest"
def _manifest_referenced_files(pools_root: Path) -> set[Path]:
"""Resolved absolute paths of every file listed in any *.manifest under *pools_root*."""
referenced: set[Path] = set()
if not pools_root.is_dir():
return referenced
for manifest_path in pools_root.rglob(f"*{MANIFEST_SUFFIX}"):
try:
referenced.update(_resolve_manifest_files(manifest_path))
except OSError:
continue
return referenced
def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
"""(total .root files, count with a same-named .parquet under any schema) for one gen."""
if not raw_gen_dir.is_dir():
return 0, 0
parquet_stems: set[tuple[str, str]] = set()
if processed_gen_dir.is_dir():
for schema_dir in processed_gen_dir.iterdir():
if not schema_dir.is_dir():
continue
for detector_dir in schema_dir.iterdir():
if not detector_dir.is_dir():
continue
for f in detector_dir.iterdir():
if f.is_file() and f.suffix == ".parquet":
parquet_stems.add((detector_dir.name, f.stem))
total = 0
referenced = 0
for detector_dir in raw_gen_dir.iterdir():
if not detector_dir.is_dir():
continue
for f in detector_dir.iterdir():
if f.is_file() and f.suffix == ".root":
total += 1
if (detector_dir.name, f.stem) in parquet_stems:
referenced += 1
return total, referenced
def _referenced_parquet_count(
schema_dir: Path, manifest_referenced: set[Path]
) -> tuple[int, int]:
"""(total .parquet files, count listed in at least one manifest) for one schema dir."""
if not schema_dir.is_dir():
return 0, 0
total = 0
referenced = 0
for detector_dir in schema_dir.iterdir():
if not detector_dir.is_dir():
continue
for f in detector_dir.iterdir():
if f.is_file() and f.suffix == ".parquet":
total += 1
if f.resolve() in manifest_referenced:
referenced += 1
return total, referenced
def _human_size(n: int) -> str:
size = float(n)
for unit in ("B", "KB", "MB", "GB", "TB"):
@@ -166,14 +241,30 @@ def _human_size(n: int) -> str:
return f"{size:.1f} TB"
_ROW_WIDTH = 48
_LABEL_WIDTH = 26
_COUNT_WIDTH = 20
_ROW_WIDTH = _LABEL_WIDTH + _COUNT_WIDTH + 10
def _row(label: str, size_bytes: int, indent: int = 0, level: str | None = None) -> str:
def _count_str(count: int, referenced: int | None = None) -> str:
files = "file" if count == 1 else "files"
if referenced is not None:
return f"{count} {files} ({referenced} ref)"
return f"{count} {files}"
def _row(
label: str,
size_bytes: int,
indent: int = 0,
level: str | None = None,
count: int | None = None,
referenced: int | None = None,
) -> str:
text = " " * indent + label
count_str = _count_str(count, referenced) if count is not None else ""
size_str = _human_size(size_bytes)
pad = max(1, _ROW_WIDTH - len(text) - len(size_str))
row = text + " " * pad + size_str
row = f"{text:<{_LABEL_WIDTH}}{count_str:<{_COUNT_WIDTH}}{size_str:>10}"
return _colorize(row, level) if level else row
@@ -182,7 +273,9 @@ def print_status(root: Path) -> None:
if not raw_root.is_dir():
print(f"no raw/ tree found under {root}")
return
manifest_referenced = _manifest_referenced_files(root / "pools")
grand_total = 0
grand_files = 0
for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
kind = kind_dir.name
gens = sorted(
@@ -192,10 +285,13 @@ def print_status(root: Path) -> None:
)
print(_colorize(f"{kind}/", "kind"))
kind_total = 0
kind_files = 0
for gen in gens:
gen_tag = f"gen{gen}"
raw_size = _du(raw_root / kind / gen_tag)
schema_dir = root / "processed" / kind / gen_tag
raw_gen_dir = raw_root / kind / gen_tag
raw_size = _du(raw_gen_dir)
processed_gen_dir = root / "processed" / kind / gen_tag
schema_dir = processed_gen_dir
schemas = sorted(
int(m.group(1))
for m in (
@@ -206,29 +302,60 @@ def print_status(root: Path) -> None:
if m
)
schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
schema_counts = {
s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced)
for s in schemas
}
processed_size = sum(schema_sizes.values())
processed_files = sum(c[0] for c in schema_counts.values())
processed_referenced = sum(c[1] for c in schema_counts.values())
raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
gen_total = raw_size + processed_size
gen_files = raw_files + processed_files
kind_total += gen_total
kind_files += gen_files
print(_row(gen_tag, gen_total, indent=1, level="gen"))
print(_row("raw", raw_size, indent=2, level="bucket"))
print(_row("processed", processed_size, indent=2, level="bucket"))
print(_row(gen_tag, gen_total, indent=1, level="gen", count=gen_files))
print(
_row(
"raw", raw_size, indent=2, level="bucket",
count=raw_files, referenced=raw_referenced,
)
)
print(
_row(
"processed", processed_size, indent=2, level="bucket",
count=processed_files, referenced=processed_referenced,
)
)
if schemas:
for s in schemas:
print(_row(f"schema{s}", schema_sizes[s], indent=3, level="schema"))
s_total, s_referenced = schema_counts[s]
print(
_row(
f"schema{s}", schema_sizes[s], indent=3, level="schema",
count=s_total, referenced=s_referenced,
)
)
else:
print(_colorize(" (none)", "schema"))
print(_row(f"{kind} total", kind_total, indent=1, level="gen"))
print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
print()
grand_total += kind_total
grand_files += kind_files
derived_size = _du(root / "derived")
pools_size = _du(root / "pools")
derived_dir = root / "derived"
pools_dir = root / "pools"
derived_size = _du(derived_dir)
pools_size = _du(pools_dir)
derived_files = _count_files(derived_dir)
pools_files = _count_files(pools_dir)
grand_total += derived_size + pools_size
print(_row("derived/", derived_size, level="root"))
print(_row("pools/", pools_size, level="root"))
grand_files += derived_files + pools_files
print(_row("derived/", derived_size, level="root", count=derived_files))
print(_row("pools/", pools_size, level="root", count=pools_files))
print("-" * _ROW_WIDTH)
print(_row("grand total", grand_total, level="root"))
print(_row("grand total", grand_total, level="root", count=grand_files))
# ---------------------------------------------------------------------------