add test coverage for single-partition

rapidsai · Jan 11, 2025 · 22cebeb · 22cebeb
1 parent 69f6336
commit 22cebeb
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/python/cudf_polars/tests/experimental/test_groupby.py b/python/cudf_polars/tests/experimental/test_groupby.py
@@ -37,6 +37,21 @@ def test_groupby(df, engine, op, keys):
     assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
+@pytest.mark.parametrize("op", ["sum", "mean", "len"])
+@pytest.mark.parametrize("keys", [("y",), ("y", "z")])
+def test_groupby_single_partitions(df, op, keys):
+    q = getattr(df.group_by(*keys), op)()
+    assert_gpu_result_equal(
+        q,
+        engine=pl.GPUEngine(
+            raise_on_fail=True,
+            executor="dask-experimental",
+            executor_options={"max_rows_per_partition": 1e9},
+        ),
+        check_row_order=False,
+    )
+
+
 @pytest.mark.parametrize("op", ["sum", "mean", "len", "count"])
 @pytest.mark.parametrize("keys", [("y",), ("y", "z")])
 def test_groupby_agg(df, engine, op, keys):