cudf_benchmarks/bench_sort.py at master · skirui-source/cudf_benchmarks · GitHub

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
import cudf
import cupy as cp
import pytest


@pytest.mark.parametrize("cls", [cudf.Series, cudf.Index])
@pytest.mark.parametrize("N", [1_000, 100_000, 10_000_000])
def test_series_index_argsort(benchmark, cls, N):
    obj = cls(cp.random.rand(N))
    benchmark(obj.argsort)


@pytest.mark.parametrize("cls", [cudf.Series, cudf.Index])
@pytest.mark.parametrize("N", [1_000, 100_000, 10_000_000])
def test_series_index_sort_values(benchmark, cls, N):
    obj = cls(cp.random.rand(N))
    benchmark(obj.sort_values)


@pytest.mark.parametrize("N", [1_000, 100_000, 10_000_000])
@pytest.mark.parametrize("n", [10])
def test_series_nsmallest(benchmark, N, n):
    ser = cudf.Series(cp.random.rand(N))
    benchmark(ser.nsmallest, n)


@pytest.mark.parametrize("N", [1_000, 100_000, 10_000_000])
@pytest.mark.parametrize("ncol", [5, 10])
# @pytest.mark.parametrize("ncol_sort", [1, 2, 3])
@pytest.mark.parametrize("ncol_sort", [1])
def test_dataframe_argsort(benchmark, N, ncol, ncol_sort):
    df = cudf.DataFrame({i: cp.random.rand(N) for i in range(ncol)})
    benchmark(df.argsort)


@pytest.mark.parametrize("N", [1_000, 100_000, 10_000_000])
@pytest.mark.parametrize("ncol", [5, 10])
# @pytest.mark.parametrize("ncol_sort", [1, 2, 3])
@pytest.mark.parametrize("ncol_sort", [1])
def test_dataframe_sort_values(benchmark, N, ncol, ncol_sort):
    df = cudf.DataFrame({i: cp.random.rand(N) for i in range(ncol)})
    benchmark(df.sort_values, [i for i in range(ncol_sort)])


@pytest.mark.parametrize("N", [1_000, 100_000, 10_000_000])
@pytest.mark.parametrize("ncol", [5, 10])
# @pytest.mark.parametrize("ncol_sort", [1, 2, 3])
@pytest.mark.parametrize("ncol_sort", [1])
@pytest.mark.parametrize("n", [10])
def test_dataframe_nsmallest(benchmark, N, ncol, ncol_sort, n):
    df = cudf.DataFrame({i: cp.random.rand(N) for i in range(ncol)})
    benchmark(df.nsmallest, n, [i for i in range(ncol_sort)])


def test_rangeindex_argsort(benchmark, rangeindex):
    benchmark(rangeindex.argsort)