ws_dataset: add multi-file scan_ipc fast path to sql_select

support_scripts/bench_sql_select_fastpath.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python
+#%%
+"""bench_sql_select_fastpath.py — fast vs slow path timing for sql_select().
+
+`_parse_sql_queries_polars` has two code paths:
+
+- fast path: one `pl.scan_ipc(paths_list)` per col_dir, plan = ~N_col_dirs nodes.
+- slow path: one `scan_ipc` per (shard, col_dir), plan = N_shards × N_col_dirs.
+
+The slow path is selected when the caller asks for `__shard_offset__`, passes
+a `shard_pipe`, or any shard is invalid in some col_dir.
+
+This script just times both, no profiler. N_RUNS each, best-of reported.
+"""
+
+import gc
+import time
+
+import wsds
+
+# ---- config ----
+DATASET_PATH = "/mnt/weka/data-wsds/data-ar/indices/source"   # <-- change me
+N_RUNS = 3
+
+QUERIES_FAST = (
+    "__key__",
+    "__shard_path__",
+    "load_duration AS audio_duration",
+    "duration",
+    "duration_seconds",
+    "est_duration",
+    "inspected_duration",
+    "speech_duration",
+)
+QUERIES_SLOW = QUERIES_FAST + ("__shard_offset__ AS offset",)
+
+# %%
+ds = wsds.WSDataset(str(DATASET_PATH))
+print(f"dataset: {DATASET_PATH}")
+print(f"shards:  {len(ds.get_shard_list()):,}\n")
+
+
+def reset():
+    if hasattr(ds, "_validated_shards"):
+        ds._validated_shards.clear()
+    gc.collect()
+
+
+def time_runs(label, queries):
+    runs = []
+    rows = None
+    for _ in range(N_RUNS):
+        reset()
+        t0 = time.perf_counter()
+        df = ds.sql_select(*queries, shard_subsample=1)
+        runs.append((time.perf_counter() - t0) * 1000)
+        rows = len(df)
+    best = min(runs)
+    print(f"  {label:30s} best={best:>7,.0f} ms   rows={rows:,}   runs={[f'{r:.0f}ms' for r in runs]}")
+    return best
+
+
+# %%
+fast_ms = time_runs("FAST (multi-file scan)", QUERIES_FAST)
+slow_ms = time_runs("SLOW (per-shard scan)",  QUERIES_SLOW)
+
+print()
+print(f"  speedup: {slow_ms / fast_ms:.2f}x")
+
+# %%

wsds/ws_dataset.py

@@ -298,6 +298,27 @@ def _parse_sql_queries_polars(self, *queries, shard_subsample=1, rng=None, shard
         # Prefetch shard tails concurrently to warm up the filesystem cache
         verified_shard_list = validate_shards(self, shard_list, list(column_dirs.keys()))
 
+        # Fast path: when no shard_pipe, no __shard_offset__, and every shard is
+        # valid in every col_dir, we can collapse the per-shard scan loop into
+        # one multi-file `pl.scan_ipc(paths_list)` per col_dir. The plan goes
+        # from ~N_shards × N_col_dirs scan nodes to N_col_dirs — the polars
+        # optimizer scales super-linearly with plan size, so this is ~2× faster
+        # on wide workloads. Falls back to the per-shard path below when any
+        # condition fails.
+        if (
+            shard_pipe is None
+            and "__shard_offset__" not in needed_special_columns
+            and verified_shard_list
+            and all(ok for _, ok in verified_shard_list)
+        ):
+            return exprs, self._build_multifile_plan(
+                [s for s, _ in verified_shard_list],
+                column_dirs,
+                exprs,
+                key_column_dir,
+                "__shard_path__" in needed_special_columns,
+            )
+
         row_merge = []
         column_dir_samples = {}
         missing = defaultdict(list)
@@ -350,6 +371,34 @@ def _parse_sql_queries_polars(self, *queries, shard_subsample=1, rng=None, shard
 
         return exprs, pl.concat(row_merge)
 
+    def _build_multifile_plan(self, shards, column_dirs, exprs, key_column_dir, needs_shard_path):
+        """Fast path for `_parse_sql_queries_polars`: one `pl.scan_ipc(paths_list)`
+        per col_dir, horizontally concatenated. See the comment in the caller
+        for why this is faster than the per-shard scan loop.
+
+        Pre-conditions enforced by the caller:
+        - all shards are valid in every col_dir (no `pl.defer` NULL-fill path)
+        - no per-shard `shard_pipe` (semantics would be lost across the union)
+        - `__shard_offset__` is not requested (multi-file scan's row index is
+          global rather than per-file)
+        """
+        per_dir_frames = []
+        for column_dir, fields in column_dirs.items():
+            paths = [self.get_shard_path(column_dir, s) for s in shards]
+            is_key_dir = column_dir == key_column_dir
+            # `__shard_path__` is synthesized by polars via `include_file_paths`;
+            # strip it (and any stray `__shard_offset__`) from the file-column
+            # select list and re-add `__shard_path__` after the scan if needed.
+            select_fields = [f for f in fields if f not in ("__shard_path__", "__shard_offset__")]
+            df = pl.scan_ipc(
+                paths,
+                include_file_paths="__shard_path__" if (is_key_dir and needs_shard_path) else None,
+            )
+            if is_key_dir and needs_shard_path:
+                select_fields = select_fields + ["__shard_path__"]
+            per_dir_frames.append(df.select(select_fields))
+        return pl.concat(per_dir_frames, how="horizontal").select(exprs)
+
     def _check_for_subsampling(self, shard_subsample):
         if shard_subsample is None:
             # Check if we're running inside a PyTorch DataLoader worker