WIP: Add dask query 8 [skip ci]

IndexError: Column(s) nation already selected

Author: milesgranger

tests/tpch/test_dask.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 
 import dask_expr as dd
+from dask.dataframe import Aggregation
 
 
 def test_query_1(client, dataset_path, fs):
@@ -217,6 +218,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)
     ]
+    # TODO: This is wrong.
     lineitem_filtered["l_year"] = 1  # lineitem_filtered["l_shipdate"].dt.year
     lineitem_filtered["revenue"] = lineitem_filtered["l_extendedprice"] * (
         1.0 - lineitem_filtered["l_discount"]
@@ -297,3 +299,103 @@ def test_query_7(client, dataset_path, fs):
         by=["supp_nation", "cust_nation", "l_year"],
         ascending=True,
     ).compute()
+
+
+def test_query_8(client, dataset_path, fs):
+    var1 = datetime.strptime("1995-01-01", "%Y-%m-%d")
+    var2 = datetime.strptime("1997-01-01", "%Y-%m-%d")
+
+    supplier_ds = dd.read_parquet(dataset_path + "supplier", filesystem=fs)
+    lineitem_ds = dd.read_parquet(dataset_path + "lineitem", filesystem=fs)
+    orders_ds = dd.read_parquet(dataset_path + "orders", filesystem=fs)
+    customer_ds = dd.read_parquet(dataset_path + "customer", filesystem=fs)
+    nation_ds = dd.read_parquet(dataset_path + "nation", filesystem=fs)
+    region_ds = dd.read_parquet(dataset_path + "region", filesystem=fs)
+    part_ds = dd.read_parquet(dataset_path + "part", filesystem=fs)
+
+    part_filtered = part_ds[part_ds["p_type"] == "ECONOMY ANODIZED STEEL"][
+        ["p_partkey"]
+    ]
+
+    lineitem_filtered = lineitem_ds[["l_partkey", "l_suppkey", "l_orderkey"]]
+    lineitem_filtered["volume"] = lineitem_ds["l_extendedprice"] * (
+        1.0 - lineitem_ds["l_discount"]
+    )
+    total = part_filtered.merge(
+        lineitem_filtered,
+        left_on="p_partkey",
+        right_on="l_partkey",
+        how="inner",
+    )[["l_suppkey", "l_orderkey", "volume"]]
+
+    supplier_filtered = supplier_ds[["s_suppkey", "s_nationkey"]]
+    total = total.merge(
+        supplier_filtered,
+        left_on="l_suppkey",
+        right_on="s_suppkey",
+        how="inner",
+    )[["l_orderkey", "volume", "s_nationkey"]]
+
+    orders_filtered = orders_ds[
+        (orders_ds["o_orderdate"] >= var1) & (orders_ds["o_orderdate"] < var2)
+    ]
+
+    orders_filtered["o_year"] = orders_filtered["o_orderdate"].dt.year
+    orders_filtered = orders_filtered[["o_orderkey", "o_custkey", "o_year"]]
+    total = total.merge(
+        orders_filtered,
+        left_on="l_orderkey",
+        right_on="o_orderkey",
+        how="inner",
+    )[["volume", "s_nationkey", "o_custkey", "o_year"]]
+
+    customer_filtered = customer_ds[["c_custkey", "c_nationkey"]]
+    total = total.merge(
+        customer_filtered,
+        left_on="o_custkey",
+        right_on="c_custkey",
+        how="inner",
+    )[["volume", "s_nationkey", "o_year", "c_nationkey"]]
+
+    n1_filtered = nation_ds[["n_nationkey", "n_regionkey"]]
+    n2_filtered = nation_ds[["n_nationkey", "n_name"]].rename(
+        columns={"n_name": "nation"}
+    )
+    total = total.merge(
+        n1_filtered,
+        left_on="c_nationkey",
+        right_on="n_nationkey",
+        how="inner",
+    )[["volume", "s_nationkey", "o_year", "n_regionkey"]]
+
+    total = total.merge(
+        n2_filtered,
+        left_on="s_nationkey",
+        right_on="n_nationkey",
+        how="inner",
+    )[["volume", "o_year", "n_regionkey", "nation"]]
+
+    region_filtered = region_ds[region_ds["r_name"] == "AMERICA"][["r_regionkey"]]
+    total = total.merge(
+        region_filtered,
+        left_on="n_regionkey",
+        right_on="r_regionkey",
+        how="inner",
+    )[["volume", "o_year", "nation"]]
+
+    def chunk_udf(df):
+        denominator = df["volume"]
+        df = df[df["nation"] == "BRAZIL"]
+        numerator = df["volume"]
+        return (numerator, denominator)
+
+    def agg_udf(x):
+        return round(x[0].sum() / x[1].sum(), 2)
+
+    agg = Aggregation(name="mkt_share", chunk=chunk_udf, agg=agg_udf)
+
+    total["mkt_share"] = total.groupby(["o_year"]).agg(agg)
+    total = total.rename(columns={"o_year": "o_year", "x": "mkt_share"})
+    total = total.sort_values(by=["o_year"], ascending=[True])
+
+    total.compute()