|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +# SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +"""Benchmark input files. |
| 5 | +
|
| 6 | +Run ``python -m benchmark.data <dir>`` to write the files before a measured run, then set |
| 7 | +``VORTEX_BENCH_DATA_DIR=<dir>``. Under CodSpeed simulation all of the process runs on Valgrind, |
| 8 | +so data written ahead of time keeps the slow setup out of the run. |
| 9 | +""" |
| 10 | + |
| 11 | +import random |
| 12 | +import sys |
| 13 | +from array import array |
| 14 | +from pathlib import Path |
| 15 | +from typing import Any |
| 16 | + |
| 17 | +import pyarrow as pa |
| 18 | +import pyarrow.compute as pc |
| 19 | + |
| 20 | +import vortex as vx |
| 21 | + |
| 22 | +LENGTH = 100_000 |
| 23 | + |
| 24 | +COLUMN_SETS: dict[str, tuple[str, ...]] = { |
| 25 | + "int": ("x",), |
| 26 | + "int_str": ("x", "y"), |
| 27 | + "int_float": ("x", "z"), |
| 28 | + "int_str_float": ("x", "y", "z"), |
| 29 | +} |
| 30 | + |
| 31 | + |
| 32 | +def _hex_strings(length: int) -> pa.StringArray: |
| 33 | + # 32 random hex characters for each row, built from buffers without a Python loop per row. |
| 34 | + width = 32 |
| 35 | + chars = random.Random(0).randbytes(length * width // 2).hex().encode() |
| 36 | + offsets = array("i", range(0, length * width + 1, width)) |
| 37 | + |
| 38 | + return pa.StringArray.from_buffers(length, pa.py_buffer(offsets), pa.py_buffer(chars)) |
| 39 | + |
| 40 | + |
| 41 | +def table(columns: tuple[str, ...], length: int = LENGTH) -> pa.Table: |
| 42 | + x = pa.array(range(length), type=pa.int64()) |
| 43 | + arrays: list[pa.Array[Any] | pa.ChunkedArray[Any]] = [x] |
| 44 | + |
| 45 | + if "y" in columns: |
| 46 | + arrays.append(_hex_strings(length)) |
| 47 | + if "z" in columns: |
| 48 | + arrays.append(pc.sqrt(x)) |
| 49 | + |
| 50 | + return pa.table(arrays, names=list(columns)) |
| 51 | + |
| 52 | + |
| 53 | +def write(directory: Path, name: str) -> Path: |
| 54 | + """Write the file for ``name`` into ``directory``, unless it is already there.""" |
| 55 | + path = directory / f"{name}.vortex" |
| 56 | + |
| 57 | + if not path.exists(): |
| 58 | + directory.mkdir(parents=True, exist_ok=True) |
| 59 | + vx.io.write(table(COLUMN_SETS[name]), str(path)) |
| 60 | + |
| 61 | + return path |
| 62 | + |
| 63 | + |
| 64 | +if __name__ == "__main__": |
| 65 | + for name in COLUMN_SETS: |
| 66 | + print(write(Path(sys.argv[1]), name)) |
0 commit comments