diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 1a2732b..5b7afd1 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -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)) # ---------------------------------------------------------------------------