feat: parallel map reduce (#1102)

mandel · web-flow · commit 218342efe1bd · 2025-08-08T14:31:10.000-04:00
Signed-off-by: Louis Mandel &lt;lmandel@us.ibm.com&gt;
diff --git a/pdl-live-react/src/pdl_ast.d.ts b/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
 }
 /**
diff --git a/src/pdl/pdl-schema.json b/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": [
             {
diff --git a/src/pdl/pdl_ast.py b/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
 
diff --git a/src/pdl/pdl_interpreter.py b/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:
diff --git a/tests/test_for.py b/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"

Original file line number	Diff line number	Diff line change
`@@ -2491,6 +2491,11 @@ export interface MapBlock {`
`2491`	`2491`	`*`
`2492`	`2492`	`*/`
`2493`	`2493`	`join?: JoinText \| JoinArray \| JoinObject \| JoinLastOf \| JoinReduce`
	`2494`	`+ /**`
	`2495`	+ * 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`.
	`2496`	`+ *`
	`2497`	`+ */`
	`2498`	`+ maxWorkers?: number \| null`
`2494`	`2499`	`pdl__trace?: PdlTrace`
`2495`	`2500`	`}`
`2496`	`2501`	`/**`