Skip to content

Commit

Permalink
clean (#349)
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli authored Jun 29, 2024
1 parent a2b2ec8 commit 5c018ca
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 15 deletions.
21 changes: 12 additions & 9 deletions narwhals/_pandas_like/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,31 +190,34 @@ def with_columns(
**named_exprs: IntoPandasExpr,
) -> Self:
index = self._dataframe.index
new_series = evaluate_into_exprs(self, *exprs, **named_exprs)
new_columns = evaluate_into_exprs(self, *exprs, **named_exprs)
# If the inputs are all Expressions which return full columns
# (as opposed to scalars), we can use a fast path (concat, instead of assign).
# We can't use the fastpath if any input is not an expression (e.g.
# if it's a Series) because then we might be changing its flags.
# See `test_memmap` for an example of where this is necessary.
fast_path = (
all(s.len() > 1 for s in new_series)
all(s.len() > 1 for s in new_columns)
and all(isinstance(x, PandasExpr) for x in exprs)
and all(isinstance(x, PandasExpr) for (_, x) in named_exprs.items())
)

if fast_path:
new_names = {s.name: s for s in new_series}
new_column_name_to_new_column_map = {s.name: s for s in new_columns}
to_concat = []
# Make sure to preserve column order
for s in self._dataframe.columns:
if s in new_names:
for name in self._dataframe.columns:
if name in new_column_name_to_new_column_map:
to_concat.append(
validate_dataframe_comparand(index, new_names.pop(s))
validate_dataframe_comparand(
index, new_column_name_to_new_column_map.pop(name)
)
)
else:
to_concat.append(self._dataframe.loc[:, s])
to_concat.append(self._dataframe.loc[:, name])
to_concat.extend(
validate_dataframe_comparand(index, new_names[s]) for s in new_names
validate_dataframe_comparand(index, new_column_name_to_new_column_map[s])
for s in new_column_name_to_new_column_map
)

df = horizontal_concat(
Expand All @@ -223,7 +226,7 @@ def with_columns(
)
else:
df = self._dataframe.assign(
**{s.name: validate_dataframe_comparand(index, s) for s in new_series}
**{s.name: validate_dataframe_comparand(index, s) for s in new_columns}
)
return self._from_dataframe(df)

Expand Down
10 changes: 4 additions & 6 deletions narwhals/_pandas_like/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,15 +75,13 @@ def __init__(
*,
implementation: str,
) -> None:
"""Parameters
----------
df
DataFrame this column originates from.
"""

self._name = series.name
self._series = series
self._implementation = implementation

# In pandas, copy-on-write becomes the default in version 3.
# So, before that, we need to explicitly avoid unnecessary
# copies by using `copy=False` sometimes.
self._use_copy_false = False
if self._implementation == "pandas":
pd = get_pandas()
Expand Down

0 comments on commit 5c018ca

Please sign in to comment.