Filtering/mapping on one column is very slow

I'm currently using the `wikipedia` dataset— I'm tokenizing the articles with the `tokenizers` library using `map()` and also adding a new `num_tokens` column to the dataset as part of that map operation.

I want to be able to _filter_ the dataset based on this `num_tokens` column, but even when I specify `input_columns=['num_tokens']`, it seems that the entirety of each row is loaded into memory, which makes the operation take much longer than it should. Indeed, `filter` currently just calls `map`, and I found that in `_map_single` on lines 1690-1704 of `arrow_dataset.py`, the method is just grabbing slices of _all the rows_ of the dataset and then passing only the specified columns to the map function. It seems that, when the user passes a value for `input_columns`, the `map` function should create a temporary pyarrow table by selecting just those columns, and then get slices from that table. Or something like that— I'm not very familiar with the pyarrow API.

I know that in the meantime I can sort of get around this by simply only returning the rows that match my filter criterion from the tokenizing function I pass to `map()`, but I actually _also_ want to map on just the `num_tokens` column in order to compute batches with a roughly uniform number of tokens per batch. I would also ideally like to be able to change my minimum and maximum article lengths without having to re-tokenize the entire dataset.

PS: This is definitely not a "dataset request." I'm realizing that I don't actually know how to remove labels from my own issues on other people's repos, if that is even possible.

Comments:
Hi ! Yes we are working on making `filter` significantly faster. You can look at related PRs here: #2060 #2178 

I think you can expect to have the fast version of `filter` available next week.

We'll make it only select one column, and we'll also make the overall filtering operation way faster by avoiding many arrow<->python conversions especially during writing.

I'll let you know how it goes !

@lhoestq Thanks for the response— it's great to hear that we'll be getting a much faster `filter` method soon. However, my use case does also involve using `map` over a single column in order to pre-compute roughly uniformly sized batches, and right now that is also very slow. Is there any plan to make `map` faster for single column operations?

If that's not a priority for the maintainers right now, I could try my hand at adding the feature, but I can't guarantee I would do a good job given my lack of familiarity with pyarrow.

Currently the optimal setup for single-column computations is probably to do something like
```python
result = dataset.map(f, input_columns="my_col", remove_columns=dataset.column_names)
```
This has two advantages:
- input_columns="my_col" allows to only read the column "my_col"
- remove_columns=dataset.column_names makes `map` only keep the output of your function `f`, and it drops the other columns of the dataset instead of keeping them.

Let me know if it improves speed on your side.

You can also get more speed by using `batched=True` and setting `num_proc=` for multiprocessing

Hi @lhoestq ,

I'm hijacking this issue, because I'm currently trying to do the approach you recommend:

> Currently the optimal setup for single-column computations is probably to do something like
> 
> ```python
> result = dataset.map(f, input_columns="my_col", remove_columns=dataset.column_names)
> ```

Here is my code: (see edit, in which I added a simplified version

```
This is the error:
```bash
pyarrow.lib.ArrowInvalid: Column 1 named tokens expected length 8964 but got length 1000
```
I wonder why this error occurs, when I delete every column? Can you give me a hint?

### Edit:
I preprocessed my dataset before (using map with the features argument) and saved it to disk. May this be part of the error?  I can iterate over the
complete dataset and print every sample before calling map. There seems to be no other problem with the dataset.

I tried to simplify the code that crashes:

```python
# works
log.debug(dataset.column_names)
log.debug(dataset)
for i, sample in enumerate(dataset):
    log.debug(i, sample)

# crashes
counted_dataset = dataset.map(
    lambda x: {"a": list(range(20))},
    input_columns=column,
    remove_columns=dataset.column_names,
    load_from_cache_file=False,
    num_proc=num_workers,
    batched=True,
)
```

```
pyarrow.lib.ArrowInvalid: Column 1 named tokens expected length 20 but got length 1000
```

Edit2: 

May this be a problem with a schema I set when preprocessing the dataset before? I tried to add the `features` argument to the function and then I get a new error:

```python
# crashes
counted_dataset = dataset.map(
    lambda x: {"a": list(range(20))},
    input_columns=column,
    remove_columns=dataset.column_names,
    load_from_cache_file=False,
    num_proc=num_workers,
    batched=True,
    features=datasets.Features(
        {
              "a": datasets.Sequence(datasets.Value("int32"))
         }
    )
)
```

```
 File "env/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1704, in _map_single
    writer.write_batch(batch)
  File "env/lib/python3.8/site-packages/datasets/arrow_writer.py", line 312, in write_batch
    col_type = schema.field(col).type if schema is not None else None
  File "pyarrow/types.pxi", line 1341, in pyarrow.lib.Schema.field
KeyError: 'Column tokens does not exist in schema'
```

Hi ! Can you open a separate issue for that ?
Also if you could provide a google colab or a sample code to reproduce this issue that would be helpful.
On my side I was not able to reproduce this error.

@lhoestq Sorry I'm just responding now. I'm currently using your recommendation for the map on a single column, and I've gotten it to be fast enough to sort of work for my use case by just setting `num_proc=10`, although it's still quite slow. It's clear that it is still loading the entirety of each row into memory and then discarding everything except the selected column, instead of exploiting the columnar data format to only load the selected column.

My code is like this:
```
 self.dataset = self.dataset.sort('num_tokens')
 batch_dataset = self.dataset.map(
	compute_uniform_sized_batches,
	batched=True, batch_size=10_000, num_proc=10, input_columns=['num_tokens'],
	remove_columns=get_columns_all_equal(self.dataset),
	with_indices=True,
	fn_kwargs=dict(max_size=tokens_per_batch)
)
self.batches = {
	name: list(zip(split['start'], split['length']))
	for name, split in batch_dataset.items()
}
```
I find that the processes with higher IDs take significantly longer to complete, presumably because the dataset is sorted by article length and they're loading the entire article text into memory, instead of just the 'num_tokens' column.

I should note that my batching procedure would work best if I just used `batch_size=None` and loaded the whole column into memory at once, but I found that this was intolerably slow and gave me no progress information, so I'm using the less than ideal `batch_size=10_000`.

Hi @norabelrose ! I'm glad you managed to make this work on your side.
Regarding memory usage, you can try to drop the columns that you don't want to use for your `map` for now.

In the future we'll try to find a way to not load unnecessary columns in memory in `map`. Currently the way it works is that it gets the batch as a python dict, then it updates it using the output of your mapping function, and finally it removes columns from `remove_columns`. Therefore for a moment some columns are loaded in memory even if you remove them or don't use them for your mapping function.

It would be nice to have a way to optimize memory for cases such as yours !

@lhoestq After looking through the source code, it looks like the following solution has at least some chance of working:
- refactor `Dataset.map()` so that the `input_columns` parameter is implemented by using the `self.formatted_as()` context manager with `columns=input_columns`
- change `Dataset._getitem()` so that it passes `self._data.drop(drop_columns)` to the `query_table()` function whenever `format_columns` is non-None and `output_all_columns` is False, instead of `self._data` itself

Looks like a great direction :)
Note that `query_table` doesn't bring data into memory. Only `format_table` does.
Also the dataset may already have a format with `columns=` already defined so we would need to define the formatted `input_dataset` like:
```python
# before the `map` main for loop
input_columns = input_columns if input_columns is not None else self.column_names
if not self._output_all_columns:
    columns = [col for col in input_columns if self._format_columns is None or col in self._format_columns]
    input_dataset = self.with_format(
        type=self._format_type,
        columns=columns
    )
else:
    # in this case we could find a way to filter both format_columns and unformatted columns eventually
    input_dataset = self
# then input_dataset can be used in the main for loop of `map`
```

EDIT: oh and regarding streaming format versus file format for arrow, we plan to start using the file format #1933 at one point (though I'm not sure if it would improve performance)

Good to know about `query_table` not bringing anything into memory. I was under the impression that it did because a while back I looked at my `map` operation in pdb and it looked like it was spending forever in line 93 of formatting.py, `return pa.concat_tables(....)`, although that was before the `fast_slice` interpolation search was implemented, so it may have had more to do with the slow ChunkedArray slice implementation than anything else.

If `query_table` is I/O free then the fix may be as simple as just adding this to line 1779 of arrow_dataset.py:
```python
# Only load the columns we actually need
if input_columns:
  stack.enter_context(self.formatted_as(
    self._format_type,
    columns=input_columns,
    output_all_columns=False,
    **self._format_kwargs
  ))
```
It's not clear to me why the `[col for col in input_columns if self._format_columns is None or col in self._format_columns]` check would be necessary— it seems like either `input_columns` should simply temporarily override the `_format_columns` within the `map` operation, or we should throw an error if there are any conflicts. Currently it doesn't look like this case is checked for at all within `map`, but maybe I'm just missing it.

`query_table` simply slices/concatenates parts of the table. The actual data inside the table is not brought in memory.
Also I'm more in favor of declaring `input_dataset = self.with_format(...)` since `formatted_as` may update the dataset fingerprint of `self`, which is not expected when someone runs `map`.

> It's not clear to me why the [col for col in input_columns if self._format_columns is None or col in self._format_columns] check would be necessary— it seems like either input_columns should simply temporarily override the _format_columns within the map operation, or we should throw an error if there are any conflicts. Currently it doesn't look like this case is checked for at all within map, but maybe I'm just missing it.

Actually yes we can just use input_columns. And we do need to add a check to make sure there are not conflicts or this could lead to confusing errors.

That sounds good to me! I just submitted a PR (#2246) implementing your approach. I also changed how `_query_table` handles Iterable keys since it still seemed like `pa.concat_tables` was taking a long time to create the table for each batch. Now my whole `map()` operation takes 1 min 46 seconds where it used to take somewhere on the order of 10 minutes.