From 76dc72e22aa1d2492e52f2da4fa7930df2d60d6f Mon Sep 17 00:00:00 2001 From: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:12:43 -0400 Subject: [PATCH] feat: add opt-in multithreaded fill to synthetic data generators The synthetic generators are dominated by single-threaded numpy standard_normal draws (~0.9 GB/s), leaving cores idle. Add an opt-in n_workers parameter to scan_synthetic_regression and scan_synthetic_panel that fills the x/eps Gaussian draws in parallel. Draws are generated in fixed-size blocks (n_workers * fetch_size rows), each split across worker threads that fill disjoint slices in place via out=. standard_normal releases the GIL, so the fill scales across cores (~3x for regression, universe-dependent for panels). Reproducibility is re-keyed to (seed, n_workers, fetch_size) and stays independent of the batch_size Polars picks, so is_pure remains sound. n_workers=1 (default) keeps the exact serial path and its byte-for-byte output. Per-worker seed sub-streams are spawned lazily per block to keep seed storage at O(n_workers) and preserve bounded, streaming memory. Signed-off-by: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com> --- .../io_sources/lazy_data_generator.py | 147 ++++++++++++------ .../io_sources/test_lazy_data_generator.py | 102 ++++++++++++ 2 files changed, 203 insertions(+), 46 deletions(-) diff --git a/polars_io_tools/io_sources/lazy_data_generator.py b/polars_io_tools/io_sources/lazy_data_generator.py index d4f038a..05039c7 100644 --- a/polars_io_tools/io_sources/lazy_data_generator.py +++ b/polars_io_tools/io_sources/lazy_data_generator.py @@ -24,6 +24,19 @@ MeanComputer = Callable[[np.random.Generator, np.random.Generator, np.ndarray, int, int], tuple[np.ndarray, dict]] +def _parallel_standard_normal(rngs: list[np.random.Generator], shape: tuple[int, int], pool) -> np.ndarray: + """Fill a ``shape`` array with standard-normal draws, one contiguous row-slice per generator. + + Rows are split across ``rngs`` and each slice is filled in place via ``out=`` on ``pool``; + ``standard_normal`` releases the GIL, so the fill runs across cores. Following NumPy's + documented multithreaded-generation pattern; slices are disjoint, so the writes never race. + """ + out = np.empty(shape, dtype=np.float64) + chunks = np.array_split(out, len(rngs), axis=0) + list(pool.map(lambda rng, chunk: rng.standard_normal(out=chunk), rngs, chunks)) + return out + + def _check_reserved(name: str, kind: str, *, n_features: int, n_responses: int, use_weights: bool, extra: tuple[str, ...] = ()) -> None: """Raise ValueError if ``name`` collides with a generated column name.""" reserved = {f"x{i}" for i in range(n_features)} | {f"y{i}" for i in range(n_responses)} | ({"weight"} if use_weights else set()) | set(extra) @@ -46,11 +59,14 @@ def _register_source( mean_computer: MeanComputer, extras_schema: dict, chunk_sizes: np.ndarray | None = None, + n_workers: int = 1, explain_name: str, explain_detail: str | None = None, ) -> pl.LazyFrame: feature_cols = [f"x{i}" for i in range(n_features)] response_cols = [f"y{i}" for i in range(n_responses)] + if n_workers < 1: + raise ValueError(f"n_workers must be >= 1, got {n_workers}") schema_fields: dict[str, pl.DataType] = {name: pl.Float64() for name in feature_cols + response_cols} schema_fields.update(extras_schema) if use_weights: @@ -65,62 +81,95 @@ def source_generator( ) -> Iterator[pl.DataFrame]: bs = batch_size or fetch_size x_ss, eps_ss, w_ss, aux_ss, group_ss = np.random.SeedSequence(seed).spawn(5) - x_rng = np.random.default_rng(x_ss) - eps_rng = np.random.default_rng(eps_ss) w_rng = np.random.default_rng(w_ss) aux_rng = np.random.default_rng(aux_ss) group_rng = np.random.default_rng(group_ss) + # x/eps are the dominant cost; optionally fill them in parallel. w/aux/group + # stay single-stream (cheap, and already batch_size-independent). n_workers==1 + # keeps the exact serial streams, so default output is byte-for-byte unchanged. + threaded = n_workers > 1 and n_samples > 0 + pool = None + if threaded: + from concurrent.futures import ThreadPoolExecutor + + # Fixed-size blocks (independent of Polars' batch_size) filled in parallel by one + # persistent Generator per worker, so row i depends only on (seed, n_workers, fetch_size). + block_rows = min(n_workers * fetch_size, n_samples) + x_rngs = [np.random.default_rng(s) for s in x_ss.spawn(n_workers)] + eps_rngs = [np.random.default_rng(s) for s in eps_ss.spawn(n_workers)] + pool = ThreadPoolExecutor(max_workers=n_workers) + x_block = eps_block = None + block_pos = block_rows # trigger a fill on the first iteration + else: + x_rng = np.random.default_rng(x_ss) + eps_rng = np.random.default_rng(eps_ss) + remaining_gen = n_samples remaining_deliver = n_rows if n_rows is not None else n_samples chunk_idx = 0 rows_left_in_chunk = int(chunk_sizes[0]) if chunk_sizes is not None and len(chunk_sizes) > 0 else 0 - while remaining_gen > 0 and remaining_deliver > 0: - if chunk_sizes is not None: - k = min(bs, rows_left_in_chunk, remaining_gen) - row_in_chunk = int(chunk_sizes[chunk_idx]) - rows_left_in_chunk - else: - k = min(bs, remaining_gen) - row_in_chunk = 0 - x_batch = x_rng.standard_normal(size=(k, n_features)) - if use_weights: - w_batch = w_rng.uniform(weights_low, weights_high, size=k) - eps_scale = epsilon_scale / np.sqrt(w_batch)[:, None] - else: - w_batch = None - eps_scale = epsilon_scale - eps_batch = eps_rng.standard_normal(size=(k, n_responses)) * eps_scale + epsilon_loc - mean_batch, extras = mean_computer(aux_rng, group_rng, x_batch, chunk_idx, row_in_chunk) - y_batch = mean_batch + eps_batch - - data: dict = {name: x_batch[:, i] for i, name in enumerate(feature_cols)} - for i, name in enumerate(response_cols): - data[name] = y_batch[:, i] - data.update(extras) - if use_weights: - data["weight"] = w_batch - - df = pl.DataFrame(data) - if predicate is not None: - df = df.filter(predicate) - if with_columns is not None: - df = df.select(with_columns) - - if df.height > remaining_deliver: - df = df.head(remaining_deliver) - remaining_deliver -= df.height - remaining_gen -= k # Always k, not df.height: RNG must advance by the full generated batch to keep row-i values batch-independent. - - if chunk_sizes is not None: - rows_left_in_chunk -= k - if rows_left_in_chunk == 0 and chunk_idx + 1 < len(chunk_sizes): - chunk_idx += 1 - rows_left_in_chunk = int(chunk_sizes[chunk_idx]) - - log.debug("scan_synthetic: yielded %d rows; remaining_gen=%d, remaining_deliver=%d", df.height, remaining_gen, remaining_deliver) - yield df + try: + while remaining_gen > 0 and remaining_deliver > 0: + if chunk_sizes is not None: + k = min(bs, rows_left_in_chunk, remaining_gen) + row_in_chunk = int(chunk_sizes[chunk_idx]) - rows_left_in_chunk + else: + k = min(bs, remaining_gen) + row_in_chunk = 0 + if threaded: + if block_pos == block_rows: + x_block = _parallel_standard_normal(x_rngs, (block_rows, n_features), pool) + eps_block = _parallel_standard_normal(eps_rngs, (block_rows, n_responses), pool) + block_pos = 0 + k = min(k, block_rows - block_pos) # keep each batch within one fixed block + x_batch = x_block[block_pos : block_pos + k] + eps_raw = eps_block[block_pos : block_pos + k] + block_pos += k + else: + x_batch = x_rng.standard_normal(size=(k, n_features)) + eps_raw = eps_rng.standard_normal(size=(k, n_responses)) + if use_weights: + w_batch = w_rng.uniform(weights_low, weights_high, size=k) + eps_scale = epsilon_scale / np.sqrt(w_batch)[:, None] + else: + w_batch = None + eps_scale = epsilon_scale + eps_batch = eps_raw * eps_scale + epsilon_loc + mean_batch, extras = mean_computer(aux_rng, group_rng, x_batch, chunk_idx, row_in_chunk) + y_batch = mean_batch + eps_batch + + data: dict = {name: x_batch[:, i] for i, name in enumerate(feature_cols)} + for i, name in enumerate(response_cols): + data[name] = y_batch[:, i] + data.update(extras) + if use_weights: + data["weight"] = w_batch + + df = pl.DataFrame(data) + if predicate is not None: + df = df.filter(predicate) + if with_columns is not None: + df = df.select(with_columns) + + if df.height > remaining_deliver: + df = df.head(remaining_deliver) + remaining_deliver -= df.height + remaining_gen -= k # Always k, not df.height: RNG must advance by the full generated batch to keep row-i values batch-independent. + + if chunk_sizes is not None: + rows_left_in_chunk -= k + if rows_left_in_chunk == 0 and chunk_idx + 1 < len(chunk_sizes): + chunk_idx += 1 + rows_left_in_chunk = int(chunk_sizes[chunk_idx]) + + log.debug("scan_synthetic: yielded %d rows; remaining_gen=%d, remaining_deliver=%d", df.height, remaining_gen, remaining_deliver) + yield df + finally: + if pool is not None: + pool.shutdown(wait=False) return register_io_source_with_is_pure( source_generator, schema=schema, is_pure=seed is not None, explain_name=explain_name, explain_detail=explain_detail @@ -163,6 +212,7 @@ def scan_synthetic_regression( chunk_key: str | None = None, n_chunks: int | None = None, seed: int | None = None, + n_workers: int = 1, fetch_size: int = 10_000, description: str | None = None, ) -> pl.LazyFrame: @@ -191,6 +241,7 @@ def scan_synthetic_regression( chunk_key: If provided together with ``n_chunks``, emits a monotonic ``Int64`` column with this name whose values are ``0, 1, ..., n_chunks - 1``. Must not collide with a generated column name. n_chunks: Number of contiguous chunks to split ``n_samples`` into. Required when ``chunk_key`` is set. Must satisfy ``1 <= n_chunks <= n_samples``. seed: Seed for ``np.random.default_rng``. If None, uses fresh entropy per call (and the source is registered with ``is_pure=False``). + n_workers: Number of threads used to fill the ``x``/``y`` Gaussian draws (the dominant cost) in parallel. ``1`` (default) keeps the fully serial path and its exact output. With ``n_workers > 1`` the draws are generated in fixed-size blocks split across threads, so reproducibility is keyed on ``(seed, n_workers, fetch_size)`` and remains independent of the ``batch_size`` Polars chooses; values differ from the serial path. Peak generation memory grows with ``n_workers`` (a block holds ``n_workers * fetch_size`` rows). Most useful when ``n_samples`` is large. fetch_size: Default number of rows generated per batch when Polars does not provide a ``batch_size``. Must be >= 1. Defaults to 10_000. description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). """ @@ -260,6 +311,7 @@ def mean_computer( mean_computer=mean_computer, extras_schema=extras_schema, chunk_sizes=chunk_sizes, + n_workers=n_workers, explain_name="scan_synthetic_regression", explain_detail=description, ) @@ -282,6 +334,7 @@ def scan_synthetic_panel( epsilon_loc: float = 0.0, epsilon_scale: float = 1.0, seed: int | None = None, + n_workers: int = 1, fetch_size: int = 10_000, description: str | None = None, ) -> pl.LazyFrame: @@ -310,6 +363,7 @@ def scan_synthetic_panel( epsilon_loc: Mean of the Gaussian noise. Defaults to 0.0. epsilon_scale: When ``use_weights=False``, the noise stddev for every row. When ``use_weights=True``, the *reference* stddev at ``w=1``; actual per-row noise is ``N(loc, (epsilon_scale/√w)²)``. Must be >= 0. Defaults to 1.0. seed: Seed for ``np.random.default_rng``. If None, uses fresh entropy per call (and the source is registered with ``is_pure=False``). + n_workers: Number of threads used to fill the ``x``/``y`` Gaussian draws (the dominant cost) in parallel. ``1`` (default) keeps the fully serial path and its exact output. With ``n_workers > 1`` the draws are generated in fixed-size blocks split across threads, so reproducibility is keyed on ``(seed, n_workers, fetch_size)`` and remains independent of the ``batch_size`` Polars chooses; values differ from the serial path. Peak generation memory grows with ``n_workers`` (a block holds ``n_workers * fetch_size`` rows). For panels the per-date batch is ``n_symbols`` rows, so parallelism only helps when ``n_symbols`` is large. fetch_size: Default number of rows generated per batch when Polars does not provide a ``batch_size``. Must be >= 1. Defaults to 10_000. description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). """ @@ -446,6 +500,7 @@ def mean_computer( mean_computer=mean_computer, extras_schema=extras_schema, chunk_sizes=chunk_sizes, + n_workers=n_workers, explain_name="scan_synthetic_panel", explain_detail=description, ) diff --git a/polars_io_tools/tests/io_sources/test_lazy_data_generator.py b/polars_io_tools/tests/io_sources/test_lazy_data_generator.py index 6ab5c7b..75b5cea 100644 --- a/polars_io_tools/tests/io_sources/test_lazy_data_generator.py +++ b/polars_io_tools/tests/io_sources/test_lazy_data_generator.py @@ -342,3 +342,105 @@ def test_panel_betas_1d_promotion(): X = df.select("x0", "x1", "x2").to_numpy() y_expected = X @ np.array([1.0, 2.0, 3.0]) assert np.allclose(df["y0"].to_numpy(), y_expected, atol=1e-12) + + +def test_n_workers_determinism_and_schema_parity(): + # Fixed (seed, n_workers) reproduces exactly; schema/shape match the serial path. + kwargs = {"n_samples": 12_345, "n_features": 6, "n_responses": 2, "use_weights": True, "seed": 7} + serial = scan_synthetic_regression(**kwargs, n_workers=1).collect() + a = scan_synthetic_regression(**kwargs, n_workers=4).collect() + b = scan_synthetic_regression(**kwargs, n_workers=4).collect() + assert a.equals(b) + assert a.schema == serial.schema + assert a.shape == serial.shape + # Threading changes the RNG stream layout, so values differ from the serial path. + assert not a.equals(serial) + + +def test_n_workers_batch_size_independence(): + # With n_workers > 1, reproducibility is keyed on (seed, n_workers, fetch_size): blocks are + # sized n_workers * fetch_size, so fetch_size is part of the contract (like seed). What must + # still hold for is_pure is independence from the *runtime* batch_size Polars chooses at a + # fixed fetch_size -- verified here by comparing the in-memory and streaming engines. + kwargs = {"n_samples": 4000, "n_features": 4, "seed": 42, "n_workers": 8, "fetch_size": 512} + a = scan_synthetic_regression(**kwargs).collect() + b = scan_synthetic_regression(**kwargs).collect(engine="streaming") + assert a.equals(b) + + +def test_n_workers_beta_recovery(): + betas = np.array([1.0, -2.0, 0.5, 3.0]) + df = scan_synthetic_regression(n_samples=200_000, n_features=4, n_responses=1, betas=betas, epsilon_scale=0.1, seed=3, n_workers=8).collect() + features = df.select("x0", "x1", "x2", "x3").to_numpy() + labels = df.select("y0").to_numpy() + fitted = LinearRegression().fit(features, labels).coef_ + assert np.max(np.abs(fitted.reshape(-1) - betas)) < 0.02 + + +def test_n_workers_pushdowns(): + kwargs = {"n_samples": 5000, "n_features": 5, "n_responses": 2, "seed": 1, "n_workers": 8} + assert scan_synthetic_regression(**kwargs).head(321).collect().height == 321 + assert scan_synthetic_regression(**kwargs).select(["x0", "y1"]).collect().columns == ["x0", "y1"] + + +def test_n_workers_panel_parity(): + kwargs = { + "start_date": date(2020, 1, 1), + "end_date": date(2020, 2, 1), + "n_symbols": 500, + "n_features": 4, + "seed": 5, + } + serial = scan_synthetic_panel(**kwargs, n_workers=1).collect() + a = scan_synthetic_panel(**kwargs, n_workers=4).collect() + b = scan_synthetic_panel(**kwargs, n_workers=4).collect() + assert a.equals(b) + assert a.schema == serial.schema + assert a.shape == serial.shape + + +def test_n_workers_invalid(): + with pytest.raises(ValueError, match="n_workers must be >= 1"): + scan_synthetic_regression(n_samples=10, n_features=2, seed=1, n_workers=0) + with pytest.raises(ValueError, match="n_workers must be >= 1"): + scan_synthetic_panel(start_date=date(2020, 1, 1), end_date=date(2020, 1, 3), n_features=2, seed=1, n_workers=0) + + +def test_n_workers_multi_block_batch_size_independence(): + # n_samples spans several fixed blocks (block_rows = n_workers * fetch_size = 4 * 250 = 1000), + # exercising the block-cursor refill and boundary cap. Output must be identical whether Polars + # drives it in memory or via the streaming engine (different runtime batch sizes) at fixed + # fetch_size, and must not depend on that batch size. + kwargs = {"n_samples": 3300, "n_features": 3, "n_responses": 2, "use_weights": True, "seed": 9, "n_workers": 4, "fetch_size": 250} + a = scan_synthetic_regression(**kwargs).collect() + b = scan_synthetic_regression(**kwargs).collect(engine="streaming") + assert a.equals(b) + assert a.height == 3300 + + +def test_n_workers_panel_batch_size_independence_full(): + # Threaded panel with categories, group_by, and weights: x/eps come from the block filler + # while aux/group/weight stay per-batch serial streams -- verify they stay row-aligned and + # the whole frame is independent of the runtime batch_size (in-memory vs streaming engine). + kwargs = { + "start_date": date(2020, 1, 1), + "end_date": date(2020, 3, 1), + "n_symbols": 137, # not a divisor of block_rows, so date chunks straddle blocks + "n_features": 4, + "categories": [["A", "B", "C"]], + "group_by": ("g", ["x", "y"]), + "use_weights": True, + "seed": 21, + "n_workers": 8, + "fetch_size": 64, + } + a = scan_synthetic_panel(**kwargs).collect() + b = scan_synthetic_panel(**kwargs).collect(engine="streaming") + assert a.equals(b) + + +def test_n_workers_head_pushdown_terminates_early(): + # An early n_rows pushdown must return promptly and shut its pool down via the generator's + # finally, without generating the full declared dataset. + df = scan_synthetic_regression(n_samples=50_000_000, n_features=4, seed=1, n_workers=4).head(5).collect() + assert df.height == 5