-
Notifications
You must be signed in to change notification settings - Fork 200
Expand file tree
/
Copy pathbench.py
More file actions
62 lines (47 loc) · 1.87 KB
/
Copy pathbench.py
File metadata and controls
62 lines (47 loc) · 1.87 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright the Vortex contributors
"""
Simple benchmarks
"""
from __future__ import annotations
import io
from collections.abc import Callable
import pyarrow as pa
import pyarrow.parquet as pq
import pytest
import vortex as vx
@pytest.fixture(params=[10, 100])
def vortex_array(request): # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
rows: list[dict[str, list[int | None]]] = []
for _ in range(1_000):
r: dict[str, list[int | None]] = {}
for col in range(request.param): # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
# Create large arrays of length 100 for each column.
r[f"col{col}"] = [1, 2, None, 4] * 25
rows.append(r)
return vx.array(rows)
@pytest.fixture(params=[10, 100])
def arrow_array(request): # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
rows: list[dict[str, list[int | None]]] = []
for _ in range(1_000):
r: dict[str, list[int | None]] = {}
for col in range(request.param): # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
# Create large arrays of length 100 for each column.
r[f"col{col}"] = [1, 2, None, 4] * 25
rows.append(r)
return pa.Table.from_pylist(rows)
def test_compress_vortex(
benchmark: Callable[[Callable[[], None]], None],
vortex_array: vx.Array,
):
def compress():
_ = vx.compress(vortex_array)
benchmark(compress)
def test_compress_parquet(
benchmark: Callable[[Callable[[], None]], None], arrow_array: pa.Array[pa.Scalar[pa.DataType]]
):
def compress():
# write to bytes in memory.
bout = io.BytesIO()
pq.write_table(arrow_array, bout) # pyright: ignore[reportArgumentType, reportUnknownMemberType]
benchmark(compress)