Repository: dask/dask

Issue:
Performance with isin function on large filter list
I am opening a separate issue from #3198 as there seem to be two separate threads in that issue. Here's reproduce-able code:



```python
import dask.dataframe as dd
import pandas as pd

filter_list = list(range(2, 600000, 2))

for n in [10, 100, 1000]:
    df = dd.from_pandas(pd.DataFrame.from_dict({'A':[1, 2, 3, 4, 5] * 10_000}), npartitions=n)
    %time res = df.A.describe().compute()
```


    CPU times: user 51.9 ms, sys: 4.07 ms, total: 56 ms
    Wall time: 45.3 ms
    CPU times: user 231 ms, sys: 10.5 ms, total: 242 ms
    Wall time: 198 ms
    CPU times: user 2.32 s, sys: 167 ms, total: 2.49 s
    Wall time: 2.01 s


```python
for n in [10, 100, 1000]:
    df = dd.from_pandas(pd.DataFrame.from_dict({'A': [1, 2, 3, 4, 5] * 10_000}), npartitions=n)
    %time res = df[df['A'].isin(filter_list)].compute()
```

    CPU times: user 4.71 s, sys: 44.1 ms, total: 4.75 s
    Wall time: 4.73 s
    CPU times: user 45.2 s, sys: 313 ms, total: 45.5 s
    Wall time: 45.3 s
    CPU times: user 7min 46s, sys: 3.46 s, total: 7min 50s
    Wall time: 7min 48s

```python
for n in [10, 100, 1000, 10000]:
    df = dd.from_pandas(pd.DataFrame.from_dict({'A': [1, 2, 3, 4, 5] * 100_000}), npartitions=n) 
    mask = df.map_partitions(lambda x: x[x.A.isin(filter_list)])
    %time res = df.loc[mask.index].compute()
```

    CPU times: user 243 ms, sys: 12 ms, total: 255 ms
    Wall time: 200 ms
    CPU times: user 2.2 s, sys: 41.4 ms, total: 2.24 s
    Wall time: 1.86 s
    CPU times: user 20.6 s, sys: 214 ms, total: 20.8 s
    Wall time: 17.6 s
    CPU times: user 4min 47s, sys: 3.47 s, total: 4min 51s
    Wall time: 4min 11s

Direct use of `.isin` balloons in time to compute as the number of partitions increases. My observation from watching my system monitor, and as reported in #3198, is that the ballooning time comes from a single threaded processes that happens before the main processing kicks in. My guess is that this is part of building the task graph, but I don't know.

In my real world data processing, using map_partitions wasn't sufficient to get adequate performance and I needed to repartition.


Discussion:
Thanks for the excellent issue report, I'll look into this this morning.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1669,7 +1669,20 @@ def isna(self):
 
     @derived_from(pd.DataFrame)
     def isin(self, values):
-        return elemwise(M.isin, self, list(values))
+        if is_dataframe_like(self._meta):
+            # DataFrame.isin does weird alignment stuff
+            bad_types = (_Frame, pd.Series, pd.DataFrame)
+        else:
+            bad_types = (_Frame,)
+        if isinstance(values, bad_types):
+            raise NotImplementedError(
+                "Passing a %r to `isin`" % typename(type(values))
+            )
+        meta = self._meta_nonempty.isin(values)
+        # We wrap values in a delayed for two reasons:
+        # - avoid serializing data in every task
+        # - avoid cost of traversal of large list in optimizations
+        return self.map_partitions(M.isin, delayed(values), meta=meta)
 
     @derived_from(pd.DataFrame)
     def astype(self, dtype):
@@ -2128,7 +2141,8 @@ def nsmallest(self, n=5, split_every=None):
 
     @derived_from(pd.Series)
     def isin(self, values):
-        return elemwise(M.isin, self, list(values))
+        # Added just to get the different docstring for Series
+        return super(Series, self).isin(values)
 
     @insert_meta_param_description(pad=12)
     @derived_from(pd.Series)


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
@@ -795,13 +795,28 @@ def test_unique():
 
 
 def test_isin():
-    # Series test
-    assert_eq(d.a.isin([0, 1, 2]), full.a.isin([0, 1, 2]))
-    assert_eq(d.a.isin(pd.Series([0, 1, 2])),
-              full.a.isin(pd.Series([0, 1, 2])))
+    f_list = [1, 2, 3]
+    f_series = pd.Series(f_list)
+    f_dict = {'a': [0, 3], 'b': [1, 2]}
+
+    # Series
+    assert_eq(d.a.isin(f_list), full.a.isin(f_list))
+    assert_eq(d.a.isin(f_series), full.a.isin(f_series))
+    with pytest.raises(NotImplementedError):
+        d.a.isin(d.a)
+
+    # Index
+    da.utils.assert_eq(d.index.isin(f_list), full.index.isin(f_list))
+    da.utils.assert_eq(d.index.isin(f_series), full.index.isin(f_series))
+    with pytest.raises(NotImplementedError):
+        d.a.isin(d.a)
 
     # DataFrame test
-    assert_eq(d.isin([0, 1, 2]), full.isin([0, 1, 2]))
+    assert_eq(d.isin(f_list), full.isin(f_list))
+    assert_eq(d.isin(f_dict), full.isin(f_dict))
+    for obj in [d, f_series, full]:
+        with pytest.raises(NotImplementedError):
+            d.isin(obj)
 
 
 def test_len():