Filtering refactor
Opened by theo-m (closed)

fix https://github.com/huggingface/datasets/issues/2032

benchmarking is somewhat inconclusive, currently running on `book_corpus` with:

```python
    bc = load_dataset("bookcorpus")
    now = time.time()
    bc.filter(lambda x: len(x["text"]) < 64)
    elapsed = time.time() - now
    print(elapsed)
```

this branch does it in 233 seconds, master in 1409 seconds.

Comments: I thought at first that the multiproc test was not relevant now that we do stuff only in memory, but I think there's something that's actually broken, my tiny benchmark on bookcorpus runs forever (2hrs+) when I add `num_proc=4` as a kwarg, will investigate 👀 

I'm not familiar with the caching you describe for `.map`, I'll look it up.

turns out the multi proc issue is also on master, I won't fix it in this PR but opened #2071 to track the problem.

tracemalloc outputs from this script:

```python
import logging
import sys
import time
import tracemalloc

from datasets import load_dataset, set_caching_enabled


if __name__ == "__main__":
    set_caching_enabled(False)
    logging.basicConfig(level=logging.DEBUG)

    tracemalloc.start()
    bc = load_dataset("bookcorpus")

    now = time.time()
    try:
        snapshot1 = tracemalloc.take_snapshot()
        bc["train"].filter(lambda x: len(x["text"]) < 64, num_proc=int(sys.argv[1]))
    except Exception as e:
        print(f"cancelled: {e}")
        exit(1)
    snapshot2 = tracemalloc.take_snapshot()
    tracemalloc.stop()
    elapsed = time.time() - now

    print(elapsed)
    top_stats = snapshot2.compare_to(snapshot1, "lineno")

    print("[ Top 10 differences ]")
    for stat in top_stats[:10]:
        print(stat)

```


This branch:

```
    ssh://theo@35.205.12.130:22/home/theo/.local/share/miniconda3/envs/datasets/bin/python -u benchmark_filter.py 1
    DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): s3.amazonaws.com:443
    DEBUG:urllib3.connectionpool:https://s3.amazonaws.com:443 "HEAD /datasets.huggingface.co/datasets/datasets/bookcorpus/bookcorpus.py HTTP/1.1" 200 0
    DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): raw.githubusercontent.com:443
    DEBUG:urllib3.connectionpool:https://raw.githubusercontent.com:443 "HEAD /huggingface/datasets/master/datasets/bookcorpus/bookcorpus.py HTTP/1.1" 200 0
    DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): raw.githubusercontent.com:443
    DEBUG:urllib3.connectionpool:https://raw.githubusercontent.com:443 "HEAD /huggingface/datasets/master/datasets/bookcorpus/dataset_infos.json HTTP/1.1" 200 0
    WARNING:datasets.builder:Reusing dataset bookcorpus (/home/theo/.cache/huggingface/datasets/bookcorpus/plain_text/1.0.0/af844be26c089fb64810e9f2cd841954fd8bd596d6ddd26326e4c70e2b8c96fc)
      0%|                                                 | 0/74005 [00:00<?, ?ba/s]2021-03-23 10:23:20.051255: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
    2021-03-23 10:23:20.051304: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
    DEBUG:tensorflow:Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
    100%|████████████████████████████████████| 74005/74005 [12:16<00:00, 100.54ba/s]
    815.6356580257416
    [ Top 10 differences ]
    <frozen importlib._bootstrap_external>:580: size=38.0 MiB (+33.7 MiB), count=326226 (+307928), average=122 B
    <frozen importlib._bootstrap>:219: size=7643 KiB (+7553 KiB), count=26372 (+25473), average=297 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/torch/__init__.py:427: size=1291 KiB (+1291 KiB), count=5924 (+5924), average=223 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/abc.py:85: size=1039 KiB (+1026 KiB), count=3428 (+3384), average=310 B
    <frozen importlib._bootstrap_external>:64: size=917 KiB (+891 KiB), count=5300 (+5132), average=177 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/collections/__init__.py:456: size=720 KiB (+709 KiB), count=3403 (+3349), average=217 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/tensorflow/python/util/tf_export.py:346: size=607 KiB (+607 KiB), count=3962 (+3962), average=157 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/linecache.py:137: size=998 KiB (+487 KiB), count=9551 (+4517), average=107 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/tensorflow/python/util/tf_decorator.py:241: size=367 KiB (+367 KiB), count=5225 (+5225), average=72 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/tensorflow/python/util/decorator_utils.py:114: size=359 KiB (+359 KiB), count=330 (+330), average=1114 B
```

On master:
```
    ssh://theo@35.205.12.130:22/home/theo/.local/share/miniconda3/envs/datasets/bin/python -u benchmark_filter.py 1
    DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): s3.amazonaws.com:443
    DEBUG:urllib3.connectionpool:https://s3.amazonaws.com:443 "HEAD /datasets.huggingface.co/datasets/datasets/bookcorpus/bookcorpus.py HTTP/1.1" 200 0
    DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): raw.githubusercontent.com:443
    DEBUG:urllib3.connectionpool:https://raw.githubusercontent.com:443 "HEAD /huggingface/datasets/master/datasets/bookcorpus/bookcorpus.py HTTP/1.1" 200 0
    DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): raw.githubusercontent.com:443
    DEBUG:urllib3.connectionpool:https://raw.githubusercontent.com:443 "HEAD /huggingface/datasets/master/datasets/bookcorpus/dataset_infos.json HTTP/1.1" 200 0
    WARNING:datasets.builder:Reusing dataset bookcorpus (/home/theo/.cache/huggingface/datasets/bookcorpus/plain_text/1.0.0/af844be26c089fb64810e9f2cd841954fd8bd596d6ddd26326e4c70e2b8c96fc)
      0%|                                                 | 0/74005 [00:00<?, ?ba/s]2021-03-23 12:26:47.219622: W tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
    2021-03-23 12:26:47.219669: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
    DEBUG:tensorflow:Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
    100%|███████████████████████████████████| 74005/74005 [1:02:17<00:00, 19.80ba/s]
    3738.870892047882
    [ Top 10 differences ]
    <frozen importlib._bootstrap_external>:580: size=38.0 MiB (+33.7 MiB), count=326221 (+307919), average=122 B
    <frozen importlib._bootstrap>:219: size=7648 KiB (+7557 KiB), count=26455 (+25555), average=296 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/torch/__init__.py:427: size=1291 KiB (+1291 KiB), count=5924 (+5924), average=223 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/abc.py:85: size=1039 KiB (+1026 KiB), count=3429 (+3385), average=310 B
    <frozen importlib._bootstrap_external>:64: size=917 KiB (+891 KiB), count=5300 (+5132), average=177 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/collections/__init__.py:456: size=720 KiB (+709 KiB), count=3403 (+3349), average=217 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/tensorflow/python/util/tf_export.py:346: size=607 KiB (+607 KiB), count=3962 (+3962), average=157 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/linecache.py:137: size=1000 KiB (+489 KiB), count=9569 (+4535), average=107 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/tensorflow/python/util/tf_decorator.py:241: size=367 KiB (+367 KiB), count=5225 (+5225), average=72 B
    /home/theo/.local/share/miniconda3/envs/datasets/lib/python3.8/site-packages/tensorflow/python/util/decorator_utils.py:114: size=359 KiB (+359 KiB), count=330 (+330), average=1114 B
```

I'm not concluding much, it seems nothing is really happening to memory on `pyarrow::Table.filter`? 

Cool ! Maybe it increases the memory a bit but what's brought in memory is not the resulting Table but something else (not sure what though).
What's the length of the resulting dataset ?
You can also take a look at `pyarrow.total_allocated_memory()` to show how much memory is being used by pyarrow

```diff
diff --git a/benchmarks/benchmark_filter.py b/benchmarks/benchmark_filter.py
index 4b9efd4e..a862c204 100644
--- a/benchmarks/benchmark_filter.py
+++ b/benchmarks/benchmark_filter.py
@@ -1,6 +1,9 @@
 import logging
 import sys
 import time
+import tracemalloc
+
+import pyarrow as pa
 
 from datasets import load_dataset, set_caching_enabled
 
@@ -9,13 +12,28 @@ if __name__ == "__main__":
     set_caching_enabled(False)
     logging.basicConfig(level=logging.DEBUG)
 
+    tracemalloc.start()
     bc = load_dataset("bookcorpus")
 
     now = time.time()
     try:
+        snapshot1 = tracemalloc.take_snapshot()
+        pamem1 = pa.total_allocated_bytes()
         bc["train"].filter(lambda x: len(x["text"]) < 64, num_proc=int(sys.argv[1]))
+        pamem2 = pa.total_allocated_bytes()
+        snapshot2 = tracemalloc.take_snapshot()
     except Exception as e:
         print(f"cancelled: {e}")
+        exit(1)
+    tracemalloc.stop()
     elapsed = time.time() - now
 
     print(elapsed)
+    top_stats = snapshot2.compare_to(snapshot1, "lineno")
+
+    print("[ Top 10 differences ]")
+    for stat in top_stats[:10]:
+        print(stat)
+
+    print("[ pyarrow reporting ]")
+    print(f"before: ({pamem1}) after: ({pamem2})")
```

this yields 0-0, does not seem like a good tool 😛 and the documentation is [quite mysterious.](https://arrow.apache.org/docs/python/generated/pyarrow.total_allocated_bytes.html)

Personally if I use your script to benchmark on this branch
```python
bc = load_dataset("bookcorpus", split="train[:1%]")
bc = bc.filter(lambda x: len(x["text"]) < 64)
```

then I get
```
[ pyarrow reporting ]
before: (0) after: (15300672)
```

Maybe you got 0-0 because the filter output is directly garbage collected, since you didn't do
```python
bc["train"] = bc["train"].filter(...)
```
Can you try again on your side just to make sure ?

Even if the documentation doesn't say much, `pa.total_allocated_bytes` if pretty useful, and also very consistent.
It tracks the number of bytes used for arrow data.

> Maybe you got 0-0 because the filter output is directly garbage collected, since you didn't do
> 
> ```python
> bc["train"] = bc["train"].filter(...)
> ```
Nice catch! I get 1.74GB for this branch

Looks like we may need to write the filtered table on the disk then.

The other option is to slice the table to keep only the good rows and concatenate them but this is too slow at the moment since slicing is O(n) until #1803 is fixed. I'll work on this issue this afternoon

From investigation it looks like the lib's `Table.filter` cannot send its output to memorymap, asked a question on the mailing list, see [here](https://lists.apache.org/thread.html/r8cd8591ce83a967eb0097a7f31785ac2f3ee95ea371c8c5beb0720ad%40%3Cuser.arrow.apache.org%3E)

closing in favor of #2836