chore: update tpch code (#1389)

* chore: update tpch code

* update q5

* add 14, 17, 20

* add q22

* add q21

* reverse q7 q10 q21

Author: Genesis929

third_party/bigframes_vendored/tpch/queries/q14.py

@@ -18,7 +18,7 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
     var1 = date(1995, 9, 1)
     var2 = date(1995, 10, 1)
 
-    merged = lineitem.merge(part, left_on="L_PARTKEY", right_on="P_PARTKEY")
+    merged = part.merge(lineitem, left_on="P_PARTKEY", right_on="L_PARTKEY")
 
     filtered = merged[(merged["L_SHIPDATE"] >= var1) & (merged["L_SHIPDATE"] < var2)]
 

third_party/bigframes_vendored/tpch/queries/q17.py

@@ -19,7 +19,7 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
 
     filtered_part = part[(part["P_BRAND"] == VAR1) & (part["P_CONTAINER"] == VAR2)]
     q1 = bpd.merge(
-        filtered_part, lineitem, how="left", left_on="P_PARTKEY", right_on="L_PARTKEY"
+        lineitem, filtered_part, how="right", left_on="L_PARTKEY", right_on="P_PARTKEY"
     )
 
     grouped = (

third_party/bigframes_vendored/tpch/queries/q20.py

@@ -49,8 +49,10 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
         partsupp, left_on="P_PARTKEY", right_on="PS_PARTKEY"
     )
 
-    final_join = joined_parts.merge(
-        q1, left_on=["PS_SUPPKEY", "P_PARTKEY"], right_on=["L_SUPPKEY", "L_PARTKEY"]
+    final_join = q1.merge(
+        joined_parts,
+        left_on=["L_SUPPKEY", "L_PARTKEY"],
+        right_on=["PS_SUPPKEY", "P_PARTKEY"],
     )
     final_filtered = final_join[final_join["PS_AVAILQTY"] > final_join["SUM_QUANTITY"]]
 

third_party/bigframes_vendored/tpch/queries/q22.py

@@ -15,37 +15,24 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
     )
 
     country_codes = ["13", "31", "23", "29", "30", "18", "17"]
-
     customer["CNTRYCODE"] = customer["C_PHONE"].str.slice(0, 2)
+    customer = customer[customer["CNTRYCODE"].isin(country_codes)]
 
     avg_acctbal = (
-        customer[
-            (customer["CNTRYCODE"].isin(country_codes)) & (customer["C_ACCTBAL"] > 0)
-        ][["C_ACCTBAL"]]
+        customer[customer["C_ACCTBAL"] > 0.0][["C_ACCTBAL"]]
         .mean()
         .rename("AVG_ACCTBAL")
     )
-
-    orders_unique = orders["O_CUSTKEY"].unique(keep_order=False).to_frame()
-
-    matched_customers = customer.merge(
-        orders_unique, left_on="C_CUSTKEY", right_on="O_CUSTKEY"
-    )
-    matched_customers["IS_IN_ORDERS"] = True
-
-    customer = customer.merge(
-        matched_customers[["C_CUSTKEY", "IS_IN_ORDERS"]], on="C_CUSTKEY", how="left"
-    )
-    customer["IS_IN_ORDERS"] = customer["IS_IN_ORDERS"].fillna(False)
     customer = customer.merge(avg_acctbal, how="cross")
 
-    filtered_customers = customer[
-        (customer["CNTRYCODE"].isin(country_codes))
-        & (customer["C_ACCTBAL"] > customer["AVG_ACCTBAL"])
-        & (~customer["IS_IN_ORDERS"])
-    ]
+    filtered_customer = customer[customer["C_ACCTBAL"] > customer["AVG_ACCTBAL"]]
 
-    result = filtered_customers.groupby("CNTRYCODE", as_index=False).agg(
+    orders_unique = orders["O_CUSTKEY"].unique(keep_order=False).to_frame()
+    filtered_customer = filtered_customer.merge(
+        orders_unique, left_on="C_CUSTKEY", right_on="O_CUSTKEY", how="left"
+    )
+    filtered_customer = filtered_customer[filtered_customer["O_CUSTKEY"].isnull()]
+    result = filtered_customer.groupby("CNTRYCODE", as_index=False).agg(
         NUMCUST=bpd.NamedAgg(column="C_CUSTKEY", aggfunc="count"),
         TOTACCTBAL=bpd.NamedAgg(column="C_ACCTBAL", aggfunc="sum"),
     )

third_party/bigframes_vendored/tpch/queries/q3.py

@@ -23,11 +23,13 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
 
     fcustomer = customer[customer["C_MKTSEGMENT"] == "BUILDING"]
 
-    jn1 = fcustomer.merge(orders, left_on="C_CUSTKEY", right_on="O_CUSTKEY")
-    jn2 = jn1.merge(lineitem, left_on="O_ORDERKEY", right_on="L_ORDERKEY")
+    filtered_orders = orders[orders["O_ORDERDATE"] < date_var]
+    filtered_lineitem = lineitem[lineitem["L_SHIPDATE"] > date_var]
 
-    jn2 = jn2[jn2["O_ORDERDATE"] < date_var]
-    jn2 = jn2[jn2["L_SHIPDATE"] > date_var]
+    jn1 = filtered_lineitem.merge(
+        filtered_orders, left_on="L_ORDERKEY", right_on="O_ORDERKEY"
+    )
+    jn2 = fcustomer.merge(jn1, left_on="C_CUSTKEY", right_on="O_CUSTKEY")
     jn2["REVENUE"] = jn2["L_EXTENDEDPRICE"] * (1 - jn2["L_DISCOUNT"])
 
     gb = jn2.groupby(["O_ORDERKEY", "O_ORDERDATE", "O_SHIPPRIORITY"], as_index=False)

third_party/bigframes_vendored/tpch/queries/q5.py

@@ -35,20 +35,20 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
     var2 = date(1994, 1, 1)
     var3 = date(1995, 1, 1)
 
+    region = region[region["R_NAME"] == var1]
+    orders = orders[(orders["O_ORDERDATE"] >= var2) & (orders["O_ORDERDATE"] < var3)]
+    lineitem["REVENUE"] = lineitem["L_EXTENDEDPRICE"] * (1.0 - lineitem["L_DISCOUNT"])
+
     jn1 = region.merge(nation, left_on="R_REGIONKEY", right_on="N_REGIONKEY")
     jn2 = jn1.merge(customer, left_on="N_NATIONKEY", right_on="C_NATIONKEY")
-    jn3 = jn2.merge(orders, left_on="C_CUSTKEY", right_on="O_CUSTKEY")
-    jn4 = jn3.merge(lineitem, left_on="O_ORDERKEY", right_on="L_ORDERKEY")
+    jn3 = orders.merge(jn2, left_on="O_CUSTKEY", right_on="C_CUSTKEY")
+    jn4 = lineitem.merge(jn3, left_on="L_ORDERKEY", right_on="O_ORDERKEY")
     jn5 = jn4.merge(
         supplier,
         left_on=["L_SUPPKEY", "N_NATIONKEY"],
         right_on=["S_SUPPKEY", "S_NATIONKEY"],
     )
 
-    jn5 = jn5[jn5["R_NAME"] == var1]
-    jn5 = jn5[(jn5["O_ORDERDATE"] >= var2) & (jn5["O_ORDERDATE"] < var3)]
-    jn5["REVENUE"] = jn5["L_EXTENDEDPRICE"] * (1.0 - jn5["L_DISCOUNT"])
-
     gb = jn5.groupby("N_NAME", as_index=False)["REVENUE"].sum()
     result_df = gb.sort_values("REVENUE", ascending=False)