Simplify some broadcasts

rapidsai · Oct 3, 2024 · 3ab5bc3 · 3ab5bc3
1 parent 9f7cb38
commit 3ab5bc3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py
@@ -1007,7 +1007,9 @@ def __init__(
     def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
         """Evaluate and return a dataframe."""
         df = self.df.evaluate(cache=cache)
-        sort_keys = broadcast(*(k.evaluate(df) for k in self.by))
+        sort_keys = broadcast(
+            *(k.evaluate(df) for k in self.by), target_length=df.num_rows
+        )
         # TODO: More robust identification here.
         keys_in_result = {
             k.name: i
@@ -1080,7 +1082,7 @@ def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
         df = self.df.evaluate(cache=cache)
         # This can reorder things.
         columns = broadcast(
-            *(df.select(list(self.schema.keys())).columns), target_length=df.num_rows
+            *(df.column_map[name] for name in self.schema), target_length=df.num_rows
         )
         return DataFrame(columns)