Fix streaming tar files from canonical datasets
State: closed | Author: albertvillanova (MEMBER)

Previous PR #2800 implemented support to stream remote tar files when passing the parameter `data_files`: they required a glob string `"*"`.

However, this glob string creates an error when streaming canonical datasets (with a `join` after the `open`).

This PR fixes this issue and allows streaming tar files both from:
- canonical datasets scripts and
- data files.

This PR also adds support for compressed tar files: `.tar.gz`, `.tar.bz2`,...


Comments: In case it's relevant for this PR, I'm finding that I cannot stream the `bookcorpus` dataset (using the `master` branch of `datasets`), which is a `.tar.bz2` file:

```python
from datasets import load_dataset

books_dataset_streamed = load_dataset("bookcorpus", split="train", streaming=True)
# Throws a 404 HTTP error
next(iter(books_dataset_streamed))
```

The full stack trace is:

```
---------------------------------------------------------------------------
ClientResponseError                       Traceback (most recent call last)
<ipython-input-11-5ebbbe110b13> in <module>()
----> 1 next(iter(books_dataset_streamed))

11 frames
/usr/local/lib/python3.7/dist-packages/datasets/iterable_dataset.py in __iter__(self)
    339 
    340     def __iter__(self):
--> 341         for key, example in self._iter():
    342             if self.features:
    343                 # we encode the example for ClassLabel feature types for example

/usr/local/lib/python3.7/dist-packages/datasets/iterable_dataset.py in _iter(self)
    336         else:
    337             ex_iterable = self._ex_iterable
--> 338         yield from ex_iterable
    339 
    340     def __iter__(self):

/usr/local/lib/python3.7/dist-packages/datasets/iterable_dataset.py in __iter__(self)
     76 
     77     def __iter__(self):
---> 78         for key, example in self.generate_examples_fn(**self.kwargs):
     79             yield key, example
     80 

/root/.cache/huggingface/modules/datasets_modules/datasets/bookcorpus/44662c4a114441c35200992bea923b170e6f13f2f0beb7c14e43759cec498700/bookcorpus.py in _generate_examples(self, directory)
     98         for txt_file in files:
     99             with open(txt_file, mode="r", encoding="utf-8") as f:
--> 100                 for line in f:
    101                     yield _id, {"text": line.strip()}
    102                     _id += 1

/usr/local/lib/python3.7/dist-packages/fsspec/implementations/http.py in read(self, length)
    496         else:
    497             length = min(self.size - self.loc, length)
--> 498         return super().read(length)
    499 
    500     async def async_fetch_all(self):

/usr/local/lib/python3.7/dist-packages/fsspec/spec.py in read(self, length)
   1481             # don't even bother calling fetch
   1482             return b""
-> 1483         out = self.cache._fetch(self.loc, self.loc + length)
   1484         self.loc += len(out)
   1485         return out

/usr/local/lib/python3.7/dist-packages/fsspec/caching.py in _fetch(self, start, end)
    374         ):
    375             # First read, or extending both before and after
--> 376             self.cache = self.fetcher(start, bend)
    377             self.start = start
    378         elif start < self.start:

/usr/local/lib/python3.7/dist-packages/fsspec/asyn.py in wrapper(*args, **kwargs)
     86     def wrapper(*args, **kwargs):
     87         self = obj or args[0]
---> 88         return sync(self.loop, func, *args, **kwargs)
     89 
     90     return wrapper

/usr/local/lib/python3.7/dist-packages/fsspec/asyn.py in sync(loop, func, timeout, *args, **kwargs)
     67         raise FSTimeoutError
     68     if isinstance(result[0], BaseException):
---> 69         raise result[0]
     70     return result[0]
     71 

/usr/local/lib/python3.7/dist-packages/fsspec/asyn.py in _runner(event, coro, result, timeout)
     23         coro = asyncio.wait_for(coro, timeout=timeout)
     24     try:
---> 25         result[0] = await coro
     26     except Exception as ex:
     27         result[0] = ex

/usr/local/lib/python3.7/dist-packages/fsspec/implementations/http.py in async_fetch_range(self, start, end)
    535                 # range request outside file
    536                 return b""
--> 537             r.raise_for_status()
    538             if r.status == 206:
    539                 # partial content, as expected

/usr/local/lib/python3.7/dist-packages/aiohttp/client_reqrep.py in raise_for_status(self)
   1003                 status=self.status,
   1004                 message=self.reason,
-> 1005                 headers=self.headers,
   1006             )
   1007 

ClientResponseError: 404, message='Not Found', url=URL('https://storage.googleapis.com/huggingface-nlp/datasets/bookcorpus/bookcorpus.tar.bz2/books_large_p1.txt')
```

Let me know if this is unrelated and I'll open a separate issue :)

Environment info:

```
- `datasets` version: 1.11.1.dev0
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.11
- PyArrow version: 3.0.0
```

@lewtun: `.tar.compression-extension` files are not supported yet. That is the objective of this PR.

> @lewtun: `.tar.compression-extension` files are not supported yet. That is the objective of this PR.

thanks for the context and the great work on the streaming features (right now i'm writing the streaming section of the HF course, so am acting like a beta tester 😄)

@lewtun this PR fixes previous issue with xjoin:

Given:
```python
xjoin(
    "https://storage.googleapis.com/huggingface-nlp/datasets/bookcorpus/bookcorpus.tar.bz2",
    "books_large_p1.txt"
)
```

- Before it gave: 
  `"https://storage.googleapis.com/huggingface-nlp/datasets/bookcorpus/bookcorpus.tar.bz2/books_large_p1.txt"`
  thus raising the 404 error

- Now it gives:
  `tar://books_large_p1.txt::https://storage.googleapis.com/huggingface-nlp/datasets/bookcorpus/bookcorpus.tar.bz2`
  (this is the expected format for `fsspec`) and additionally passes the parameter `compression="bz2"`.
  See: https://github.com/huggingface/datasets/pull/2806/files#diff-97bb2d08db65ce3b679aefc43cadad76d053c1e58ecc315e49b80873d0fbdabeR15

closing in favor of #3066