Merge pull request #3 from marenwestermann/benchmark-q10

feat: query 10 implementation

Author: MarcoGorelli

execute/q10.py

@@ -0,0 +1,64 @@
+from queries import q10
+import os
+
+import pandas as pd
+import polars as pl
+
+pd.options.mode.copy_on_write = True
+pd.options.future.infer_string = True
+
+customer = os.path.join("data", "customer.parquet")
+nation = os.path.join("data", "nation.parquet")
+lineitem = os.path.join("data", "lineitem.parquet")
+orders = os.path.join("data", "orders.parquet")
+
+IO_FUNCS = {
+    'pandas': lambda x: pd.read_parquet(x, engine='pyarrow'),
+    'pandas[pyarrow]': lambda x: pd.read_parquet(x, engine='pyarrow', dtype_backend='pyarrow'),
+    'polars[eager]': lambda x: pl.read_parquet(x),
+    'polars[lazy]': lambda x: pl.scan_parquet(x),
+}
+
+tool = 'pandas'
+fn = IO_FUNCS[tool]
+print(
+    q10.query(
+        fn(customer),
+        fn(nation),
+        fn(lineitem),
+        fn(orders)
+    )
+)
+
+tool = 'pandas[pyarrow]'
+fn = IO_FUNCS[tool]
+print(
+    q10.query(
+        fn(customer),
+        fn(nation),
+        fn(lineitem),
+        fn(orders)
+    )
+)
+
+tool = 'polars[eager]'
+fn = IO_FUNCS[tool]
+print(
+    q10.query(
+        fn(customer),
+        fn(nation),
+        fn(lineitem),
+        fn(orders)
+    )
+)
+
+tool = 'polars[lazy]'
+fn = IO_FUNCS[tool]
+print(
+    q10.query(
+        fn(customer),
+        fn(nation),
+        fn(lineitem),
+        fn(orders)
+    ).collect()
+)

queries/q10.py

@@ -0,0 +1,54 @@
+from typing import Any
+from datetime import datetime
+import narwhals as nw
+
+def query(
+    customer_ds_raw: Any,
+    nation_ds_raw: Any,
+    lineitem_ds_raw: Any,
+    orders_ds_raw: Any,
+) -> Any:
+
+    nation_ds = nw.from_native(nation_ds_raw)
+    line_item_ds = nw.from_native(lineitem_ds_raw)
+    orders_ds = nw.from_native(orders_ds_raw)
+    customer_ds = nw.from_native(customer_ds_raw)
+    
+    var1 = datetime(1993, 10, 1)
+    var2 = datetime(1994, 1, 1)
+
+    result = (
+        customer_ds.join(orders_ds, left_on="c_custkey", right_on="o_custkey")
+        .join(line_item_ds, left_on="o_orderkey", right_on="l_orderkey")
+        .join(nation_ds, left_on="c_nationkey", right_on="n_nationkey")
+        .filter(nw.col("o_orderdate").is_between(var1, var2, closed="left"))
+        .filter(nw.col("l_returnflag") == "R")
+        .with_columns(
+            (nw.col("l_extendedprice") * (1 - nw.col("l_discount")))
+            .alias("revenue")
+        )
+        .group_by(
+            "c_custkey",
+            "c_name",
+            "c_acctbal",
+            "c_phone",
+            "n_name",
+            "c_address",
+            "c_comment",
+        )
+        .agg(nw.sum("revenue"))
+        .select(
+            "c_custkey",
+            "c_name",
+            "revenue",
+            "c_acctbal",
+            "n_name",
+            "c_address",
+            "c_phone",
+            "c_comment",
+        )
+        .sort(by="revenue", descending=True)
+        .head(20)
+    )
+
+    return nw.to_native(result)