Beam datasets
State: closed | Author: lhoestq (MEMBER)

# Beam datasets

## Intro

Beam Datasets are using beam pipelines for preprocessing (basically lots of `.map` over objects called PCollections).
The advantage of apache beam is that you can choose which type of runner you want to use to preprocess your data. The main runners are:
- the `DirectRunner` to run the pipeline locally (default). However I encountered memory issues for big datasets (like the french or english wikipedia). Small dataset work fine
- Google Dataflow. I didn't play with it.
- Spark or Flink, two well known data processing frameworks. I tried to use the Spark/Flink local runners provided by apache beam for python and wasn't able to make them work properly though...

## From tfds beam datasets to our own beam datasets

Tensorflow datasets used beam and a complicated pipeline to shard the TFRecords files.
To allow users to download beam datasets and not having to preprocess them, they also allow to download the already preprocessed datasets from their google storage (the beam pipeline doesn't run in that case).

On our side, we replace TFRecords by something else. Arrow or Parquet do the job but I chose Parquet as: 1) there is a builtin apache beam parquet writer that is quite convenient, and 2) reading parquet from the pyarrow library is also simple and effective (there is a mmap option !)

Moreover we don't shard datasets in many many files like tfds (they were doing probably doing that mainly because of the limit of 2Gb per TFRecord file). Therefore we have a simpler pipeline that saves each split into one parquet file. We also removed the utilities to use their google storage (for now maybe ? we'll have to discuss it).

## Main changes

- Added a BeamWriter to save the output of beam pipelines into parquet files and fill dataset infos
- Create a ParquetReader and refactor a bit the arrow_reader.py

\> **With this, we can now try to add beam datasets from tfds**

I already added the wikipedia one, and I will also try to add the Wiki40b dataset

## Test the wikipedia script

You can download and run the beam pipeline for wikipedia (using the `DirectRunner` by default) like this:

```
>>> import nlp
>>> nlp.load("datasets/nlp/wikipedia", dataset_config="20200501.frr")
```

This wikipedia dataset (lang: frr, North Frisian) is a small one (~10Mb), but feel free to try bigger ones (and fill 20Gb of swap memory if you try the english one lol)

## Next

Should we allow to download preprocessed datasets from the tfds google storage ?
Should we try to optimize the beam pipelines to run locally without memory issues ?
Should we try other data processing frameworks for big datasets, like spark ?


## About this PR

It should be merged after #25 

-----------------

I'd be happy to have your feedback and your ideas to improve the processing of big datasets like wikipedia :)

Comments: Right now the changes are a bit hard to read as the one from #25 are also included. You can wait until #25 is merged before looking at the implementation details

Nice!! I tested it a bit and works quite well. I will do a my review once the #25 will be merged because there are several overlaps.

At least I can share my thoughts on your **Next** section:
1) I don't think it is a good thing to rely on tfds preprocessed datasets uploaded in their online storage, because they might be updated or deleted at any moment by Google and then possibly break our own processing.
2) Improves the pipeline is always a good direction, but in the meantime we might also share the preprocessed dataset in S3 storage. Which might be another way to see 1), instead of downloading Google preprocessed datasets, using our own ones.
3) Apache Beam can be easily integrated in Spark, so I don't see the need to replace Beam by Spark.

Ok I've merged #25 so you can rebase or merge if you want.

I fully agree with @jplu notes for the "next section".

Don't hesitate to use some credit on Google Dataflow if you think it would be useful to give it a try.

Pr is ready for review !

New minor changes:
- re-added the csv dataset builder (it was on my branch from #25 but disappeared from master)
- move the csv script and the wikipedia script to "under construction" for now
- some renaming in the `nlp-cli test` command