refactor: consolidate order by string equivalence tests into parameterized test

Author: kosiew

python/tests/test_dataframe.py

@@ -925,51 +925,66 @@ def test_window_frame_defaults_match_postgres(partitioned_df):
     assert df_2.sort(col_a).to_pydict() == expected
 
 
-def test_last_value_order_by_string_equivalence(partitioned_df):
-    expr = f.last_value(column("a")).over(
-        Window(
-            partition_by=[column("c")],
-            order_by=[column("b")],
-            window_frame=WindowFrame("rows", None, None),
+def _build_last_value_df(df):
+    return df.select(
+        f.last_value(column("a"))
+        .over(
+            Window(
+                partition_by=[column("c")],
+                order_by=[column("b")],
+                window_frame=WindowFrame("rows", None, None),
+            )
         )
-    )
-    string = f.last_value(column("a")).over(
-        Window(
-            partition_by=[column("c")],
-            order_by="b",
-            window_frame=WindowFrame("rows", None, None),
+        .alias("expr"),
+        f.last_value(column("a"))
+        .over(
+            Window(
+                partition_by=[column("c")],
+                order_by="b",
+                window_frame=WindowFrame("rows", None, None),
+            )
         )
+        .alias("str"),
     )
-    df = partitioned_df.select(expr.alias("expr"), string.alias("str"))
-    table = pa.Table.from_batches(df.collect())
-    assert table.column("expr").to_pylist() == table.column("str").to_pylist()
 
 
-def test_nth_value_order_by_string_equivalence(partitioned_df):
-    expr = f.nth_value(column("b"), 3).over(Window(order_by=[column("a")]))
-    string = f.nth_value(column("b"), 3).over(Window(order_by="a"))
-    df = partitioned_df.select(expr.alias("expr"), string.alias("str"))
-    table = pa.Table.from_batches(df.collect())
-    assert table.column("expr").to_pylist() == table.column("str").to_pylist()
+def _build_nth_value_df(df):
+    return df.select(
+        f.nth_value(column("b"), 3).over(Window(order_by=[column("a")])).alias("expr"),
+        f.nth_value(column("b"), 3).over(Window(order_by="a")).alias("str"),
+    )
 
 
-def test_rank_order_by_string_equivalence(partitioned_df):
-    expr = f.rank(order_by=[column("b")])
-    string = f.rank(order_by="b")
-    df = partitioned_df.select(expr.alias("expr"), string.alias("str"))
-    table = pa.Table.from_batches(df.collect())
-    assert table.column("expr").to_pylist() == table.column("str").to_pylist()
+def _build_rank_df(df):
+    return df.select(
+        f.rank(order_by=[column("b")]).alias("expr"),
+        f.rank(order_by="b").alias("str"),
+    )
 
 
-def test_array_agg_order_by_string_equivalence(partitioned_df):
-    df = partitioned_df.aggregate(
+def _build_array_agg_df(df):
+    return df.aggregate(
         [column("c")],
         [
             f.array_agg(column("a"), order_by=[column("a")]).alias("expr"),
             f.array_agg(column("a"), order_by="a").alias("str"),
         ],
     ).sort(column("c"))
+
+
+@pytest.mark.parametrize(
+    ("builder", "expected"),
+    [
+        pytest.param(_build_last_value_df, [3, 3, 3, 3, 6, 6, 6], id="last_value"),
+        pytest.param(_build_nth_value_df, [None, None, 7, 7, 7, 7, 7], id="nth_value"),
+        pytest.param(_build_rank_df, [1, 1, 3, 3, 5, 6, 6], id="rank"),
+        pytest.param(_build_array_agg_df, [[0, 1, 2, 3], [4, 5, 6]], id="array_agg"),
+    ],
+)
+def test_order_by_string_equivalence(partitioned_df, builder, expected):
+    df = builder(partitioned_df)
     table = pa.Table.from_batches(df.collect())
+    assert table.column("expr").to_pylist() == expected
     assert table.column("expr").to_pylist() == table.column("str").to_pylist()