c

py-polars/tests/unit/datatypes/test_float.py

@@ -99,6 +99,7 @@ def test_hash() -> None:
 
 
 def test_group_by() -> None:
+    # Test num_groups_proxy
     df = (
         pl.Series(
             "x",
@@ -113,25 +114,27 @@ def test_group_by() -> None:
         )
         .to_frame()
         .with_row_index()
+        .with_columns(a=pl.lit("a"))
     )
 
     expect = pl.Series("index", [[0, 1], [2, 3], [4], [5]], dtype=pl.List(pl.UInt32))
 
-    for maintain_order in (True, False):
-        for drop_nulls in (True, False):
-            out = df
-            if drop_nulls:
-                out = out.drop_nulls()
-
-            out = (
-                out.group_by("x", maintain_order=maintain_order)
-                .agg("index")
-                .sort(pl.col("index").list.get(0))
-                .select("index")
-                .to_series()
-            )
-
-            if drop_nulls:
-                assert_series_equal(expect.head(3), out)
-            else:
-                assert_series_equal(expect, out)
+    for group_keys in (("x",), ("x", "a")):
+        for maintain_order in (True, False):
+            for drop_nulls in (True, False):
+                out = df
+                if drop_nulls:
+                    out = out.drop_nulls()
+
+                out = (
+                    out.group_by(group_keys, maintain_order=maintain_order)
+                    .agg("index")
+                    .sort(pl.col("index").list.get(0))
+                    .select("index")
+                    .to_series()
+                )
+
+                if drop_nulls:
+                    assert_series_equal(expect.head(3), out)
+                else:
+                    assert_series_equal(expect, out)