diff --git a/support_scripts/bench_sql_select_fastpath.py b/support_scripts/bench_sql_select_fastpath.py new file mode 100644 index 0000000..355e247 --- /dev/null +++ b/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") + +# %% diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py index f6ccce1..6dfd4a7 100644 --- a/wsds/ws_dataset.py +++ b/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