Skip to content

Commit 4c2b60e

Browse files
authored
perf: simplify pandas-like with_columns (#1366)
* wip try using fastpath always * simplify
1 parent 8f2e4a9 commit 4c2b60e

File tree

1 file changed

+20
-39
lines changed

1 file changed

+20
-39
lines changed

narwhals/_pandas_like/dataframe.py

+20-39
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
from typing import overload
1010

1111
from narwhals._expression_parsing import evaluate_into_exprs
12-
from narwhals._pandas_like.expr import PandasLikeExpr
1312
from narwhals._pandas_like.utils import broadcast_series
1413
from narwhals._pandas_like.utils import convert_str_slice_to_int_slice
1514
from narwhals._pandas_like.utils import create_compliant_series
@@ -411,46 +410,28 @@ def with_columns(
411410
if not new_columns and len(self) == 0:
412411
return self
413412

414-
# If the inputs are all Expressions
415-
# (as opposed to Series), we can use a fast path (concat, instead of assign).
416-
# We can't use the fastpath if any input is not an expression (e.g.
417-
# if it's a Series) because then we might be changing its flags.
418-
# See `test_memmap` for an example of where this is necessary.
419-
fast_path = all(isinstance(x, PandasLikeExpr) for x in exprs) and all(
420-
isinstance(x, PandasLikeExpr) for (_, x) in named_exprs.items()
421-
)
422-
423-
if fast_path:
424-
new_column_name_to_new_column_map = {s.name: s for s in new_columns}
425-
to_concat = []
426-
# Make sure to preserve column order
427-
for name in self._native_frame.columns:
428-
if name in new_column_name_to_new_column_map:
429-
to_concat.append(
430-
validate_dataframe_comparand(
431-
index, new_column_name_to_new_column_map.pop(name)
432-
)
413+
new_column_name_to_new_column_map = {s.name: s for s in new_columns}
414+
to_concat = []
415+
# Make sure to preserve column order
416+
for name in self._native_frame.columns:
417+
if name in new_column_name_to_new_column_map:
418+
to_concat.append(
419+
validate_dataframe_comparand(
420+
index, new_column_name_to_new_column_map.pop(name)
433421
)
434-
else:
435-
to_concat.append(self._native_frame[name])
436-
to_concat.extend(
437-
validate_dataframe_comparand(index, new_column_name_to_new_column_map[s])
438-
for s in new_column_name_to_new_column_map
439-
)
440-
441-
df = horizontal_concat(
442-
to_concat,
443-
implementation=self._implementation,
444-
backend_version=self._backend_version,
445-
)
446-
else:
447-
# This is the logic in pandas' DataFrame.assign
448-
if self._backend_version < (2,):
449-
df = self._native_frame.copy(deep=True)
422+
)
450423
else:
451-
df = self._native_frame.copy(deep=False)
452-
for s in new_columns:
453-
df[s.name] = validate_dataframe_comparand(index, s)
424+
to_concat.append(self._native_frame[name])
425+
to_concat.extend(
426+
validate_dataframe_comparand(index, new_column_name_to_new_column_map[s])
427+
for s in new_column_name_to_new_column_map
428+
)
429+
430+
df = horizontal_concat(
431+
to_concat,
432+
implementation=self._implementation,
433+
backend_version=self._backend_version,
434+
)
454435
return self._from_native_frame(df)
455436

456437
def rename(self, mapping: dict[str, str]) -> Self:

0 commit comments

Comments
 (0)