Remove TODO note

rapidsai · Oct 1, 2024 · 7ef586b · 7ef586b
1 parent 20f1a94
commit 7ef586b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/python/cudf_polars/cudf_polars/containers/column.py b/python/cudf_polars/cudf_polars/containers/column.py
@@ -84,8 +84,7 @@ def sorted_like(self, like: Column, /) -> Self:
             is_sorted=like.is_sorted, order=like.order, null_order=like.null_order
         )
 
-    # TODO: Return Column once #16272 is fixed.
-    def astype(self, dtype: plc.DataType) -> plc.Column:
+    def astype(self, dtype: plc.DataType) -> Column:
         """
         Return the backing column as the requested dtype.
 
@@ -109,8 +108,8 @@ def astype(self, dtype: plc.DataType) -> plc.Column:
         the current one.
         """
         if self.obj.type() != dtype:
-            return plc.unary.cast(self.obj, dtype)
-        return self.obj
+            return Column(plc.unary.cast(self.obj, dtype)).sorted_like(self)
+        return self
 
     def copy_metadata(self, from_: pl.Series, /) -> Self:
         """

diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py
@@ -1205,7 +1205,7 @@ def evaluate(self, *, cache: MutableMapping[int, DataFrame]) -> DataFrame:
             ).columns()
             value_column = plc.concatenate.concatenate(
                 [
-                    df.column_map[pivotee].astype(self.schema[value_name])
+                    df.column_map[pivotee].astype(self.schema[value_name]).obj
                     for pivotee in pivotees
                 ]
             )