chore: update tpch q21 (#1225)

Author: Genesis929

third_party/bigframes_vendored/tpch/queries/q21.py

@@ -1,6 +1,4 @@
-# Contains code from https://github.com/pola-rs/tpch/blob/main/queries/duckdb/q21.py
-
-import typing
+# Contains code from https://github.com/pola-rs/tpch/blob/main/queries/polars/q21.py
 
 import bigframes
 import bigframes.pandas as bpd
@@ -24,39 +22,38 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
         index_col=bigframes.enums.DefaultIndexKind.NULL,
     )
 
-    nation = nation[nation["N_NAME"] == "SAUDI ARABIA"]
-    orders = orders[orders["O_ORDERSTATUS"] == "F"]
-
-    l1 = lineitem[lineitem["L_RECEIPTDATE"] > lineitem["L_COMMITDATE"]][
-        ["L_ORDERKEY", "L_SUPPKEY"]
-    ]
+    var1 = "SAUDI ARABIA"
 
-    l2 = lineitem.groupby("L_ORDERKEY", as_index=False).agg(
-        NUNIQUE_COL=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="nunique")
+    q1 = lineitem.groupby("L_ORDERKEY", as_index=False).agg(
+        N_SUPP_BY_ORDER=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="size")
     )
-    l2 = l2[l2["NUNIQUE_COL"] > 1][["L_ORDERKEY"]]
+    q1 = q1[q1["N_SUPP_BY_ORDER"] > 1]
 
-    l3 = l1.groupby("L_ORDERKEY", as_index=False).agg(
-        NUNIQUE_COL=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="nunique")
-    )
-    l3 = l3[l3["NUNIQUE_COL"] == 1][["L_ORDERKEY"]]
+    lineitem_filtered = lineitem[lineitem["L_RECEIPTDATE"] > lineitem["L_COMMITDATE"]]
 
-    l1 = l1.merge(l2, on="L_ORDERKEY", how="inner").merge(
-        l3, on="L_ORDERKEY", how="inner"
+    q1 = q1.merge(lineitem_filtered, on="L_ORDERKEY")
+
+    q_final = q1.groupby("L_ORDERKEY", as_index=False).agg(
+        N_SUPP_BY_ORDER_FINAL=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="size")
     )
 
-    merged = supplier.merge(nation, left_on="S_NATIONKEY", right_on="N_NATIONKEY")
-    merged = merged.merge(l1, left_on="S_SUPPKEY", right_on="L_SUPPKEY")
-    merged = merged.merge(orders, left_on="L_ORDERKEY", right_on="O_ORDERKEY")
+    q_final = q_final.merge(q1, on="L_ORDERKEY")
+    q_final = q_final.merge(supplier, left_on="L_SUPPKEY", right_on="S_SUPPKEY")
+    q_final = q_final.merge(nation, left_on="S_NATIONKEY", right_on="N_NATIONKEY")
+    q_final = q_final.merge(orders, left_on="L_ORDERKEY", right_on="O_ORDERKEY")
+
+    q_final = q_final[
+        (q_final["N_SUPP_BY_ORDER_FINAL"] == 1)
+        & (q_final["N_NAME"] == var1)
+        & (q_final["O_ORDERSTATUS"] == "F")
+    ]
 
-    result = merged.groupby("S_NAME", as_index=False).agg(
+    q_final = q_final.groupby("S_NAME", as_index=False).agg(
         NUMWAIT=bpd.NamedAgg(column="L_SUPPKEY", aggfunc="size")
     )
 
-    result = (
-        typing.cast(bpd.DataFrame, result)
-        .sort_values(["NUMWAIT", "S_NAME"], ascending=[False, True])
-        .head(100)
-    )
+    q_final = q_final.sort_values(
+        by=["NUMWAIT", "S_NAME"], ascending=[False, True]
+    ).head(100)
 
-    result.to_gbq()
+    q_final.to_gbq()