diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..aca63a8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,33 @@ +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST +.pytest_cache/ +.coverage +.tox/ +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ diff --git a/tests/test_shard_from_audio.py b/tests/test_shard_from_audio.py index 9a9c31c..7505398 100644 --- a/tests/test_shard_from_audio.py +++ b/tests/test_shard_from_audio.py @@ -1,5 +1,4 @@ import struct -from pathlib import Path import pyarrow as pa import pyarrow.ipc @@ -202,3 +201,31 @@ def test_shard_naming(self, tmp_path): shard_names = sorted(p.name for p in output_dir.glob("*.wsds")) assert shard_names == ["audio-00000.wsds", "audio-00001.wsds"] + + def test_sort_files_parameter(self, tmp_path): + """sort_files=True produces deterministic ordering.""" + input_dir = tmp_path / "in" + output_dir_sorted = tmp_path / "out_sorted" + output_dir_unsorted = tmp_path / "out_unsorted" + input_dir.mkdir() + + # Create files with names that would sort differently than filesystem order + stems = ["zebra", "alpha", "beta", "gamma"] + for stem in stems: + make_wav(input_dir / f"{stem}.wav") + + # Run with sort_files=True + shard_from_audio_dir(str(input_dir), str(output_dir_sorted), sort_files=True) + + # Verify sorted output has keys in alphabetical order + shards_sorted = _collect_shards(output_dir_sorted) + keys_sorted = [k for keys, _, _ in shards_sorted for k in keys] + assert keys_sorted == sorted(stems) + + # Run without sorting (default) + shard_from_audio_dir(str(input_dir), str(output_dir_unsorted), sort_files=False) + + # Verify unsorted output has same keys (but possibly different order) + shards_unsorted = _collect_shards(output_dir_unsorted) + keys_unsorted = [k for keys, _, _ in shards_unsorted for k in keys] + assert sorted(keys_unsorted) == sorted(stems) diff --git a/wsds/ws_tools.py b/wsds/ws_tools.py index b0b98da..4b1dd30 100644 --- a/wsds/ws_tools.py +++ b/wsds/ws_tools.py @@ -6,7 +6,7 @@ import sys from collections.abc import Callable from pathlib import Path -import numpy as np + import polars as pl import pyarrow as pa @@ -327,7 +327,7 @@ def keys(dataset: Path, verbose=False, skip_audio=True): try: if not pl.scan_ipc(shard_fname).select((pl.col("__key__") == expected_keys).all()).collect().item(): tqdm.write(f"Shard {shard} in {subdir} has keys that don't match the index.") - except pl.exceptions.ShapeError as err: + except pl.exceptions.ShapeError: tqdm.write(f"Shard {shard} in {subdir} has {pl.scan_ipc(shard_fname).select(pl.len()).collect().item()} keys while we expect {len(expected_keys)}.") reader = pa.RecordBatchFileReader(pa.memory_map(str(shard_fname))) batch_size = int(reader.schema.metadata[b'batch_size']) @@ -611,6 +611,7 @@ def shard_from_audio_dir( key_fn: Callable[[str], str] | None = None, write_key_mapping: bool = False, key_prefix: str = "", + sort_files: bool = False, ): """Write batched Feather (.wsds) shards with up to N audio files each. @@ -622,6 +623,9 @@ def shard_from_audio_dir( key_prefix: Optional prefix to prepend to keys before hashing. Useful to avoid collisions when processing multiple directories with files that share the same names (e.g., "egyptian", "saudi"). + sort_files: If True, sorts all files before processing (loads full list + into memory). If False (default), iterates without sorting for + better performance with large corpora. """ from tqdm import tqdm @@ -629,8 +633,13 @@ def shard_from_audio_dir( output_dir.mkdir(parents=True, exist_ok=True) exts = (".wav", ".flac", ".mp3", ".m4a", ".ogg", ".opus") - all_files = sorted(p for p in input_dir.rglob("*") if p.suffix.lower() in exts) - print(f"[INFO] Found {len(all_files):,} audio files under {input_dir}") + file_iter = (p for p in input_dir.rglob("*") if p.suffix.lower() in exts) + if sort_files: + files = sorted(file_iter) + print(f"[INFO] Found {len(files):,} audio files under {input_dir}") + else: + files = file_iter + print(f"[INFO] Processing audio files under {input_dir}") MAX_ARROW_BYTES = 2_140_000_000 # ~2.1 GB Arrow cell limit @@ -647,7 +656,7 @@ def flush_batch(): shard_idx += 1 batch = [] - for path in tqdm(all_files, ncols=90, desc="Writing WSDS shards"): + for path in tqdm(files, ncols=90, desc="Writing WSDS shards"): rel_path = path.relative_to(input_dir).with_suffix('') stem = str(rel_path) if key_prefix: