address code review

python/cudf_polars/cudf_polars/experimental/shuffle.py

@@ -6,7 +6,6 @@
 
 import json
 import operator
-from functools import reduce
 from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
@@ -167,25 +166,20 @@ def _(
     # be handled separately.
     from cudf_polars.experimental.parallel import PartitionInfo
 
-    # Extract child partitioning
-    children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True)
-    partition_info = reduce(operator.or_, _partition_info)
-    pi = partition_info[children[0]]
-
-    # Check if we are already shuffled or update partition_info
-    if ir.keys == pi.partitioned_on:
-        # Already shuffled!
-        new_node = children[0]
-    else:
-        new_node = ir.reconstruct(children)
-        partition_info[new_node] = PartitionInfo(
-            # Default shuffle preserves partition count
-            count=pi.count,
-            # Add partitioned_on info
-            partitioned_on=ir.keys,
-        )
-
-    return new_node, partition_info
+    (child,) = ir.children
+
+    new_child, pi = rec(child)
+    if ir.keys == pi[new_child].partitioned_on:
+        # Already shuffled
+        return new_child, pi
+    new_node = ir.reconstruct([new_child])
+    pi[new_node] = PartitionInfo(
+        # Default shuffle preserves partition count
+        count=pi[new_child].count,
+        # Add partitioned_on info
+        partitioned_on=ir.keys,
+    )
+    return new_node, pi
 
 
 @generate_ir_tasks.register(Shuffle)