chore: Use isin() in tpch benchmarks (#1400)

Author: TrevorBergeron

third_party/bigframes_vendored/tpch/queries/q16.py

@@ -20,22 +20,16 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
 
     var1 = "Brand#45"
 
-    supplier = (
-        supplier[
-            ~supplier["S_COMMENT"].str.contains("Customer.*Complaints", regex=True)
-        ]["S_SUPPKEY"]
-        .unique(keep_order=False)
-        .to_frame()
-    )
+    supplier = supplier[
+        ~supplier["S_COMMENT"].str.contains("Customer.*Complaints", regex=True)
+    ]["S_SUPPKEY"]
 
     q_filtered = part.merge(partsupp, left_on="P_PARTKEY", right_on="PS_PARTKEY")
     q_filtered = q_filtered[q_filtered["P_BRAND"] != var1]
     q_filtered = q_filtered[~q_filtered["P_TYPE"].str.contains("MEDIUM POLISHED")]
     q_filtered = q_filtered[q_filtered["P_SIZE"].isin([49, 14, 23, 45, 19, 3, 36, 9])]
 
-    final_df = q_filtered.merge(
-        supplier, left_on=["PS_SUPPKEY"], right_on=["S_SUPPKEY"]
-    )
+    final_df = q_filtered[q_filtered["PS_SUPPKEY"].isin(supplier)]
 
     grouped = final_df.groupby(["P_BRAND", "P_TYPE", "P_SIZE"], as_index=False)
     result = grouped.agg(

third_party/bigframes_vendored/tpch/queries/q18.py

@@ -22,14 +22,13 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
 
     var1 = 300
 
+    # order with over 300 items
     q1 = lineitem.groupby("L_ORDERKEY", as_index=False).agg(
         SUM_QUANTITY=bpd.NamedAgg(column="L_QUANTITY", aggfunc="sum")
     )
     q1 = q1[q1["SUM_QUANTITY"] > var1]
 
-    filtered_orders = orders.merge(
-        q1, left_on="O_ORDERKEY", right_on="L_ORDERKEY", how="inner"
-    )
+    filtered_orders = orders[orders["O_ORDERKEY"].isin(q1["L_ORDERKEY"])]
 
     result = filtered_orders.merge(
         lineitem, left_on="O_ORDERKEY", right_on="L_ORDERKEY"

third_party/bigframes_vendored/tpch/queries/q20.py

@@ -44,21 +44,19 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
 
     filtered_parts = part[part["P_NAME"].str.startswith(var4)]
 
-    filtered_parts = filtered_parts["P_PARTKEY"].unique(keep_order=False).to_frame()
-    joined_parts = filtered_parts.merge(
-        partsupp, left_on="P_PARTKEY", right_on="PS_PARTKEY"
-    )
+    filtered_parts = filtered_parts["P_PARTKEY"]
+    joined_parts = partsupp[partsupp["PS_PARTKEY"].isin(filtered_parts)]
 
     final_join = q1.merge(
         joined_parts,
         left_on=["L_SUPPKEY", "L_PARTKEY"],
-        right_on=["PS_SUPPKEY", "P_PARTKEY"],
+        right_on=["PS_SUPPKEY", "PS_PARTKEY"],
     )
-    final_filtered = final_join[final_join["PS_AVAILQTY"] > final_join["SUM_QUANTITY"]]
-
-    final_filtered = final_filtered["PS_SUPPKEY"].unique(keep_order=False).to_frame()
+    final_filtered = final_join[final_join["PS_AVAILQTY"] > final_join["SUM_QUANTITY"]][
+        "PS_SUPPKEY"
+    ]
 
-    final_result = final_filtered.merge(q3, left_on="PS_SUPPKEY", right_on="S_SUPPKEY")
+    final_result = q3[q3["S_SUPPKEY"].isin(final_filtered)]
     final_result = final_result[["S_NAME", "S_ADDRESS"]].sort_values(by="S_NAME")
 
     next(final_result.to_pandas_batches(max_results=1500))

third_party/bigframes_vendored/tpch/queries/q22.py

@@ -27,11 +27,9 @@ def q(project_id: str, dataset_id: str, session: bigframes.Session):
 
     filtered_customer = customer[customer["C_ACCTBAL"] > customer["AVG_ACCTBAL"]]
 
-    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()]
+    filtered_customer = filtered_customer[
+        ~filtered_customer["C_CUSTKEY"].isin(orders["O_CUSTKEY"])
+    ]
     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"),