datasets

1
import pytest
2

3
from datasets.parallel import ParallelBackendConfig, parallel_backend
4
from datasets.utils.py_utils import map_nested
5

6
from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows
7

8

9
def add_one(i):  # picklable for multiprocessing
10
    return i + 1
11

12

13
@require_dill_gt_0_3_2
14
@require_joblibspark
15
@require_not_windows
16
def test_parallel_backend_input():
17
    with parallel_backend("spark"):
18
        assert ParallelBackendConfig.backend_name == "spark"
19

20
    lst = [1, 2, 3]
21
    with pytest.raises(ValueError):
22
        with parallel_backend("unsupported backend"):
23
            map_nested(add_one, lst, num_proc=2)
24

25
    with pytest.raises(ValueError):
26
        with parallel_backend("unsupported backend"):
27
            map_nested(add_one, lst, num_proc=-1)
28

29

30
@require_dill_gt_0_3_2
31
@require_joblibspark
32
@require_not_windows
33
@pytest.mark.parametrize("num_proc", [2, -1])
34
def test_parallel_backend_map_nested(num_proc):
35
    s1 = [1, 2]
36
    s2 = {"a": 1, "b": 2}
37
    s3 = {"a": [1, 2], "b": [3, 4]}
38
    s4 = {"a": {"1": 1}, "b": 2}
39
    s5 = {"a": 1, "b": 2, "c": 3, "d": 4}
40
    expected_map_nested_s1 = [2, 3]
41
    expected_map_nested_s2 = {"a": 2, "b": 3}
42
    expected_map_nested_s3 = {"a": [2, 3], "b": [4, 5]}
43
    expected_map_nested_s4 = {"a": {"1": 2}, "b": 3}
44
    expected_map_nested_s5 = {"a": 2, "b": 3, "c": 4, "d": 5}
45

46
    with parallel_backend("spark"):
47
        assert map_nested(add_one, s1, num_proc=num_proc) == expected_map_nested_s1
48
        assert map_nested(add_one, s2, num_proc=num_proc) == expected_map_nested_s2
49
        assert map_nested(add_one, s3, num_proc=num_proc) == expected_map_nested_s3
50
        assert map_nested(add_one, s4, num_proc=num_proc) == expected_map_nested_s4
51
        assert map_nested(add_one, s5, num_proc=num_proc) == expected_map_nested_s5
52