Only use thread for a batch of rows not each row

Author: kesmit13

singlestoredb/config.py

@@ -444,6 +444,13 @@
     environ=['SINGLESTOREDB_EXT_FUNC_TIMEOUT'],
 )
 
+register_option(
+    'external_function.concurrency_limit', 'int', check_int, 1,
+    'Specifies the maximum number of subsets of a batch of rows '
+    'to process simultaneously.',
+    environ=['SINGLESTOREDB_EXT_FUNC_CONCURRENCY_LIMIT'],
+)
+
 #
 # Debugging options
 #

singlestoredb/functions/ext/asgi.py

@@ -53,6 +53,7 @@
 from typing import Callable
 from typing import Dict
 from typing import Iterable
+from typing import Iterator
 from typing import List
 from typing import Optional
 from typing import Sequence
@@ -286,11 +287,21 @@ def identity(x: Any) -> Any:
     return x
 
 
+def chunked(seq: Sequence[Any], max_chunks: int) -> Iterator[Sequence[Any]]:
+    """Yield up to max_chunks chunks from seq, splitting as evenly as possible."""
+    n = len(seq)
+    if max_chunks <= 0 or max_chunks > n:
+        max_chunks = n
+    chunk_size = (n + max_chunks - 1) // max_chunks  # ceil division
+    for i in range(0, n, chunk_size):
+        yield seq[i:i + chunk_size]
+
+
 async def run_in_parallel(
     func: Callable[..., Any],
     params_list: Sequence[Sequence[Any]],
     cancel_event: threading.Event,
-    limit: int = 10,
+    limit: int = get_option('external_function.concurrency_limit'),
     transformer: Callable[[Any], Any] = identity,
 ) -> List[Any]:
     """"
@@ -306,26 +317,39 @@ async def run_in_parallel(
         The event to check for cancellation
     limit : int
         The maximum number of concurrent tasks to run
+    transformer : Callable[[Any], Any]
+        A function to transform the results
 
     Returns
     -------
     List[Any]
         The results of the function calls
 
     """
-    semaphore = asyncio.Semaphore(limit)
+    is_async = asyncio.iscoroutinefunction(func)
 
-    async def worker(params: Sequence[Any]) -> Any:
-        async with semaphore:
+    async def call(batch: Sequence[Any]) -> Any:
+        """Loop over batches of parameters and call the function."""
+        res = []
+        for params in batch:
             cancel_on_event(cancel_event)
-            if asyncio.iscoroutinefunction(func):
-                return transformer(await func(*params))
+            if is_async:
+                res.append(transformer(await func(*params)))
             else:
-                return transformer(await to_thread(func, *params))
+                res.append(transformer(func(*params)))
+        return res
+
+    async def thread_call(batch: Sequence[Any]) -> Any:
+        if is_async:
+            return await call(batch)
+        return await to_thread(lambda: asyncio.run(call(batch)))
+
+    # Create tasks in chunks to limit concurrency
+    tasks = [thread_call(batch) for batch in chunked(params_list, limit)]
 
-    tasks = [worker(p) for p in params_list]
+    results = await asyncio.gather(*tasks)
 
-    return await asyncio.gather(*tasks)
+    return list(itertools.chain.from_iterable(results))
 
 
 def build_udf_endpoint(