fix groupby

rapidsai · Nov 21, 2024 · 709d0d6 · 709d0d6
1 parent 242e79d
commit 709d0d6
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/cudf_polars/cudf_polars/experimental/groupby.py b/python/cudf_polars/cudf_polars/experimental/groupby.py
@@ -119,6 +119,8 @@ def lower_groupby_node(
         ir.options,
         *children,
     )
+    child_count = partition_info[children[0]].count
+    partition_info[gb_pwise] = PartitionInfo(count=child_count)
 
     gb_tree = GroupByTree(
         ir.schema,
@@ -128,6 +130,7 @@ def lower_groupby_node(
         ir.options,
         gb_pwise,
     )
+    partition_info[gb_tree] = PartitionInfo(count=1)
 
     schema = ir.schema
     output_exprs = []