Add multiprocessing (#552)

* add multiprocessing

* test multiprocessing and fix empty tables issues

* improve logging

* minor

* fix unit

* fix tqdm in notebook for multiprocessing

Author: lhoestq

src/nlp/__init__.py

@@ -24,7 +24,7 @@
 from pyarrow import total_allocated_bytes
 
 from . import datasets
-from .arrow_dataset import Dataset
+from .arrow_dataset import Dataset, concatenate_datasets
 from .arrow_reader import ArrowReader, ReadInstruction
 from .arrow_writer import ArrowWriter
 from .builder import ArrowBasedBuilder, BeamBasedBuilder, BuilderConfig, DatasetBuilder, GeneratorBasedBuilder
@@ -43,7 +43,7 @@
 )
 from .info import DatasetInfo, MetricInfo
 from .inspect import inspect_dataset, inspect_metric, list_datasets, list_metrics
-from .load import concatenate_datasets, import_main_class, load_dataset, load_metric, prepare_module
+from .load import import_main_class, load_dataset, load_metric, prepare_module
 from .metric import Metric
 from .splits import NamedSplit, Split, SplitBase, SplitDict, SplitGenerator, SplitInfo, SubSplitInfo, percent
 from .utils import *

src/nlp/arrow_dataset.py

@@ -159,12 +159,13 @@ def _read_files(self, files) -> pa.Table:
                 skip/take indicates which example read in the file: `ds.slice(skip, take)`
         """
         assert len(files) > 0 and all(isinstance(f, dict) for f in files), "please provide valid file informations"
-        pa_batches = []
+        pa_tables = []
         for f_dict in files:
             pa_table: pa.Table = self._get_dataset_from_filename(f_dict)
-            pa_batches.extend(pa_table.to_batches())
-        assert len(pa_batches) > 0, "tried to read an empty arrow table"
-        pa_table = pa.Table.from_batches(pa_batches)
+            pa_tables.append(pa_table)
+        pa_tables = [t for t in pa_tables if len(t) > 0]
+        pa_tables = pa_tables or [pa.Table.from_batches([], schema=pa.schema(self._info.features.type))]
+        pa_table = pa.concat_tables(pa_tables)
         return pa_table
 
     def get_file_instructions(self, name, instruction, split_infos):

src/nlp/arrow_reader.py

@@ -134,6 +134,7 @@ def __init__(
         disable_nullable: bool = False,
         update_features: bool = False,
         with_metadata: bool = True,
+        unit: str = "examples",
     ):
         if path is None and stream is None:
             raise ValueError("At least one of path and stream must be provided.")
@@ -161,6 +162,7 @@ def __init__(
         self.writer_batch_size = writer_batch_size or DEFAULT_MAX_BATCH_SIZE
         self.update_features = update_features
         self.with_metadata = with_metadata
+        self.unit = unit
 
         self._num_examples = 0
         self._num_bytes = 0
@@ -290,8 +292,9 @@ def finalize(self, close_stream=True):
         if close_stream:
             self.stream.close()
         logger.info(
-            "Done writing %s examples in %s bytes %s.",
+            "Done writing %s %s in %s bytes %s.",
             self._num_examples,
+            self.unit,
             self._num_bytes,
             self._path if self._path else "",
         )

src/nlp/arrow_writer.py

@@ -123,8 +123,11 @@ def _fingerprint(func):
 
         @wraps(func)
         def wrapper(*args, **kwargs):
-            self: "Dataset" = args[0]
-            args = args[1:]
+            if args:
+                self: "Dataset" = args[0]
+                args = args[1:]
+            else:
+                self: "Dataset" = kwargs.pop("self")
             kwargs_for_fingerprint = dict(kwargs)
             kwargs_for_fingerprint.update(zip(func.__code__.co_varnames, args))
 
@@ -145,7 +148,7 @@ def wrapper(*args, **kwargs):
                 new_inplace_history_item = (func.__name__, deepcopy(args), deepcopy(kwargs))
             else:
                 for fingerprint_name in fingerprint_names:  # transforms like `train_test_split` have several hashes
-                    if fingerprint_name not in kwargs:
+                    if kwargs.get(fingerprint_name) is None:
                         kwargs_for_fingerprint["fingerprint_name"] = fingerprint_name
                         kwargs[fingerprint_name] = update_fingerprint(self._fingerprint, func, kwargs_for_fingerprint)
 

src/nlp/fingerprint.py

@@ -24,19 +24,16 @@
 import shutil
 from hashlib import sha256
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
-import numpy as np
-import pyarrow as pa
 from filelock import FileLock
 
 from .arrow_dataset import Dataset
 from .builder import DatasetBuilder
 from .dataset_dict import DatasetDict
 from .features import Features
-from .fingerprint import update_fingerprint
-from .info import DATASET_INFOS_DICT_FILE_NAME, DatasetInfo
+from .info import DATASET_INFOS_DICT_FILE_NAME
 from .metric import Metric
 from .splits import Split
 from .utils.download_manager import GenerateMode
@@ -560,110 +557,3 @@ def load_dataset(
         builder_instance._save_infos()
 
     return ds
-
-
-def concatenate_datasets(
-    dsets: List["Dataset"],
-    info: Optional[Any] = None,
-    split: Optional[Any] = None,
-):
-    """
-    Converts a list of :obj:``nlp.Dataset`` with the same schema into a single :obj:``nlp.Dataset``.
-
-    Args:
-        dsets (:obj:``List[nlp.Dataset]``): A list of Datasets to concatenate
-        info (:obj:``nlp.DatasetInfo``, `optional`, defaults to :obj:``None``): If specified, the dataset info containing info like
-            description, citation, etc.
-        split (:obj:``nlp.NamedSplit``, `optional`, defaults to :obj:``None``): If specified, the name of the dataset split.
-    """
-    if not all([dset.features.type == dsets[0].features.type for dset in dsets]):
-        raise ValueError("Features must match for all datasets")
-
-    # Datasets tables should all come from disk or memory, but not a mix
-
-    dsets_in_memory = [not dset._data_files for dset in dsets]
-    if any(dset_in_memory != dsets_in_memory[0] for dset_in_memory in dsets_in_memory):
-        raise ValueError(
-            "Datasets should ALL come from memory, or should ALL come from disk.\n"
-            "However datasets {} come from memory and datasets {} come from disk.".format(
-                [i for i in range(len(dsets)) if dsets_in_memory[i]],
-                [i for i in range(len(dsets)) if not dsets_in_memory[i]],
-            )
-        )
-
-    # Concatenate tables
-
-    table = pa.concat_tables([dset._data for dset in dsets])
-    data_files = [f for dset in dsets for f in dset._data_files]
-    inplace_history = [h for dset in dsets for h in dset._inplace_history]
-
-    def apply_offset_to_indices_table(table, offset):
-        if offset == 0:
-            return table
-        else:
-            array = table["indices"]
-            if isinstance(array, pa.ChunkedArray):
-                new_array = pa.array(np.concatenate([c.to_numpy() for c in array.chunks]) + offset, pa.uint64())
-            else:
-                new_array = pa.array(array.to_numpy() + offset, pa.uint64())
-            return pa.Table.from_arrays([new_array], names=["indices"])
-
-    # Concatenate indices if they exist
-
-    if any(dset._indices is not None for dset in dsets):
-
-        # Datasets indices tables should all come from disk or memory, but not a mix
-        # Datasets with no indices tables are replaced with a dataset with an indicies table in memory
-
-        indices_mappings_in_memory = [not dset._indices_data_files for dset in dsets]
-        if any(
-            indices_mapping_in_memory != indices_mappings_in_memory[0]
-            for indices_mapping_in_memory in indices_mappings_in_memory
-        ):
-            raise ValueError(
-                "Datasets' indices should ALL come from memory, or should ALL come from disk.\n"
-                "However datasets' indices {} come from memory and datasets' indices {} come from disk.".format(
-                    [i for i in range(len(dsets)) if indices_mappings_in_memory[i]],
-                    [i for i in range(len(dsets)) if not indices_mappings_in_memory[i]],
-                )
-            )
-        indices_in_memory = indices_mappings_in_memory[0]
-
-        # Create missing indices tables in memory
-
-        if indices_in_memory:
-            for i in range(len(dsets)):
-                if dsets[i]._indices is None:
-                    dsets[i] = dsets[i].select(range(len(dsets[i])))
-        assert all(dset._indices is not None for dset in dsets), "each dataset should have an indices table"
-
-        # An offset needs to be applied to the indices before concatenating
-
-        indices_tables = []
-        offset = 0
-        for dset in dsets:
-            indices_tables.append(apply_offset_to_indices_table(dset._indices, offset))
-            offset += len(dset._data)
-
-        # Concatenate indices
-
-        indices_table = pa.concat_tables(indices_tables)
-        indices_data_files = None if indices_in_memory else [f for dset in dsets for f in dset._indices_data_files]
-    else:
-        indices_table = None
-        indices_data_files = None
-    if info is None:
-        info = DatasetInfo.from_merge([dset.info for dset in dsets])
-    fingerprint = update_fingerprint(
-        "".join(dset._fingerprint for dset in dsets), concatenate_datasets, {"info": info, "split": split}
-    )
-    return Dataset(
-        table,
-        info=info,
-        split=split,
-        data_files=data_files,
-        indices_table=indices_table,
-        indices_data_files=indices_data_files,
-        fingerprint=fingerprint,
-        inplace_history=inplace_history,
-    )

src/nlp/load.py

@@ -20,6 +20,14 @@ def __getstate__(self):
         raise pickle.PicklingError()
 
 
+def picklable_map_function(x):
+    return {"id": int(x["filename"].split("_")[-1])}
+
+
+def picklable_filter_function(x):
+    return int(x["filename"].split("_")[-1]) < 10
+
+
 class BaseDatasetTest(TestCase):
     def _create_dummy_dataset(self, multiple_columns=False):
         if multiple_columns:
@@ -493,6 +501,34 @@ def func(x, i):
                 Features({"filename": Value("string"), "name": Value("string"), "id": Value("int64")}),
             )
 
+    def test_map_multiprocessing(self):
+        dset = self._create_dummy_dataset()
+
+        self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
+        fingerprint = dset._fingerprint
+        dset_test = dset.map(picklable_map_function, num_proc=2)
+        self.assertEqual(len(dset_test), 30)
+        self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
+        self.assertDictEqual(
+            dset_test.features,
+            Features({"filename": Value("string"), "id": Value("int64")}),
+        )
+        self.assertEqual(len(dset_test._data_files), 0)
+        self.assertNotEqual(dset_test._fingerprint, fingerprint)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            dset = dset.map(cache_file_name=os.path.join(tmp_dir, "test.arrow"))
+            fingerprint = dset._fingerprint
+            dset_test = dset.map(picklable_map_function, num_proc=3)
+            self.assertEqual(len(dset_test), 30)
+            self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
+            self.assertDictEqual(
+                dset_test.features,
+                Features({"filename": Value("string"), "id": Value("int64")}),
+            )
+            self.assertEqual(len(dset_test._data_files), 3)
+            self.assertNotEqual(dset_test._fingerprint, fingerprint)
+
     def test_new_features(self):
         dset = self._create_dummy_dataset()
 
@@ -642,6 +678,27 @@ def test_filter(self):
             self.assertDictEqual(dset_filter_even_num.features, Features({"filename": Value("string")}))
             self.assertNotEqual(dset_filter_even_num._fingerprint, fingerprint)
 
+    def test_filter_multiprocessing(self):
+        dset = self._create_dummy_dataset()
+
+        fingerprint = dset._fingerprint
+        dset_filter_first_ten = dset.filter(picklable_filter_function, num_proc=2)
+        self.assertEqual(len(dset_filter_first_ten), 10)
+        self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
+        self.assertDictEqual(dset_filter_first_ten.features, Features({"filename": Value("string")}))
+        self.assertEqual(len(dset_filter_first_ten._data_files), 0)
+        self.assertNotEqual(dset_filter_first_ten._fingerprint, fingerprint)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            dset = dset.map(cache_file_name=os.path.join(tmp_dir, "test.arrow"))
+            fingerprint = dset._fingerprint
+            dset_filter_first_ten = dset.filter(picklable_filter_function, num_proc=2)
+            self.assertEqual(len(dset_filter_first_ten), 10)
+            self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
+            self.assertDictEqual(dset_filter_first_ten.features, Features({"filename": Value("string")}))
+            self.assertEqual(len(dset_filter_first_ten._data_files), 2)
+            self.assertNotEqual(dset_filter_first_ten._fingerprint, fingerprint)
+
     def test_keep_features_after_transform_specified(self):
         features = Features(
             {"tokens": Sequence(Value("string")), "labels": Sequence(ClassLabel(names=["negative", "positive"]))}