feat: parallel map reduce (#1102)

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

pdl-live-react/src/pdl_ast.d.ts

@@ -2491,6 +2491,11 @@ export interface MapBlock {
    *
    */
   join?: JoinText | JoinArray | JoinObject | JoinLastOf | JoinReduce
+  /**
+   * Maximal number of workers to execute the map in parallel. Is it is set to `0`, the execution is sequential otherwise it is given as argument to the `ThreadPoolExecutor`.
+   *
+   */
+  maxWorkers?: number | null
   pdl__trace?: PdlTrace
 }
 /**

src/pdl/pdl-schema.json

@@ -2944,6 +2944,11 @@
           },
           "description": "Define how to combine the result of each iteration.\n    "
         },
+        "maxWorkers": {
+          "$ref": "#/$defs/OptionalInt",
+          "default": null,
+          "description": "Maximal number of workers to execute the map in parallel. Is it is set to `0`, the execution is sequential otherwise it is given as argument to the `ThreadPoolExecutor`.\n    "
+        },
         "pdl__trace": {
           "anyOf": [
             {

src/pdl/pdl_ast.py

@@ -1022,6 +1022,9 @@ class MapBlock(StructuredBlock):
     join: JoinType = JoinText()
     """Define how to combine the result of each iteration.
     """
+    maxWorkers: OptionalInt = None
+    """Maximal number of workers to execute the map in parallel. Is it is set to `0`, the execution is sequential otherwise it is given as argument to the `ThreadPoolExecutor`.
+    """
     # Field for internal use
     pdl__trace: Optional[list["BlockType"]] = None
 

src/pdl/pdl_interpreter.py

@@ -11,6 +11,7 @@
 # TODO: temporarily disabling warnings to mute a pydantic warning from liteLLM
 import warnings
 from asyncio import AbstractEventLoop
+from concurrent.futures import ThreadPoolExecutor
 from functools import partial, reduce
 from itertools import count
 from os import getenv
@@ -986,11 +987,18 @@ def loop_body(iidx, items):
                         map_loc,
                     )
 
-                # with ThreadPoolExecutor(max_workers=4) as executor:
-                #     map_output = executor.map(
-                map_output = map(  # pylint: disable=bad-builtin
-                    loop_body, index_iterator, items_iterator
-                )
+                map_output: Iterable[
+                    Tuple[PdlLazy[Any], LazyMessages, ScopeType, BlockType]
+                ]
+                if block.maxWorkers == 0:
+                    map_output = map(  # pylint: disable=bad-builtin
+                        loop_body, index_iterator, items_iterator
+                    )
+                else:
+                    with ThreadPoolExecutor(block.maxWorkers) as executor:
+                        map_output = executor.map(
+                            loop_body, index_iterator, items_iterator
+                        )
                 results, _, _, traces = _split_map_output(map_output)
                 # saved_background = IndependentContext(backgrounds)
             except PDLRuntimeError as exc:

tests/test_for.py

@@ -324,6 +324,25 @@ def test_for_context():
 def test_for_reduce():
     for loop_kind1 in ["repeat", "map"]:
         prog_str = f"""
+defs:
+  plus:
+    function:
+        x: number
+        y: number
+    return: ${{ x + y }}
+for:
+  i: [1,2,3,4]
+{loop_kind1}: ${{ i }}
+join:
+  reduce: ${{ plus }}
+"""
+        result = exec_str(prog_str)
+        assert result == 10
+
+
+def test_for_reduce_python():
+    for loop_kind1 in ["repeat", "map"]:
+        prog_str = f"""
 defs:
   plus:
     lang: python
@@ -338,3 +357,15 @@ def test_for_reduce():
 """
         result = exec_str(prog_str)
         assert result == 10
+
+
+def test_map_max_workers():
+    for max_workers in [0, "null", 2, 4]:
+        prog_str = f"""
+for:
+  i: [1,2,3,4]
+map: ${{ i }}
+maxWorkers: {max_workers}
+"""
+        result = exec_str(prog_str)
+        assert result == "1234"