Repository: dask/dask

Issue:
Support assignment of scalars based on boolean indexes
With Dask DataFrame, I often want to update values for columns only if some condition holds true for the values of one or more columns in those rows.

For example:
```
>>> import pandas as pd
>>> import dask.dataframe as dd
>>> 
>>> df = pd.DataFrame({'id': [0, 1, 2, None]})
>>> ddf = dd.from_pandas(df, npartitions=1)
>>> ddf[ddf['id'] > 1] = -1
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/rgelhausen/conda/envs/jit/lib/python3.8/site-packages/dask/dataframe/core.py", line 4121, in __setitem__
    raise NotImplementedError(f"Item assignment with {type(key)} not supported")
NotImplementedError: Item assignment with <class 'dask.dataframe.core.Series'> not supported
```

This is supported in Pandas:
```
>>> df[df['id'] > 1] = -1
>>> df
    id
0  0.0
1  1.0
2 -1.0
3  NaN
```


Discussion:
We've wanted to do this for a while, but mutable semantics have made it
hard (I did this once in a PR, Julia did too I think)

The workaround is to use the where method I think.  Short term maybe we can
improve the error message to point to this method?  Longer term I'd love to
see proper support here if that's doable.

(This is a drive by comment, I apologize if I don't follow up to responses)

On Tue, Sep 28, 2021, 3:19 PM Randy Gelhausen ***@***.***>
wrote:

> With Dask DataFrame, I often want to update values for columns only if
> some condition holds true for the values of one or more columns in those
> rows.
>
> For example:
>
> >>> import pandas as pd
> >>> import dask.dataframe as dd
> >>>
> >>> df = pd.DataFrame({'id': [0, 1, 2, None]})
> >>> ddf = dd.from_pandas(df, npartitions=1)
> >>> ddf[ddf['id'] > 1] = -1
> Traceback (most recent call last):
>   File "<stdin>", line 1, in <module>
>   File "/home/rgelhausen/conda/envs/jit/lib/python3.8/site-packages/dask/dataframe/core.py", line 4121, in __setitem__
>     raise NotImplementedError(f"Item assignment with {type(key)} not supported")
> NotImplementedError: Item assignment with <class 'dask.dataframe.core.Series'> not supported
>
> This is supported in Pandas:
>
> >>> df[df['id'] > 1] = -1
> >>> df
>     id
> 0  0.0
> 1  1.0
> 2 -1.0
> 3  NaN
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/8190>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTABNA3N4TW6D5JJWOLUEIPNHANCNFSM5E6H4RVQ>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.
>
>


Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -4115,7 +4115,11 @@ def __setitem__(self, key, value):
         elif isinstance(key, pd.Index) and not isinstance(value, DataFrame):
             key = list(key)
             df = self.assign(**{k: value for k in key})
-        elif is_dataframe_like(key) or isinstance(key, DataFrame):
+        elif (
+            is_dataframe_like(key)
+            or is_series_like(key)
+            or isinstance(key, (DataFrame, Series))
+        ):
             df = self.where(~key, value)
         elif not isinstance(key, str):
             raise NotImplementedError(f"Item assignment with {type(key)} not supported")


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -4464,6 +4464,14 @@ def test_setitem_with_bool_dataframe_as_key():
     assert_eq(df, ddf)
 
 
+def test_setitem_with_bool_series_as_key():
+    df = pd.DataFrame({"A": [1, 4], "B": [3, 2]})
+    ddf = dd.from_pandas(df.copy(), 2)
+    df[df["A"] > 2] = 5
+    ddf[ddf["A"] > 2] = 5
+    assert_eq(df, ddf)
+
+
 def test_setitem_with_numeric_column_name_raises_not_implemented():
     df = pd.DataFrame({0: [1, 4], 1: [3, 2]})
     ddf = dd.from_pandas(df.copy(), 2)