Tsaucer/find window fn (#747)

* Add a search order when attempting to locate the appropriate window function

* Remove unnecessary markings

* Linting

* Code cleanup

Author: timsaucer

examples/tpch/_tests.py

@@ -21,6 +21,7 @@
 from datafusion import col, lit, functions as F
 from util import get_answer_file
 
+
 def df_selection(col_name, col_type):
     if col_type == pa.float64() or isinstance(col_type, pa.Decimal128Type):
         return F.round(col(col_name), lit(2)).alias(col_name)
@@ -29,14 +30,16 @@ def df_selection(col_name, col_type):
     else:
         return col(col_name)
 
+
 def load_schema(col_name, col_type):
     if col_type == pa.int64() or col_type == pa.int32():
         return col_name, pa.string()
     elif isinstance(col_type, pa.Decimal128Type):
         return col_name, pa.float64()
     else:
         return col_name, col_type
-    
+
+
 def expected_selection(col_name, col_type):
     if col_type == pa.int64() or col_type == pa.int32():
         return F.trim(col(col_name)).cast(col_type).alias(col_name)
@@ -45,20 +48,23 @@ def expected_selection(col_name, col_type):
     else:
         return col(col_name)
 
+
 def selections_and_schema(original_schema):
-    columns = [ (c, original_schema.field(c).type) for c in original_schema.names ]
+    columns = [(c, original_schema.field(c).type) for c in original_schema.names]
 
-    df_selections = [ df_selection(c, t) for (c, t) in columns]
-    expected_schema = [ load_schema(c, t) for (c, t) in columns]
-    expected_selections = [ expected_selection(c, t) for (c, t) in columns]
+    df_selections = [df_selection(c, t) for (c, t) in columns]
+    expected_schema = [load_schema(c, t) for (c, t) in columns]
+    expected_selections = [expected_selection(c, t) for (c, t) in columns]
 
     return (df_selections, expected_schema, expected_selections)
 
+
 def check_q17(df):
     raw_value = float(df.collect()[0]["avg_yearly"][0].as_py())
     value = round(raw_value, 2)
     assert abs(value - 348406.05) < 0.001
 
+
 @pytest.mark.parametrize(
     ("query_code", "answer_file"),
     [
@@ -72,9 +78,7 @@ def check_q17(df):
         ("q08_market_share", "q8"),
         ("q09_product_type_profit_measure", "q9"),
         ("q10_returned_item_reporting", "q10"),
-        pytest.param(
-            "q11_important_stock_identification", "q11", 
-        ),
+        ("q11_important_stock_identification", "q11"),
         ("q12_ship_mode_order_priority", "q12"),
         ("q13_customer_distribution", "q13"),
         ("q14_promotion_effect", "q14"),
@@ -97,13 +101,20 @@ def test_tpch_query_vs_answer_file(query_code: str, answer_file: str):
     if answer_file == "q17":
         return check_q17(df)
 
-    (df_selections, expected_schema, expected_selections) = selections_and_schema(df.schema())
+    (df_selections, expected_schema, expected_selections) = selections_and_schema(
+        df.schema()
+    )
 
     df = df.select(*df_selections)
 
     read_schema = pa.schema(expected_schema)
 
-    df_expected = module.ctx.read_csv(get_answer_file(answer_file), schema=read_schema, delimiter="|", file_extension=".out")
+    df_expected = module.ctx.read_csv(
+        get_answer_file(answer_file),
+        schema=read_schema,
+        delimiter="|",
+        file_extension=".out",
+    )
 
     df_expected = df_expected.select(*expected_selections)
 

examples/tpch/convert_data_to_parquet.py

@@ -117,15 +117,14 @@
 
 curr_dir = os.path.dirname(os.path.abspath(__file__))
 for filename, curr_schema in all_schemas.items():
-
     # For convenience, go ahead and convert the schema column names to lowercase
     curr_schema = [(s[0].lower(), s[1]) for s in curr_schema]
 
     # Pre-collect the output columns so we can ignore the null field we add
     # in to handle the trailing | in the file
     output_cols = [r[0] for r in curr_schema]
 
-    curr_schema = [ pyarrow.field(r[0], r[1], nullable=False) for r in curr_schema]
+    curr_schema = [pyarrow.field(r[0], r[1], nullable=False) for r in curr_schema]
 
     # Trailing | requires extra field for in processing
     curr_schema.append(("some_null", pyarrow.null()))

examples/tpch/q08_market_share.py

@@ -47,7 +47,9 @@
 
 ctx = SessionContext()
 
-df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns("p_partkey", "p_type")
+df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns(
+    "p_partkey", "p_type"
+)
 df_supplier = ctx.read_parquet(get_data_path("supplier.parquet")).select_columns(
     "s_suppkey", "s_nationkey"
 )

examples/tpch/q09_product_type_profit_measure.py

@@ -39,7 +39,9 @@
 
 ctx = SessionContext()
 
-df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns("p_partkey", "p_name")
+df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns(
+    "p_partkey", "p_name"
+)
 df_supplier = ctx.read_parquet(get_data_path("supplier.parquet")).select_columns(
     "s_suppkey", "s_nationkey"
 )

examples/tpch/q13_customer_distribution.py

@@ -41,7 +41,9 @@
 df_orders = ctx.read_parquet(get_data_path("orders.parquet")).select_columns(
     "o_custkey", "o_comment"
 )
-df_customer = ctx.read_parquet(get_data_path("customer.parquet")).select_columns("c_custkey")
+df_customer = ctx.read_parquet(get_data_path("customer.parquet")).select_columns(
+    "c_custkey"
+)
 
 # Use a regex to remove special cases
 df_orders = df_orders.filter(

examples/tpch/q14_promotion_effect.py

@@ -44,7 +44,9 @@
 df_lineitem = ctx.read_parquet(get_data_path("lineitem.parquet")).select_columns(
     "l_partkey", "l_shipdate", "l_extendedprice", "l_discount"
 )
-df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns("p_partkey", "p_type")
+df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns(
+    "p_partkey", "p_type"
+)
 
 
 # Check part type begins with PROMO

examples/tpch/q16_part_supplier_relationship.py

@@ -62,7 +62,8 @@
 # Select the parts we are interested in
 df_part = df_part.filter(col("p_brand") != lit(BRAND))
 df_part = df_part.filter(
-    F.substring(col("p_type"), lit(0), lit(len(TYPE_TO_IGNORE) + 1)) != lit(TYPE_TO_IGNORE)
+    F.substring(col("p_type"), lit(0), lit(len(TYPE_TO_IGNORE) + 1))
+    != lit(TYPE_TO_IGNORE)
 )
 
 # Python conversion of integer to literal casts it to int64 but the data for

examples/tpch/q17_small_quantity_order.py

@@ -56,7 +56,13 @@
 # Find the average quantity
 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")])
+    "avg_quantity",
+    F.window(
+        "avg",
+        [col("l_quantity")],
+        window_frame=window_frame,
+        partition_by=[col("l_partkey")],
+    ),
 )
 
 df = df.filter(col("l_quantity") < lit(0.2) * col("avg_quantity"))

examples/tpch/q20_potential_part_promotion.py

@@ -40,7 +40,9 @@
 
 ctx = SessionContext()
 
-df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns("p_partkey", "p_name")
+df_part = ctx.read_parquet(get_data_path("part.parquet")).select_columns(
+    "p_partkey", "p_name"
+)
 df_lineitem = ctx.read_parquet(get_data_path("lineitem.parquet")).select_columns(
     "l_shipdate", "l_partkey", "l_suppkey", "l_quantity"
 )

examples/tpch/q22_global_sales_opportunity.py

@@ -38,7 +38,9 @@
 df_customer = ctx.read_parquet(get_data_path("customer.parquet")).select_columns(
     "c_phone", "c_acctbal", "c_custkey"
 )
-df_orders = ctx.read_parquet(get_data_path("orders.parquet")).select_columns("o_custkey")
+df_orders = ctx.read_parquet(get_data_path("orders.parquet")).select_columns(
+    "o_custkey"
+)
 
 # The nation code is a two digit number, but we need to convert it to a string literal
 nation_codes = F.make_array(*[lit(str(n)) for n in NATION_CODES])