Repository: coiled/benchmarks (language: python)

Issue:
[TPC-H] Query results are never validated
While reviewing the implementation of the remaining TPC-H queries, I've identified multiple subtle mistakes. This is a cumbersome process that doesn't guarantee that every subtlety is caught. We should add tests that assert the correctness of our queries. DuckDB contains all the data we need for that: https://duckdb.org/docs/extensions/tpch.html#listing-expected-answers


Discussion / hints:

Fix patch:
diff --git a/ci/environment-tpch-nondask.yml b/ci/environment-tpch-nondask.yml
index f31a5db60f..fee9975ef6 100644
--- a/ci/environment-tpch-nondask.yml
+++ b/ci/environment-tpch-nondask.yml
@@ -12,5 +12,4 @@ dependencies:
   - protobuf ==4.24.4
 
   # Other TPCH tests
-  - python-duckdb ==0.9.2
   - polars ==0.19.19
diff --git a/setup.cfg b/setup.cfg
index e978e22351..d100bdec89 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -22,6 +22,7 @@ markers =
     shuffle_tasks: legacy tasks-based shuffle engine
     tpch_dask: dask implementation of the TPCH tests suite
     tpch_nondask: competitors' (not dask) implementation of the TPCH test suite
+    tpch_correctness: verify correctness of the dask implementation of the TPCH tests suite
 
 # pytest-timeout settings
 # 'thread' kills off the whole test suite. 'signal' only kills the offending test.


Test patch:
diff --git a/ci/environment-test.yml b/ci/environment-test.yml
index 8f393499a3..9e15ba3294 100644
--- a/ci/environment-test.yml
+++ b/ci/environment-test.yml
@@ -14,3 +14,5 @@ dependencies:
   - pytest-timeout
   - pytest-xdist
   - pyyaml
+  # TPC-H correctness test and DuckDB implementation
+  - python-duckdb ==0.9.2
diff --git a/tests/tpch/README.md b/tests/tpch/README.md
index 7702c92be9..afdf3e6384 100644
--- a/tests/tpch/README.md
+++ b/tests/tpch/README.md
@@ -45,7 +45,7 @@ If you want to run locally, you'll need to generate data.  Run the following
 from the **root directory** of this repository.
 
 ```
-python tests/tpch/generate-data.py --scale 10
+python tests/tpch/generate_data.py --scale 10
 ```
 
 Run Many Tests
diff --git a/tests/tpch/generate_answers.py b/tests/tpch/generate_answers.py
new file mode 100644
index 0000000000..ffb80826ca
--- /dev/null
+++ b/tests/tpch/generate_answers.py
@@ -0,0 +1,53 @@
+import pathlib
+
+import click
+import duckdb
+
+VERIFICATION_SCALE = 1
+
+
+def generate(base_path: pathlib.Path) -> pathlib.Path:
+    scale = VERIFICATION_SCALE
+    path = base_path / "answers" / f"scale-{scale}"
+    path.mkdir(parents=True, exist_ok=True)
+
+    print(f"Scale: {scale}, Path: {path}")
+
+    with duckdb.connect() as con:
+        con.install_extension("tpch")
+        con.load_extension("tpch")
+
+        print("Exporting TPC-H answers...")
+        answers = con.sql(
+            f"""
+            SELECT
+                query_nr,
+                answer
+            FROM tpch_answers()
+            WHERE scale_factor = {scale};
+            """
+        ).df()
+
+        for tpl in answers.itertuples():
+            query = tpl.query_nr
+            answer = tpl.answer
+            with (path / f"q{query}.out").open(mode="w") as f:
+                f.write(answer)
+        print("Finished exporting all answers!")
+    return path
+
+
+@click.command()
+@click.option(
+    "--path",
+    default="./tests/tpch",
+    help="Local or S3 base path, will affix '/answers' subdirectory to this path",
+)
+def main(
+    path: str,
+):
+    generate(pathlib.Path(path))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/tpch/generate-data.py b/tests/tpch/generate_data.py
similarity index 99%
rename from tests/tpch/generate-data.py
rename to tests/tpch/generate_data.py
index 0ad83a1b99..0ce8751c5d 100644
--- a/tests/tpch/generate-data.py
+++ b/tests/tpch/generate_data.py
@@ -33,7 +33,7 @@ def generate(
     path: str = "./tpch-data",
     relaxed_schema: bool = False,
     compression: CompressionCodec = CompressionCodec.SNAPPY,
-):
+) -> str:
     if str(path).startswith("s3"):
         path += "/" if not path.endswith("/") else ""
         path += f"scale-{scale}{'-strict' if not relaxed_schema else ''}/"
@@ -71,6 +71,7 @@ def generate(
                 client.gather(jobs)
     else:
         _tpch_data_gen(step=None, **kwargs)
+    return str(path)
 
 
 def retry(f):
diff --git a/tests/tpch/test_correctness.py b/tests/tpch/test_correctness.py
new file mode 100644
index 0000000000..bed30ff957
--- /dev/null
+++ b/tests/tpch/test_correctness.py
@@ -0,0 +1,102 @@
+import pathlib
+
+import pandas as pd
+import pytest
+from distributed import LocalCluster
+
+from .generate_answers import generate as generate_answers
+from .generate_data import generate as generate_data
+
+VERIFICATION_SCALE = 1
+
+pytestmark = pytest.mark.tpch_correctness
+
+
+@pytest.fixture(scope="module")
+def answers_path(tmp_path_factory):
+    path = tmp_path_factory.mktemp("answers")
+    return generate_answers(base_path=path)
+
+
+@pytest.fixture(scope="module")
+def data_path(tmp_path_factory):
+    path = tmp_path_factory.mktemp("data")
+    scale = VERIFICATION_SCALE
+    return pathlib.Path(generate_data(scale=scale, path=str(path), relaxed_schema=True))
+
+
+@pytest.fixture(scope="module")
+def cluster():
+    with LocalCluster() as cluster:
+        yield cluster
+
+
+@pytest.fixture(scope="module")
+def client(cluster, restart):
+    with cluster.get_client() as client:
+        if restart:
+            client.restart()
+        client.run(lambda: None)
+
+        yield client
+
+
+def get_expected_answer(query: int, answer_dir: pathlib.Path):
+    answer = pd.read_csv(
+        answer_dir / f"q{query}.out",
+        sep="|",
+    )
+    answer = answer.rename(columns=lambda x: x.strip())
+    if "o_orderdate" in answer.columns:
+        answer["o_orderdate"] = pd.to_datetime(answer["o_orderdate"])
+
+    return answer
+
+
+def verify_result(result: pd.DataFrame, query: int, answer_dir: pathlib.Path):
+    expected = get_expected_answer(query, answer_dir)
+    result = result.reset_index(drop=True)
+
+    # The expected answers are provided as whitespace-padded pipe-separated data.
+    # We must therefore strip both the expected as well as the actual answer.
+    for column, dtype in expected.dtypes.items():
+        if pd.api.types.is_object_dtype(dtype):
+            expected[column] = expected[column].apply(lambda x: x.strip())
+            result[column] = result[column].apply(lambda x: x.strip())
+    pd.testing.assert_frame_equal(result, expected, check_dtype=False)
+
+
+@pytest.mark.tpch_correctness
+@pytest.mark.parametrize(
+    "query",
+    [
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+        7,
+        pytest.param(8, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(9, marks=pytest.mark.xfail(reason="Result is wrong")),
+        pytest.param(10, marks=pytest.mark.xfail(reason="Result is wrong")),
+        pytest.param(11, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(12, marks=pytest.mark.xfail(reason="Result is wrong")),
+        pytest.param(13, marks=pytest.mark.xfail(reason="Result is wrong")),
+        pytest.param(14, marks=pytest.mark.xfail(reason="Result is wrong")),
+        pytest.param(15, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(16, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(17, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(18, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(19, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(20, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(21, marks=pytest.mark.skip(reason="Not implemented")),
+        pytest.param(22, marks=pytest.mark.skip(reason="Not implemented")),
+    ],
+)
+def test_dask_results(query, client, answers_path, data_path):
+    from . import test_dask
+
+    func = getattr(test_dask, f"test_query_{query}")
+    result = func(client, str(data_path) + "/", None)
+    verify_result(result, query, answers_path)
diff --git a/tests/tpch/test_dask.py b/tests/tpch/test_dask.py
index d27a7f1bfc..3aa050c372 100644
--- a/tests/tpch/test_dask.py
+++ b/tests/tpch/test_dask.py
@@ -1,5 +1,6 @@
 from datetime import datetime, timedelta
 
+import pandas as pd
 import pytest
 
 pytestmark = pytest.mark.tpch_dask
@@ -41,7 +42,9 @@ def test_query_1(client, dataset_path, fs):
         }
     )
 
-    total.reset_index().sort_values(["l_returnflag", "l_linestatus"]).compute()
+    result = total.reset_index().sort_values(["l_returnflag", "l_linestatus"]).compute()
+
+    return result
 
 
 @pytest.mark.shuffle_p2p
@@ -84,33 +87,37 @@ def test_query_2(client, dataset_path, fs):
         right_on=["P_PARTKEY_CPY", "MIN_SUPPLYCOST"],
         how="inner",
     )
-    merged_df[
-        [
-            "s_acctbal",
-            "s_name",
-            "n_name",
-            "p_partkey",
-            "p_mfgr",
-            "s_address",
-            "s_phone",
-            "s_comment",
+    result = (
+        merged_df[
+            [
+                "s_acctbal",
+                "s_name",
+                "n_name",
+                "p_partkey",
+                "p_mfgr",
+                "s_address",
+                "s_phone",
+                "s_comment",
+            ]
         ]
-    ].compute().sort_values(
-        by=[
-            "s_acctbal",
-            "n_name",
-            "s_name",
-            "p_partkey",
-        ],
-        ascending=[
-            False,
-            True,
-            True,
-            True,
-        ],
-    ).head(
-        100
+        .compute()
+        .sort_values(
+            by=[
+                "s_acctbal",
+                "n_name",
+                "s_name",
+                "p_partkey",
+            ],
+            ascending=[
+                False,
+                True,
+                True,
+                True,
+            ],
+        )
+        .head(100)
     )
+    return result
 
 
 @pytest.mark.shuffle_p2p
@@ -134,9 +141,13 @@ def test_query_3(client, dataset_path, fs):
     total = jn2.groupby(["l_orderkey", "o_orderdate", "o_shippriority"])[
         "revenue"
     ].sum()
-    total.reset_index().sort_values(["revenue"], ascending=False).head(10)[
-        ["l_orderkey", "revenue", "o_orderdate", "o_shippriority"]
-    ]
+    result = (
+        total.reset_index()
+        .sort_values(["revenue"], ascending=False)
+        .head(10)[["l_orderkey", "revenue", "o_orderdate", "o_shippriority"]]
+    )
+
+    return result
 
 
 @pytest.mark.shuffle_p2p
@@ -160,7 +171,9 @@ def test_query_4(client, dataset_path, fs):
         .reset_index()
         .sort_values(["o_orderpriority"])
     )
-    result_df.rename(columns={"o_orderkey": "order_count"}).compute()
+    result = result_df.rename(columns={"o_orderkey": "order_count"}).compute()
+
+    return result
 
 
 @pytest.mark.shuffle_p2p
@@ -190,7 +203,9 @@ def test_query_5(client, dataset_path, fs):
     )
     jn5["revenue"] = jn5.l_extendedprice * (1.0 - jn5.l_discount)
     gb = jn5.groupby("n_name")["revenue"].sum()
-    gb.reset_index().sort_values("revenue", ascending=False).compute()
+    result = gb.reset_index().sort_values("revenue", ascending=False).compute()
+
+    return result
 
 
 def test_query_6(client, dataset_path, fs):
@@ -209,7 +224,9 @@ def test_query_6(client, dataset_path, fs):
     )
 
     flineitem = line_item_ds[sel]
-    (flineitem.l_extendedprice * flineitem.l_discount).sum().compute()
+    result = (flineitem.l_extendedprice * flineitem.l_discount).sum().compute()
+
+    return pd.DataFrame({"revenue": [result]})
 
 
 @pytest.mark.shuffle_p2p
@@ -226,7 +243,7 @@ def test_query_7(client, dataset_path, fs):
     lineitem_filtered = line_item_ds[
         (line_item_ds["l_shipdate"] >= var1) & (line_item_ds["l_shipdate"] < var2)
     ]
-    lineitem_filtered["l_year"] = 1  # lineitem_filtered["l_shipdate"].dt.year
+    lineitem_filtered["l_year"] = lineitem_filtered["l_shipdate"].dt.year
     lineitem_filtered["revenue"] = lineitem_filtered["l_extendedprice"] * (
         1.0 - lineitem_filtered["l_discount"]
     )
@@ -302,11 +319,13 @@ def test_query_7(client, dataset_path, fs):
     )
     result_df.columns = ["supp_nation", "cust_nation", "l_year", "revenue"]
 
-    result_df.sort_values(
+    result = result_df.sort_values(
         by=["supp_nation", "cust_nation", "l_year"],
         ascending=True,
     ).compute()
 
+    return result
+
 
 @pytest.mark.shuffle_p2p
 def test_query_9(client, dataset_path, fs):
@@ -373,7 +392,7 @@ def test_query_9(client, dataset_path, fs):
     )
     subquery = subquery[["o_year", "nation", "amount"]]
 
-    _ = (
+    result = (
         subquery.groupby(["nation", "o_year"])
         .amount.sum()
         .round(2)
@@ -381,6 +400,8 @@ def test_query_9(client, dataset_path, fs):
         .sort_values(by=["nation", "o_year"], ascending=[True, False])
     ).compute()
 
+    return result
+
 
 @pytest.mark.shuffle_p2p
 def test_query_10(client, dataset_path, fs):
@@ -437,7 +458,7 @@ def test_query_10(client, dataset_path, fs):
         & (query.l_returnflag == "R")
     ]
     query["revenue"] = query.l_extendedprice * (1 - query.l_discount)
-    _ = (
+    result = (
         query.groupby(
             [
                 "c_custkey",
@@ -451,11 +472,14 @@ def test_query_10(client, dataset_path, fs):
         )
         .revenue.sum()
         .round(2)
-        .to_frame()
+        # .to_frame()
+        .reset_index()
         .sort_values(by=["revenue"], ascending=[False])
         .head(20)
     )
 
+    return result
+
 
 @pytest.mark.shuffle_p2p
 def test_query_12(client, dataset_path, fs):
@@ -512,14 +536,16 @@ def test_query_12(client, dataset_path, fs):
     table["low_line_count"] = 0
     table["low_line_count"] = table.low_line_count.where(~mask, 1)
 
-    _ = (
+    result = (
         table.groupby("l_shipmode")
-        .agg({"low_line_count": "sum", "high_line_count": "sum"})
+        .agg({"high_line_count": "sum", "low_line_count": "sum"})
         .reset_index()
         .sort_values(by="l_shipmode")
         .compute()
     )
 
+    return result
+
 
 @pytest.mark.shuffle_p2p
 def test_query_13(client, dataset_path, fs):
@@ -558,7 +584,7 @@ def test_query_13(client, dataset_path, fs):
         .rename(columns={0: "c_count"})[["c_custkey", "c_count"]]
     )
 
-    _ = (
+    result = (
         subquery.groupby("c_count")
         .size()
         .to_frame()
@@ -568,6 +594,8 @@ def test_query_13(client, dataset_path, fs):
         .compute()
     )
 
+    return result
+
 
 @pytest.mark.shuffle_p2p
 def test_query_14(client, dataset_path, fs):
@@ -605,7 +633,7 @@ def test_query_14(client, dataset_path, fs):
     )
 
     # aggregate promo revenue calculation
-    _ = (
+    result = (
         (
             100.00
             * table.promo_revenue.sum()
@@ -615,3 +643,5 @@ def test_query_14(client, dataset_path, fs):
         .round(2)
         .compute()
     )
+
+    return pd.DataFrame({"promo_revenue": [result]})