Function make_dataset

python/benchmarks/benchmark_data_processing.py:147–169

(rows_per_symbol: int, symbols: list[str])

def make_dataset(rows_per_symbol: int, symbols: list[str]) -> pl.DataFrame:
    frames: list[pl.DataFrame] = []
    ts_base = pl.datetime_range(
        start=pl.datetime(2020, 1, 1),
        end=pl.datetime(2020, 1, 1) + pl.duration(minutes=rows_per_symbol - 1),
        interval="1m",
        eager=True,
    )
    for symbol in symbols:
        frames.append(
            pl.DataFrame(
                {
                    "ts": ts_base,
                    "symbol": [symbol] * rows_per_symbol,
                    "open": pl.arange(0, rows_per_symbol, eager=True).cast(pl.Float64) + 100.0,
                    "high": pl.arange(0, rows_per_symbol, eager=True).cast(pl.Float64) + 100.5,
                    "low": pl.arange(0, rows_per_symbol, eager=True).cast(pl.Float64) + 99.5,
                    "close": pl.arange(0, rows_per_symbol, eager=True).cast(pl.Float64) + 100.2,
                    "volume": pl.repeat(1000.0, rows_per_symbol, eager=True),
                }
            )
        )
    return pl.concat(frames, rechunk=True)


def run_benchmarks(rows_per_symbol: int, symbols: int, iterations: int) -> dict[str, Any]: