Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

perf: concatenate less in pandas-like with_columns #1361

Closed
wants to merge 2 commits into from
Closed

perf: concatenate less in pandas-like with_columns #1361

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0285ae8
perf: concatenate less in pandas-like with_columns
MarcoGorelli Nov 13, 2024
fd5aca4
simplify
MarcoGorelli Nov 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 20 additions & 16 deletions narwhals/_pandas_like/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -422,27 +422,31 @@ def with_columns(

if fast_path:
new_column_name_to_new_column_map = {s.name: s for s in new_columns}
to_concat = []
# Make sure to preserve column order
for name in self._native_frame.columns:
if name in new_column_name_to_new_column_map:
to_concat.append(
validate_dataframe_comparand(
index, new_column_name_to_new_column_map.pop(name)
)
)
else:
to_concat.append(self._native_frame[name])
to_concat.extend(
validate_dataframe_comparand(index, new_column_name_to_new_column_map[s])
for s in new_column_name_to_new_column_map
columns_that_stay_the_same = list(
set(self.columns).difference(new_column_name_to_new_column_map)
)

frame = select_columns_by_name(
self._native_frame,
columns_that_stay_the_same,
self._backend_version,
self._implementation,
)
other = [
validate_dataframe_comparand(index, value)
for key, value in new_column_name_to_new_column_map.items()
if key not in frame.columns
]
df = horizontal_concat(
to_concat,
[frame, *other],
implementation=self._implementation,
backend_version=self._backend_version,
)
order = self.columns + [
x for x in new_column_name_to_new_column_map if x not in self.columns
]
df = select_columns_by_name(
df, order, self._backend_version, self._implementation
)
else:
# This is the logic in pandas' DataFrame.assign
if self._backend_version < (2,):
Expand Down
Loading