Skip to content

Commit

Permalink
fix: set _max_workers in ThreadPoolExecutor with non-None initializer (
Browse files Browse the repository at this point in the history
…#1254)

* Add _max_workers definition to ThreadPoolExecutor when arg is not None

* Add test for ThreadPoolExecutor for None and 1 max_workers

* style: pre-commit fixes

---------

Co-authored-by: Nick Manganelli <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
3 people authored Jul 23, 2024
1 parent 599984e commit 0a997cd
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 0 deletions.
2 changes: 2 additions & 0 deletions src/uproot/source/futures.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,8 @@ def __init__(self, max_workers: int | None = None):
import multiprocessing

self._max_workers = multiprocessing.cpu_count()
else:
self._max_workers = max_workers

self._work_queue = queue.Queue()
self._workers = []
Expand Down
44 changes: 44 additions & 0 deletions tests/test_1254_test_threadpool_executor_for_dask.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
import pytest
import skhep_testdata

import uproot

pytest.importorskip("pandas")


def test_decompression_threadpool_executor_for_dask():

class TestThreadPoolExecutor(uproot.source.futures.ThreadPoolExecutor):
def __init__(self, max_workers=None):
super().__init__(max_workers=max_workers)
self.submit_count = 0

def submit(self, task, /, *args, **kwargs):
self.submit_count += 1
super().submit(task, *args, **kwargs)

implicitexecutor = TestThreadPoolExecutor(max_workers=None)

a = uproot.dask(
{skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root"): "sample"},
decompression_executor=implicitexecutor,
)

a["i4"].compute()

assert implicitexecutor.max_workers > 0

assert implicitexecutor.submit_count > 0

explicitexecutor = TestThreadPoolExecutor(max_workers=1)

b = uproot.dask(
{skhep_testdata.data_path("uproot-sample-6.20.04-uncompressed.root"): "sample"},
decompression_executor=explicitexecutor,
)

b["i4"].compute()

assert explicitexecutor.max_workers == 1

assert explicitexecutor.submit_count > 0

0 comments on commit 0a997cd

Please sign in to comment.