Skip to content

Commit

Permalink
add modifies_index flag
Browse files Browse the repository at this point in the history
  • Loading branch information
FBruzzesi committed Aug 27, 2024
1 parent 749f177 commit a80c94a
Show file tree
Hide file tree
Showing 7 changed files with 169 additions and 20 deletions.
17 changes: 17 additions & 0 deletions narwhals/_dask/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,15 @@ def _from_native_frame(self, df: Any) -> Self:
return self.__class__(df, backend_version=self._backend_version)

def with_columns(self, *exprs: DaskExpr, **named_exprs: DaskExpr) -> Self:
n_modifies_index = sum(
getattr(e, "_modifies_index", 0)
for e in list(exprs) + list(named_exprs.values())
)

if n_modifies_index > 0:
msg = "Expressions that modify the index are not supported in `with_columns`."
raise ValueError(msg)

new_series = parse_exprs_and_named_exprs(self, *exprs, **named_exprs)
return self._from_native_frame(self._native_frame.assign(**new_series))

Expand Down Expand Up @@ -101,6 +110,14 @@ def select(
# This is a simple slice => fastpath!
return self._from_native_frame(self._native_frame.loc[:, exprs])

n_modifies_index = sum(
getattr(e, "_modifies_index", 0)
for e in list(exprs) + list(named_exprs.values())
)
if n_modifies_index > 1:
msg = "Found multiple expressions that modify the index"
raise ValueError(msg)

new_series = parse_exprs_and_named_exprs(self, *exprs, **named_exprs)

if not new_series:
Expand Down
Loading

0 comments on commit a80c94a

Please sign in to comment.