Run TPC-H on Databricks

[jacobtomlinson]

Reproduce the TPC-H work that @rjzamora has been looking into with Dask+cuDF and the TPC-H work Coiled has been doing with Dask but on Databricks.

https://tpch.coiled.io/

[rjzamora]

cc @charlesbluca

[skirui-source]

https://medium.com/coiled-hq/tpc-h-benchmarks-for-query-optimization-with-dask-expressions-c42921aafbd3

[skirui-source]

Initially I was running into this bug, which i resolved by downgrading from pytest v8.0.1, had to downgrade to v7.4.4

[skirui-source]

@jacobtomlinson just to clarify, is the goal here to:

1. run TPC-H benchmarks with dask-expr + cudf in Databricks?
2. run TPC-H benchmarks with dask + coiled in Databricks?

You also mentioned something about running benchmarks for cudf.pandas, is that part of this?

[skirui-source]

@rjzamora with your recent work merged into rapidsai/cudf, what are the steps to recreate the test environment you used to produce the benchmark results?

[jacobtomlinson]

The primary goal is dask-expr + cudf on Databricks. Secondary goal is cudf.pandas on Databricks.

[rjzamora]

Thanks for working on this @skirui-source !

what are the steps to recreate the test environment you used to produce the benchmark results?

I haven't updated my tpch-rapids branch of coiled/benchmarks to align with main, but the latest version of cudf-24.06 seems to work fine with that branch. You can also use cudf-24.04, but need to set the DASK_DATAFRAME__QUERY_PLANNING=True environment variable before running to use dask-expr under the hood (setting False like this will disable dask-expr for both 24.04 and 24.6, but it is turned on by default in 24.06).

To run locally, I started with a 24.04 or 24.06 rapids environment. Then did the following:

• Cloned https://github.com/rjzamora/coiled-benchmarks/tree/tpch-rapids
• Installed further conda dependencies: conda install -c canda-forge coiled s3fs altair alembic
• Installed further pip dependencies: pip install pytest-timeout pytest-xdist
• Change the LOCAL_SF*_PATH variables in benchmarks/tests/tpch/utils.py to point to your local data
  • For sf100, feel free to use my data on the ase machines: /datasets/rzamora/tpch-data/tables_scale_100

Then, from within benchmarks, you can execute a query:

export DASK_DATAFRAME__QUERY_PLANNING=True
pytest --benchmark  tests/tpch/test_dask.py::test_query_1 --local --rapids --scale 100

The timing result will be appended to a benchmark.db file. I've been using the following python code to convert that file to a pandas DataFrame (but, I'm sure there is a better way):

import pandas as pd
import sqlite3
con = sqlite3.connect("./benchmark.db")
df = pd.read_sql_query("SELECT * from test_run", con)
print(df[["name", "duration"]])
con.close()

Important notes on the data:

I believe I needed to jump through a few annoying hoops to both generate the data and modify the code to handle differences between my data and the s3 data used by Coiled. For example, my data currently uses a single parquet file for each table, while Coiled uses a directory of files (which definitely makes more sense beyond sf100). This is why my benchmarks branch has a _read_parquet "hack" in benchmarks/tests/tpch/dask_queries.py, where I add a ".parquet" suffix to the table path.

The primary goal is dask-expr + cudf on Databricks. Secondary goal is cudf.pandas on Databricks.

I would definitely focus on dask-expr + cudf for now. I believe pandas->arrow->pandas conversion is still a problem in cudf.pandas, and dask/dask-expr will do this a lot.

[skirui-source]

initial batch of results (tested on dgx14):

(rapids-24.06) skirui@dgx14:/datasets/skirui/GitRepos/coiled-benchmarks$ py.test --benchmark tests/tpch/test_dask.py --local --rapids --scale 10

In [3]: print(df[["name", "duration"]])
             name  duration
1    test_query_1  9.345021
2    test_query_2       NaN
3    test_query_3  6.615860
4    test_query_4  4.828701
5    test_query_5  7.316751
6    test_query_6  5.451433
7    test_query_7  7.176501
8    test_query_8  7.961834
9    test_query_9  6.981508
10  test_query_10  6.611681
11  test_query_11  3.811042
12  test_query_12  5.117650
13  test_query_13       NaN
14  test_query_14  5.732556
15  test_query_15  5.727635
16  test_query_16  4.815220
17  test_query_17  6.078448
18  test_query_18  7.685056
19  test_query_19  5.934201
20  test_query_20  5.236853
21  test_query_21  7.247067
22  test_query_22  4.751128

(rapids-24.06) skirui@dgx14:/datasets/skirui/GitRepos/coiled-benchmarks$ py.test --benchmark tests/tpch/test_dask.py --local --rapids --scale 100

In [3]: print(df[["name", "duration"]])
             name   duration
23   test_query_1  14.323813
24   test_query_2   7.302116
25   test_query_3        NaN
26   test_query_4  18.434659
27   test_query_5  19.057925
28   test_query_6   5.927966
29   test_query_7  16.077179
30   test_query_8  16.161171
31   test_query_9  27.749285
32  test_query_10  15.439956
33  test_query_11   5.020720
34  test_query_12   9.845365
35  test_query_13        NaN
36  test_query_14   6.785045
37  test_query_15  11.023882
38  test_query_16        NaN
39  test_query_17  21.482339
40  test_query_18  28.699488
41  test_query_19   7.627631
42  test_query_20        NaN
43  test_query_21  24.507490
44  test_query_22   5.938541

(rapids-24.06) skirui@dgx14:/datasets/skirui/GitRepos/coiled-benchmarks$ py.test --benchmark tests/tpch/test_dask.py --local --rapids --scale 1

             name     duration
45   test_query_1     7.262008
46   test_query_2     4.105303
47   test_query_3     4.405016
48   test_query_4          NaN
49   test_query_5     4.382684
50   test_query_6     3.901139
51   test_query_7          NaN
52   test_query_8     4.996450
53   test_query_9     4.708150
54  test_query_10          NaN
55  test_query_11     3.718096
56  test_query_12     3.900132
57  test_query_13          NaN
58  test_query_14     4.186984
59  test_query_15     4.032958
60  test_query_16          NaN
61  test_query_17     4.383530
62  test_query_18     4.013830
63  test_query_19  3597.592139
64  test_query_20     5.858231
65  test_query_21     5.473735
66  test_query_22     5.003113