Make file sorting optional in shard_from_audio_dir for large corpora

.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/

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)

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,15 +623,23 @@ 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
 
     input_dir, output_dir = Path(input_dir), Path(output_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: