Bodo/bodo/tests/test_df_lib/test_tpch.py at 505539317b853909ca854f714db5c5e69cc81542 · bodo-ai/Bodo · 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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
import numpy as np
import pandas as pd
import pytest

import benchmarks.tpch.dataframe_lib as tpch
import bodo.pandas as bd
from bodo.pandas.plan import assert_executed_plan_count
from bodo.tests.utils import _test_equal

pytestmark = pytest.mark.jit_dependency

datapath = "bodo/tests/data/tpch-test_data/parquet"


def run_tpch_query_test(query_func, plan_executions=0, ctes_created=0):
    """Run a tpch query and compare output to Pandas.

    Args:
        query_func (Callable): The callable object that takes in dataframes loaded from
          TPCH and returns an output dataframe.
        plan_executions (int, optional): Expected number of LazyPlans to be executed.
          Defaults to 0.
    """

    # Scale factor is set to 1.0 for testing purposes in query 11
    pd_kwargs = {"pd": pd}
    pd_args = [
        getattr(tpch, f"load_{key}")(datapath, **pd_kwargs)
        for key in tpch._query_to_args[int(query_func.__name__[-2:])]
        if key != "scale_factor"
    ]
    bd_args = [bd.from_pandas(df) for df in pd_args]

    pd_result = query_func(*pd_args, **pd_kwargs)

    with assert_executed_plan_count(plan_executions):
        bd_result = query_func(*bd_args)

    # We can't capture all CTEs created because the above should create
    # and execute plans if plan_executions > 0.  If those plans executed
    # above have CTEs then we can't capture them since they may occur on
    # workers.
    generated_ctes = bd_result._plan.get_cte_count()
    assert generated_ctes == ctes_created

    if isinstance(
        pd_result,
        (
            pd.DataFrame,
            pd.Series,
        ),
    ):
        _test_equal(
            bd_result,
            pd_result,
            check_pandas_types=False,
            sort_output=True,
            reset_index=True,
        )
    else:
        # For scalar or numeric results
        assert np.isclose(pd_result, bd_result)


def test_tpch_q01():
    run_tpch_query_test(tpch.tpch_q01)


def test_tpch_q02():
    run_tpch_query_test(tpch.tpch_q02, ctes_created=1)


def test_tpch_q03():
    run_tpch_query_test(tpch.tpch_q03)


def test_tpch_q04():
    run_tpch_query_test(tpch.tpch_q04)


def test_tpch_q05():
    run_tpch_query_test(tpch.tpch_q05)


def test_tpch_q06():
    run_tpch_query_test(tpch.tpch_q06, plan_executions=1)


def test_tpch_q07():
    run_tpch_query_test(tpch.tpch_q07)


def test_tpch_q08():
    run_tpch_query_test(tpch.tpch_q08)


def test_tpch_q09():
    run_tpch_query_test(tpch.tpch_q09)


def test_tpch_q10():
    run_tpch_query_test(tpch.tpch_q10)


def test_tpch_q11():
    run_tpch_query_test(tpch.tpch_q11, ctes_created=1)


def test_tpch_q12():
    run_tpch_query_test(tpch.tpch_q12)


def test_tpch_q13():
    run_tpch_query_test(tpch.tpch_q13)


def test_tpch_q14():
    run_tpch_query_test(tpch.tpch_q14, plan_executions=1)


def test_tpch_q15():
    run_tpch_query_test(tpch.tpch_q15, ctes_created=1)


def test_tpch_q16():
    run_tpch_query_test(tpch.tpch_q16)


def test_tpch_q17():
    run_tpch_query_test(tpch.tpch_q17, plan_executions=1)


def test_tpch_q18():
    run_tpch_query_test(tpch.tpch_q18)


def test_tpch_q19():
    run_tpch_query_test(tpch.tpch_q19, plan_executions=1)


def test_tpch_q20():
    run_tpch_query_test(tpch.tpch_q20)


def test_tpch_q21():
    run_tpch_query_test(tpch.tpch_q21, ctes_created=1)


def test_tpch_q22():
    run_tpch_query_test(tpch.tpch_q22, ctes_created=1)