Repository: dask/dask
Pull request: Filter by list
- [ X ] Closes #9411 
- [ X ] Tests added / passed
- [ X ] Passes `pre-commit run --all-files`

This PR addresses the fact that filtering a Dask DataFrame with `ddf.isin(L)` if L is a large list, is much slower than if L is a numpy array.  This appears to be a result of serializing lists with pickle, while serializing a numpy array uses dask's custom serializer.

Here we do a few things to reduce the serialization cost if L is a list:

1. If L is a list, we remove duplicate
2. If all the elements in L are of the same datatype, we convert L to a numpy array, otherwise we are forced to incur the serialization cost


Problem statement:
Filtering on a Python list is 2 orders of magnitude slower than filtering on a numpy array
**What happened**:
Filtering on a list of integers vs a numpy array on the example below is nearly 200 times slower 
From the example below I get

```
Pandas with list took: 0.573s, len: 3,934,884
Pandas with array took: 0.366s, len: 3,934,884
Dask with list took: 41.739s, len: 3,934,884
Dask with array took: 0.240s, len: 3,934,884
```

**What you expected to happen**:

Filtering on a list vs an array shouldn't be that different in Dask, I would expect it to be a bit slower, but not nearly 200 times slower.

**Minimal Complete Verifiable Example**:

```python
import time
import numpy as np
import pandas as pd
from dask.dataframe import from_pandas
from distributed import LocalCluster, Client


def main():
    n = 10_000_000
    rs = np.random.RandomState(42)
    a_column_unique_values = np.arange(1, n // 10)
    df = pd.DataFrame({"A": rs.choice(a_column_unique_values, n), "B": rs.random(n)})

    filter_values_list = sorted(
        rs.choice(a_column_unique_values, len(a_column_unique_values) // 2).tolist()
    )
    filter_values_array = np.array(filter_values_list, dtype="i8")
    print(f"Len filter: {len(filter_values_list):,}")

    t1 = time.time()
    tmp_df = df[df["A"].isin(filter_values_list)].copy()
    len_after_filter = len(tmp_df)
    print(f"Pandas with list took: {time.time() - t1:,.3f}s, len: {len_after_filter:,}")

    t1 = time.time()
    tmp_df = df[df["A"].isin(filter_values_array)].copy()
    len_after_filter = len(tmp_df)
    print(f"Pandas with array took: {time.time() - t1:,.3f}s, len: {len_after_filter:,}")

    with LocalCluster(n_workers=4, processes=True, threads_per_worker=1) as local_cluster:
        with Client(local_cluster):
            ddf = from_pandas(df, npartitions=8).persist()

            t1 = time.time()
            tmp_ddf = ddf[ddf["A"].isin(filter_values_list)].copy()
            len_after_filter = len(tmp_ddf)
            t2 = time.time()
            print(f"Dask with list took: {t2-t1:,.3f}s, len: {len_after_filter:,}")

            t1 = time.time()
            tmp_ddf = ddf[ddf["A"].isin(filter_values_array)].copy()
            len_after_filter = len(tmp_ddf)
            t2 = time.time()
            print(f"Dask with array took: {t2-t1:,.3f}s, len: {len_after_filter:,}")


if __name__ == "__main__":
    main()
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2022.8.1
- Python version:  3.9.13
- Operating System: MacOS
- Install method (conda, pip, source): pip

Hints from issue discussion:
@RogerThomas Thanks for reporting, this is surprising and seems like a bug.

I looked at this with @ian-r-rose, and we think this might be related to how the list is being serialized, specifically by the distributed scheduler.

A few notes:
* we see short overall times and short time _differences_ while using single machine schedulers
* for the list case, the task stream shows some long red bars called "transfer-len-chunks" -- seem to be taking a lot of time
* again or the list case, the profile plot shows many large "serialize" steps, which aren't present in the np.array case
thanks for looking into it @pavithraes, is there anything I can help with?
Hi @RogerThomas .  I've just added a PR that should improve this.  With your reproducer, on my machine, this PR reduces execution time for the list from 2 orders of magnitude greater than Pandas to 1.6x longer than in Pandas.  Most of this time is spent confirming that all of the elements in the list are of same datatype.

Would appreciate any feedback.

Commit message: Handle versions of numpy that can not create an array of lists

Patch to review:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index e9cd2795fbc..5df408e5eb5 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -3094,6 +3094,22 @@ def isin(self, 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
+        if isinstance(values, list):
+            # Motivated by https://github.com/dask/dask/issues/9411.  This appears to be
+            # caused by https://github.com/dask/distributed/issues/6368, and further
+            # exacerbated by the fact that the list contains duplicates.  This is a patch until
+            # we can create a better fix for Serialization.
+            try:
+                values = list(set(values))
+            except TypeError:
+                pass
+            if not any(is_dask_collection(v) for v in values):
+                try:
+                    values = np.fromiter(values, dtype=object)
+                except ValueError:
+                    # Numpy 1.23 supports creating arrays of iterables, while lower
+                    # version 1.21.x and 1.22.x do not
+                    pass
         return self.map_partitions(
             M.isin, delayed(values), meta=meta, enforce_metadata=False
         )
@@ -6517,7 +6533,6 @@ def map_partitions(
     dfs = [df for df in args if isinstance(df, _Frame)]
 
     meta = _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta)
-
     if all(isinstance(arg, Scalar) for arg in args):
         layer = {
             (name, 0): (
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 30995aa8003..f479749ad39 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -15,6 +15,7 @@
 import dask.array as da
 import dask.dataframe as dd
 import dask.dataframe.groupby
+from dask import delayed
 from dask.base import compute_as_if_collection
 from dask.blockwise import fuse_roots
 from dask.dataframe import _compat, methods
@@ -50,6 +51,18 @@
 )
 d = dd.DataFrame(dsk, "x", meta, [0, 5, 9, 9])
 full = d.compute()
+
+dsk2 = {
+    ("x", 0): pd.DataFrame({"a": [[1, 2, 3]], "b": [4, 5, 6]}, index=[0, 1, 3]),
+    ("x", 1): pd.DataFrame({"a": [4, 5, 6], "b": [[3, 4, 5]]}, index=[5, 6, 8]),
+    ("x", 2): pd.DataFrame({"a": [7, 8, 9], "b": [0, 0, 0]}, index=[9, 9, 9]),
+}
+meta2 = make_meta(
+    {"a": "i8", "b": "i8"}, index=pd.Index([], "i8"), parent_meta=pd.DataFrame()
+)
+d2 = dd.DataFrame(dsk2, "x", meta2, [0, 5, 9, 9])
+full2 = d2.compute()
+
 CHECK_FREQ = {}
 if dd._compat.PANDAS_GT_110:
     CHECK_FREQ["check_freq"] = False
@@ -1271,22 +1284,48 @@ def test_isin():
     f_list = [1, 2, 3]
     f_series = pd.Series(f_list)
     f_dict = {"a": [0, 3], "b": [1, 2]}
+    f_list2 = [1, "2"]
+    f_list_delayed = [delayed(1), delayed(2), delayed(3)]
+    f_list_of_lists = [[1, 2, 3], [2, 3, 4], [3, 4, 5]]
 
     # 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))
+    assert_eq(d.a.isin(f_list2), full.a.isin(f_list2))
+    assert_eq(
+        d.a.isin(f_list_delayed),
+        full.a.isin(f_list),
+    )
+    assert_eq(
+        d.a.isin(f_list_of_lists),
+        full.a.isin(f_list_of_lists),
+    )
     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))
+    da.utils.assert_eq(d.index.isin(f_list2), full.index.isin(f_list2))
+    da.utils.assert_eq(
+        d.index.isin(f_list_delayed),
+        full.index.isin(f_list),
+    )
     with pytest.raises(NotImplementedError):
         d.a.isin(d.a)
 
     # DataFrame test
     assert_eq(d.isin(f_list), full.isin(f_list))
     assert_eq(d.isin(f_dict), full.isin(f_dict))
+    assert_eq(d.isin(f_list2), full.isin(f_list2))
+    assert_eq(
+        d.isin(f_list_delayed),
+        full.isin(f_list),
+    )
+    assert_eq(
+        d.isin(f_list_of_lists),
+        full.isin(f_list_of_lists),
+    )
     for obj in [d, f_series, full]:
         with pytest.raises(NotImplementedError):
             d.isin(obj)


Code review comments:
diff_hunk: @@ -50,6 +51,18 @@
 )
 d = dd.DataFrame(dsk, "x", meta, [0, 5, 9, 9])
 full = d.compute()
+
+dsk2 = {
+    ("x", 0): pd.DataFrame({"a": [[1, 2, 3]], "b": [4, 5, 6]}, index=[0, 1, 3]),
+    ("x", 1): pd.DataFrame({"a": [4, 5, 6], "b": [[3, 4, 5]]}, index=[5, 6, 8]),
+    ("x", 2): pd.DataFrame({"a": [7, 8, 9], "b": [0, 0, 0]}, index=[9, 9, 9]),
+}
+meta2 = make_meta(
+    {"a": "i8", "b": "i8"}, index=pd.Index([], "i8"), parent_meta=pd.DataFrame()
+)
+d2 = dd.DataFrame(dsk2, "x", meta2, [0, 5, 9, 9])
+full2 = d2.compute()
path: dask/dataframe/tests/test_dataframe.py
text: @user1:
This looks identical to the above dataframe `d`, and is unused below. Is there a reason to make another?

A meta point -- the above snippet constructing a dataframe looks quite old, and for new tests I would probably advocate for using `dd.from_pandas()` rather than doing manual construction.

@author:
Done.