Repository: dask/dask

Issue:
Optionally batch reading large number of small files
Dask's `bag.read_text()` function performs poorly when given a large number of small files (>1000 * ~100kb), because it constructs a massive graph of operations that takes significant time (10-50s) to serialize. Reading chunks of files in batches would dramatically reduce the size of the graph, cutting overhead both in time and memory. This could be an optional parameter to `read_text()`.


Discussion:
Makes sense to me.  Something like a `files_per_partition=None` parameter?
If you are interested in implementing this that would be welcome.


On Tue, Sep 25, 2018 at 1:48 PM Eric Wolak <notifications@github.com> wrote:

> Dask's bag.read_text() function performs poorly when given a large number
> of small files (>1000 * ~100kb), because it constructs a massive graph of
> operations that takes significant time (10-50s) to serialize. Reading
> chunks of files in batches would dramatically reduce the size of the graph,
> cutting overhead both in time and memory. This could be an optional
> parameter to read_text().
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4012>, or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AASszAH-AKz8OuNQS-Jv-7vgX1Bn8zQMks5uemx-gaJpZM4W5KnJ>
> .
>

Yep. Got a prototype working now. PR incoming.

Resolving patch:
diff --git a/dask/bag/text.py b/dask/bag/text.py
--- a/dask/bag/text.py
+++ b/dask/bag/text.py
@@ -16,7 +16,7 @@
 def read_text(urlpath, blocksize=None, compression='infer',
               encoding=system_encoding, errors='strict',
               linedelimiter=os.linesep, collection=True,
-              storage_options=None):
+              storage_options=None, files_per_partition=None):
     """ Read lines from text files
 
     Parameters
@@ -38,6 +38,9 @@ def read_text(urlpath, blocksize=None, compression='infer',
     storage_options: dict
         Extra options that make sense to a particular storage connection, e.g.
         host, port, username, password, etc.
+    files_per_partition: None or int
+        If set, group input files into partitions of the requested size,
+        instead of one partition per file. Mutually exclusive with blocksize.
 
     Examples
     --------
@@ -61,12 +64,21 @@ def read_text(urlpath, blocksize=None, compression='infer',
     --------
     from_sequence: Build bag from Python sequence
     """
-    if blocksize is None:
-        files = open_files(urlpath, mode='rt', encoding=encoding,
-                           errors=errors, compression=compression,
-                           **(storage_options or {}))
-        blocks = [delayed(list)(delayed(file_to_blocks)(fil)) for fil in files]
+    if blocksize is not None and files_per_partition is not None:
+        raise ValueError('Only one of blocksize or files_per_partition can be set')
 
+    files = open_files(urlpath, mode='rt', encoding=encoding,
+                       errors=errors, compression=compression,
+                       **(storage_options or {}))
+    if blocksize is None:
+        if files_per_partition is None:
+            blocks = [delayed(list)(delayed(file_to_blocks)(fil)) for fil in files]
+        else:
+            blocks = []
+            for start in range(0, len(files), files_per_partition):
+                block_files = files[start:(start + files_per_partition)]
+                block_lines = delayed(concat)(delayed(map)(file_to_blocks, block_files))
+                blocks.append(block_lines)
     else:
         _, blocks = read_bytes(urlpath, delimiter=linedelimiter.encode(),
                                blocksize=blocksize, sample=False,


Test patch:
diff --git a/dask/bag/tests/test_text.py b/dask/bag/tests/test_text.py
--- a/dask/bag/tests/test_text.py
+++ b/dask/bag/tests/test_text.py
@@ -52,6 +52,17 @@ def test_read_text(fmt, bs, encoding):
         assert ''.join(line for block in L for line in block) == expected
 
 
+def test_files_per_partition():
+    files3 = {'{:02}.txt'.format(n): 'line from {:02}' for n in range(20)}
+    with filetexts(files3):
+        b = read_text('*.txt', files_per_partition=10)
+
+        l = len(b.take(100, npartitions=1))
+        assert l == 10, "10 files should be grouped into one partition"
+
+        assert b.count().compute() == 20, "All 20 lines should be read"
+
+
 def test_errors():
     with filetexts({'.test.foo': b'Jos\xe9\nAlice'}, mode='b'):
         with pytest.raises(UnicodeDecodeError):