zhengruifeng commented on code in PR #53769:
URL: https://github.com/apache/spark/pull/53769#discussion_r2692720774
##########
python/pyspark/worker.py:
##########
@@ -110,6 +112,58 @@
has_memory_profiler = False
+class ArrowBatchedUDFThreadPool:
+ """
+ Singleton ThreadPoolExecutor for Arrow UDF concurrent execution.
+
+ This class is NOT thread-safe. Its methods (get_pool, shutdown,
map_chunked) should
+ only be called sequentially from a single thread. The worker processes
batches
+ sequentially, so concurrent access to this class is not expected.
+ """
+
+ _instance: ClassVar[Optional[ThreadPoolExecutor]] = None
+ _max_workers: ClassVar[int] = 0
+
+ @classmethod
+ def get_pool(cls, max_workers: int) -> ThreadPoolExecutor:
+ """Get or create the singleton thread pool with the specified number
of workers."""
+ if cls._instance is None or cls._max_workers != max_workers:
+ if cls._instance is not None:
+ cls._instance.shutdown(wait=False)
+ cls._max_workers = max_workers
+ cls._instance = ThreadPoolExecutor(max_workers=max_workers)
+ return cls._instance
+
+ @classmethod
+ def shutdown(cls, wait: bool = True) -> None:
+ """Shutdown the thread pool."""
+ if cls._instance is not None:
+ cls._instance.shutdown(wait=wait)
+ cls._instance = None
+ cls._max_workers = 0
+
+ @classmethod
+ def map_chunked(cls, func: Callable, data: Iterable, max_workers: int) ->
list:
+ """Apply func to data using chunked parallel execution (one task per
worker)."""
+ data_list = list(data)
+ if not data_list:
+ return []
+
+ chunk_size = (len(data_list) + max_workers - 1) // max_workers
+ chunks = [data_list[i : i + chunk_size] for i in range(0,
len(data_list), chunk_size)]
+
+ def process_chunk(chunk):
+ return [func(item) for item in chunk]
+
+ pool = cls.get_pool(max_workers)
+ chunk_results = pool.map(process_chunk, chunks)
+ return list(itertools.chain.from_iterable(chunk_results))
+
+
+# Register shutdown handler to clean up thread pool on process exit
Review Comment:
IIRC, the process might be reused when `spark.python.worker.reuse=True`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
