Skip to content

Commit

Permalink
Column.copy takes an optional new_name argument
Browse files Browse the repository at this point in the history
This removes the need for rename.
  • Loading branch information
wence- committed May 30, 2024
1 parent 62f6455 commit a1f579f
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 9 deletions.
40 changes: 32 additions & 8 deletions python/cudf_polars/cudf_polars/containers/column.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,23 @@ def __init__(self, column: plc.Column, name: str):
self.order = plc.types.Order.ASCENDING
self.null_order = plc.types.NullOrder.BEFORE

def rename(self, name: str) -> Column:
"""Return a new column sharing data with a new name."""
return type(self)(self.obj, name).sorted_like(self)

def sorted_like(self, like: Column, /) -> Self:
"""Copy sortedness properties from a column onto self."""
"""
Copy sortedness properties from a column onto self.
Parameters
----------
like
The column to copy sortedness metadata from.
Returns
-------
Self with metadata set.
See Also
--------
set_sorted
"""
return self.set_sorted(
is_sorted=like.is_sorted, order=like.order, null_order=like.null_order
)
Expand Down Expand Up @@ -70,9 +81,22 @@ def set_sorted(
self.null_order = null_order
return self

def copy(self) -> Self:
"""Return a shallow copy of the column."""
return type(self)(self.obj, self.name).sorted_like(self)
def copy(self, *, new_name: str | None = None) -> Self:
"""
Return a shallow copy of the column.
Parameters
----------
new_name
Optional new name for the copied column.
Returns
-------
New column sharing data with self.
"""
return type(self)(
self.obj, self.name if new_name is None else new_name
).sorted_like(self)

def mask_nans(self) -> Self:
"""Return a copy of self with nans masked out."""
Expand Down
2 changes: 1 addition & 1 deletion python/cudf_polars/cudf_polars/containers/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ def replace_columns(self, *columns: Column) -> Self:
def rename_columns(self, mapping: Mapping[str, str]) -> Self:
"""Rename some columns."""
return type(self)(
[c.rename(mapping.get(c.name, c.name)) for c in self.columns], self.scalars
[c.copy(new_name=mapping.get(c.name)) for c in self.columns], self.scalars
)

def select_columns(self, names: Set[str]) -> list[Column]:
Expand Down

0 comments on commit a1f579f

Please sign in to comment.