update tpch examples for latest function uses

Author: timsaucer

examples/tpch/q02_minimum_cost_supplier.py

@@ -32,6 +32,7 @@
 import datafusion
 from datafusion import SessionContext, col, lit
 from datafusion import functions as F
+from datafusion.expr import Window
 from util import get_data_path
 
 # This is the part we're looking for. Values selected here differ from the spec in order to run
@@ -106,11 +107,8 @@
 window_frame = datafusion.WindowFrame("rows", None, None)
 df = df.with_column(
     "min_cost",
-    F.window(
-        "min",
-        [col("ps_supplycost")],
-        partition_by=[col("ps_partkey")],
-        window_frame=window_frame,
+    F.min(col("ps_supplycost")).over(
+        Window(partition_by=[col("ps_partkey")], window_frame=window_frame)
     ),
 )
 

examples/tpch/q11_important_stock_identification.py

@@ -29,6 +29,7 @@
 
 from datafusion import SessionContext, WindowFrame, col, lit
 from datafusion import functions as F
+from datafusion.expr import Window
 from util import get_data_path
 
 NATION = "GERMANY"
@@ -71,7 +72,7 @@
 window_frame = WindowFrame("rows", None, None)
 
 df = df.with_column(
-    "total_value", F.window("sum", [col("value")], window_frame=window_frame)
+    "total_value", F.sum(col("value")).over(Window(window_frame=window_frame))
 )
 
 # Limit to the parts for which there is a significant value based on the fraction of the total

examples/tpch/q15_top_supplier.py

@@ -31,6 +31,7 @@
 import pyarrow as pa
 from datafusion import SessionContext, WindowFrame, col, lit
 from datafusion import functions as F
+from datafusion.expr import Window
 from util import get_data_path
 
 DATE = "1996-01-01"
@@ -70,7 +71,8 @@
 # Use a window function to find the maximum revenue across the entire dataframe
 window_frame = WindowFrame("rows", None, None)
 df = df.with_column(
-    "max_revenue", F.window("max", [col("total_revenue")], window_frame=window_frame)
+    "max_revenue",
+    F.max(col("total_revenue")).over(Window(window_frame=window_frame)),
 )
 
 # Find all suppliers whose total revenue is the same as the maximum

examples/tpch/q17_small_quantity_order.py

@@ -30,6 +30,7 @@
 
 from datafusion import SessionContext, WindowFrame, col, lit
 from datafusion import functions as F
+from datafusion.expr import Window
 from util import get_data_path
 
 BRAND = "Brand#23"
@@ -58,11 +59,8 @@
 window_frame = WindowFrame("rows", None, None)
 df = df.with_column(
     "avg_quantity",
-    F.window(
-        "avg",
-        [col("l_quantity")],
-        window_frame=window_frame,
-        partition_by=[col("l_partkey")],
+    F.avg(col("l_quantity")).over(
+        Window(partition_by=[col("l_partkey")], window_frame=window_frame)
     ),
 )
 

examples/tpch/q22_global_sales_opportunity.py

@@ -28,6 +28,7 @@
 
 from datafusion import SessionContext, WindowFrame, col, lit
 from datafusion import functions as F
+from datafusion.expr import Window
 from util import get_data_path
 
 NATION_CODES = [13, 31, 23, 29, 30, 18, 17]
@@ -55,7 +56,8 @@
 # current row. We want our frame to cover the entire data frame.
 window_frame = WindowFrame("rows", None, None)
 df = df.with_column(
-    "avg_balance", F.window("avg", [col("c_acctbal")], window_frame=window_frame)
+    "avg_balance",
+    F.avg(col("c_acctbal")).over(Window(window_frame=window_frame)),
 )
 
 df.show()