Add unit test for avg as window function

timsaucer · timsaucer · commit 707eef6117fe · 2024-08-23T20:46:40.000-04:00
diff --git a/python/datafusion/tests/test_dataframe.py b/python/datafusion/tests/test_dataframe.py
@@ -317,7 +317,7 @@ def test_distinct():
         .build(),
         [-1, -1, 4],
     ),
-    # TODO update once upstream merges https://github.com/apache/datafusion-python/issues/833
+    # TODO update all aggregate functions as windows once upstream merges https://github.com/apache/datafusion-python/issues/833
     pytest.param(
         "first_value",
         f.window("first_value", [column("a")], order_by=[f.order_by(column("b"))]),
@@ -337,6 +337,11 @@ def test_distinct():
         ),
         [None, 5, 5],
     ),
+    pytest.param(
+        "avg",
+        f.window("avg", [column("b")]),
+        [4.0, 4.5, 5.0],
+    ),
 ]
 
 
