Push down filters in pandas queries

queries/pandas/q10.py

@@ -28,13 +28,15 @@ def query() -> pd.DataFrame:
         var1 = pd.Timestamp("1993-10-01")
         var2 = pd.Timestamp("1994-01-01")
 
+        orders_ds = orders_ds[
+            (orders_ds["o_orderdate"] >= var1) & (orders_ds["o_orderdate"] < var2)
+        ]
+        lineitem_ds = lineitem_ds[lineitem_ds["l_returnflag"] == "R"]
+
         jn1 = customer_ds.merge(orders_ds, left_on="c_custkey", right_on="o_custkey")
         jn2 = jn1.merge(lineitem_ds, left_on="o_orderkey", right_on="l_orderkey")
         jn3 = jn2.merge(nation_ds, left_on="c_nationkey", right_on="n_nationkey")
 
-        jn3 = jn3[(jn3["o_orderdate"] >= var1) & (jn3["o_orderdate"] < var2)]
-        jn3 = jn3[jn3["l_returnflag"] == "R"]
-
         jn3["revenue"] = jn3["l_extendedprice"] * (1 - jn3["l_discount"])
 
         gb = jn3.groupby(

queries/pandas/q12.py

@@ -24,12 +24,19 @@ def query() -> pd.DataFrame:
         var3 = pd.Timestamp("1994-01-01")
         var4 = pd.Timestamp("1995-01-01")
 
-        jn = orders_ds.merge(lineitem_ds, left_on="o_orderkey", right_on="l_orderkey")
+        lineitem_ds = lineitem_ds[lineitem_ds["l_shipmode"].isin([var1, var2])]
+        lineitem_ds = lineitem_ds[
+            lineitem_ds["l_commitdate"] < lineitem_ds["l_receiptdate"]
+        ]
+        lineitem_ds = lineitem_ds[
+            lineitem_ds["l_shipdate"] < lineitem_ds["l_commitdate"]
+        ]
+        lineitem_ds = lineitem_ds[
+            (lineitem_ds["l_receiptdate"] >= var3)
+            & (lineitem_ds["l_receiptdate"] < var4)
+        ]
 
-        jn = jn[jn["l_shipmode"].isin([var1, var2])]
-        jn = jn[jn["l_commitdate"] < jn["l_receiptdate"]]
-        jn = jn[jn["l_shipdate"] < jn["l_commitdate"]]
-        jn = jn[(jn["l_receiptdate"] >= var3) & (jn["l_receiptdate"] < var4)]
+        jn = orders_ds.merge(lineitem_ds, left_on="o_orderkey", right_on="l_orderkey")
 
         jn["high_line_count"] = jn["o_orderpriority"].isin(["1-URGENT", "2-HIGH"])
         jn["low_line_count"] = ~jn["o_orderpriority"].isin(["1-URGENT", "2-HIGH"])

queries/pandas/q14.py

@@ -22,9 +22,11 @@ def query() -> pd.DataFrame:
         var1 = pd.Timestamp("1995-09-01")
         var2 = pd.Timestamp("1995-10-01")
 
-        jn = lineitem_ds.merge(part_ds, left_on="l_partkey", right_on="p_partkey")
+        lineitem_ds = lineitem_ds[
+            (lineitem_ds["l_shipdate"] >= var1) & (lineitem_ds["l_shipdate"] < var2)
+        ]
 
-        jn = jn[(jn["l_shipdate"] >= var1) & (jn["l_shipdate"] < var2)]
+        jn = lineitem_ds.merge(part_ds, left_on="l_partkey", right_on="p_partkey")
 
         jn["revenue"] = jn["l_extendedprice"] * (1 - jn["l_discount"])
         jn["promo_revenue"] = jn["revenue"].where(

queries/pandas/q19.py

@@ -19,10 +19,10 @@ def query() -> pd.DataFrame:
         lineitem_ds = lineitem_ds_fn()
         part_ds = part_ds_fn()
 
-        jn = part_ds.merge(lineitem_ds, left_on="p_partkey", right_on="l_partkey")
+        lineitem_ds = lineitem_ds[lineitem_ds["l_shipmode"].isin(["AIR", "AIR REG"])]
+        lineitem_ds = lineitem_ds[lineitem_ds["l_shipinstruct"] == "DELIVER IN PERSON"]
 
-        jn = jn[jn["l_shipmode"].isin(["AIR", "AIR REG"])]
-        jn = jn[jn["l_shipinstruct"] == "DELIVER IN PERSON"]
+        jn = part_ds.merge(lineitem_ds, left_on="p_partkey", right_on="l_partkey")
 
         # Complex filter conditions
         cond1 = (

queries/pandas/q4.py

@@ -24,10 +24,16 @@ def query() -> pd.DataFrame:
         var1 = date(1993, 7, 1)
         var2 = date(1993, 10, 1)
 
-        jn = line_item_ds.merge(orders_ds, left_on="l_orderkey", right_on="o_orderkey")
-
-        jn = jn[(jn["o_orderdate"] >= var1) & (jn["o_orderdate"] < var2)]
-        jn = jn[jn["l_commitdate"] < jn["l_receiptdate"]]
+        orders_filtered = orders_ds[
+            (orders_ds["o_orderdate"] >= var1) & (orders_ds["o_orderdate"] < var2)
+        ]
+        lineitem_filtered = line_item_ds[
+            line_item_ds["l_commitdate"] < line_item_ds["l_receiptdate"]
+        ]
+
+        jn = lineitem_filtered.merge(
+            orders_filtered, left_on="l_orderkey", right_on="o_orderkey"
+        )
 
         jn = jn.drop_duplicates(subset=["o_orderpriority", "l_orderkey"])
 

queries/pandas/q5.py

@@ -39,6 +39,11 @@ def query() -> pd.DataFrame:
         var2 = date(1994, 1, 1)
         var3 = date(1995, 1, 1)
 
+        region_ds = region_ds[region_ds["r_name"] == var1]
+        orders_ds = orders_ds[
+            (orders_ds["o_orderdate"] >= var2) & (orders_ds["o_orderdate"] < var3)
+        ]
+
         jn1 = region_ds.merge(nation_ds, left_on="r_regionkey", right_on="n_regionkey")
         jn2 = jn1.merge(customer_ds, left_on="n_nationkey", right_on="c_nationkey")
         jn3 = jn2.merge(orders_ds, left_on="c_custkey", right_on="o_custkey")
@@ -49,8 +54,6 @@ def query() -> pd.DataFrame:
             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()

queries/pandas/q8.py

@@ -44,22 +44,35 @@ def query() -> pd.DataFrame:
         var4 = date(1995, 1, 1)
         var5 = date(1996, 12, 31)
 
-        n1 = nation_ds.loc[:, ["n_nationkey", "n_regionkey"]]
         n2 = nation_ds.loc[:, ["n_nationkey", "n_name"]]
 
-        jn1 = part_ds.merge(line_item_ds, left_on="p_partkey", right_on="l_partkey")
-        jn2 = jn1.merge(supplier_ds, left_on="l_suppkey", right_on="s_suppkey")
-        jn3 = jn2.merge(orders_ds, left_on="l_orderkey", right_on="o_orderkey")
-        jn4 = jn3.merge(customer_ds, left_on="o_custkey", right_on="c_custkey")
-        jn5 = jn4.merge(n1, left_on="c_nationkey", right_on="n_nationkey")
-        jn6 = jn5.merge(region_ds, left_on="n_regionkey", right_on="r_regionkey")
-
-        jn6 = jn6[(jn6["r_name"] == var2)]
-
-        jn7 = jn6.merge(n2, left_on="s_nationkey", right_on="n_nationkey")
-
-        jn7 = jn7[(jn7["o_orderdate"] >= var4) & (jn7["o_orderdate"] <= var5)]
-        jn7 = jn7[jn7["p_type"] == var3]
+        part_ds = part_ds[part_ds["p_type"] == var3][["p_partkey", "p_type"]]
+        orders_ds = orders_ds[
+            (orders_ds["o_orderdate"] >= var4) & (orders_ds["o_orderdate"] <= var5)
+        ][["o_orderkey", "o_orderdate", "o_custkey"]]
+        region_filtered = region_ds[region_ds["r_name"] == var2][["r_regionkey"]]
+        america_nations = nation_ds.merge(
+            region_filtered, left_on="n_regionkey", right_on="r_regionkey"
+        )[["n_nationkey"]]
+        america_customers = customer_ds.merge(
+            america_nations, left_on="c_nationkey", right_on="n_nationkey"
+        )[["c_custkey", "c_nationkey"]]
+        orders_america = orders_ds.merge(
+            america_customers, left_on="o_custkey", right_on="c_custkey"
+        )
+
+        lineitem_parts = part_ds.merge(
+            line_item_ds, left_on="p_partkey", right_on="l_partkey"
+        )[["l_orderkey", "l_suppkey", "l_extendedprice", "l_discount", "p_type"]]
+        jn = lineitem_parts.merge(
+            orders_america, left_on="l_orderkey", right_on="o_orderkey"
+        )
+        jn_sup = jn.merge(
+            supplier_ds[["s_suppkey", "s_nationkey"]],
+            left_on="l_suppkey",
+            right_on="s_suppkey",
+        )
+        jn7 = jn_sup.merge(n2, left_on="s_nationkey", right_on="n_nationkey")
 
         jn7["o_year"] = jn7["o_orderdate"].dt.year
         jn7["volume"] = jn7["l_extendedprice"] * (1.0 - jn7["l_discount"])

queries/pandas/q9.py

@@ -31,6 +31,8 @@ def query() -> pd.DataFrame:
         part_supp_ds = part_supp_ds_fn()
         supplier_ds = supplier_ds_fn()
 
+        part_ds = part_ds[part_ds["p_name"].str.contains("green", regex=False)]
+
         jn1 = part_ds.merge(part_supp_ds, left_on="p_partkey", right_on="ps_partkey")
         jn2 = jn1.merge(supplier_ds, left_on="ps_suppkey", right_on="s_suppkey")
         jn3 = jn2.merge(
@@ -41,8 +43,6 @@ def query() -> pd.DataFrame:
         jn4 = jn3.merge(orders_ds, left_on="l_orderkey", right_on="o_orderkey")
         jn5 = jn4.merge(nation_ds, left_on="s_nationkey", right_on="n_nationkey")
 
-        jn5 = jn5[jn5["p_name"].str.contains("green", regex=False)]
-
         jn5["o_year"] = jn5["o_orderdate"].dt.year
         jn5["amount"] = jn5["l_extendedprice"] * (1.0 - jn5["l_discount"]) - (
             jn5["ps_supplycost"] * jn5["l_quantity"]