diff --git a/src/thread/thread.py b/src/thread/thread.py index 0b20fcd..060057d 100644 --- a/src/thread/thread.py +++ b/src/thread/thread.py @@ -32,7 +32,7 @@ class ParallelProcessing: ... HookFunction, ) from typing_extensions import Generic, ParamSpec -from typing import List, Callable, Optional, Union, Mapping, Sequence, Tuple +from typing import List, Callable, Optional, Union, Mapping, Sequence, Tuple, Generator Threads: set['Thread'] = set() @@ -387,15 +387,19 @@ def _wrap_function(self, function: TargetFunction) -> TargetFunction: @wraps(function) def wrapper( index: int, - data_chunk: Sequence[_Dataset_T], + length: int, + data_chunk: Generator[_Dataset_T, None, None], *args: _Target_P.args, **kwargs: _Target_P.kwargs, ) -> List[_Target_T]: computed: List[Data_Out] = [] - for i, data_entry in enumerate(data_chunk): + + i = 0 + for data_entry in data_chunk: v = function(data_entry, *args, **kwargs) computed.append(v) - self._threads[index].progress = round((i + 1) / len(data_chunk), 5) + self._threads[index].progress = round((i + 1) / length, 5) + i += 1 self._completed += 1 if self._completed == len(self._threads): @@ -507,15 +511,23 @@ def start(self) -> None: i: v for i, v in self.overflow_kwargs.items() if i != 'name' and i != 'args' } - for i, data_chunk in enumerate(chunk_split(self.dataset, max_threads)): + i = 0 + for chunkStart, chunkEnd in chunk_split(len(self.dataset), max_threads): chunk_thread = Thread( target=self.function, - args=[i, data_chunk, *parsed_args, *self.overflow_args], + args=[ + i, + chunkEnd - chunkStart, + (self.dataset[x] for x in range(chunkStart, chunkEnd)), + *parsed_args, + *self.overflow_args, + ], name=name_format and name_format % i or None, **self.overflow_kwargs, ) self._threads.append(_ThreadWorker(chunk_thread, 0)) chunk_thread.start() + i += 1 # Handle abrupt exit diff --git a/src/thread/utils/algorithm.py b/src/thread/utils/algorithm.py index bc531b4..0250abf 100644 --- a/src/thread/utils/algorithm.py +++ b/src/thread/utils/algorithm.py @@ -8,10 +8,10 @@ |_ b.py """ -from typing import List, Sequence, Any +from typing import List, Tuple -def chunk_split(dataset: Sequence[Any], number_of_chunks: int) -> List[List[Any]]: +def chunk_split(dataset_length: int, number_of_chunks: int) -> List[Tuple[int, int]]: """ Splits a dataset into balanced chunks @@ -21,33 +21,32 @@ def chunk_split(dataset: Sequence[Any], number_of_chunks: int) -> List[List[Any] Parameters ---------- - :param dataset: This should be the dataset you want to split into chunks + :param dataset_length: This should be the length of the dataset you want to split into chunks :param number_of_chunks: The should be the number of chunks it will attempt to split into Returns ------- - :returns list[list[Any]]: The split dataset + :returns list[tuple[int, int]]: The chunked dataset slices Raises ------ AssertionError: The number of chunks specified is larger than the dataset size """ - length = len(dataset) assert ( - length >= number_of_chunks + dataset_length >= number_of_chunks ), 'The number of chunks specified is larger than the dataset size' - chunk_count = length // number_of_chunks - overflow = length % number_of_chunks + chunk_count = dataset_length // number_of_chunks + overflow = dataset_length % number_of_chunks i = 0 split = [] - while i < length: + while i < dataset_length: chunk_length = chunk_count + int(overflow > 0) b = i + chunk_length - split.append(dataset[i:b]) + split.append((i, b)) overflow -= 1 i = b diff --git a/tests/test_algorithm.py b/tests/test_algorithm.py new file mode 100644 index 0000000..9e5ed4d --- /dev/null +++ b/tests/test_algorithm.py @@ -0,0 +1,45 @@ +import random +from src.thread.utils import algorithm + + +def test_chunking_1(): + assert algorithm.chunk_split(5, 1) == [(0, 5)] + + +def test_chunking_2(): + assert algorithm.chunk_split(5, 2) == [(0, 3), (3, 5)] + + +def test_chunking_3(): + assert algorithm.chunk_split(100, 8) == [ + (0, 13), + (13, 26), + (26, 39), + (39, 52), + (52, 64), + (64, 76), + (76, 88), + (88, 100), + ] + + +def test_chunking_dynamic(): + dataset_length = random.randint(400, int(10e6)) + thread_count = random.randint(2, 100) + + expected_chunk_length_low = dataset_length // thread_count + expected_chunk_high = dataset_length % thread_count + + i = 0 + heap = [] + while i < dataset_length: + chunk_length = expected_chunk_length_low + int(expected_chunk_high > 0) + b = i + chunk_length + + heap.append((i, b)) + expected_chunk_high -= 1 + i = b + + assert ( + algorithm.chunk_split(dataset_length, thread_count) == heap + ), f'\nLength: {dataset_length}\nThreads: {thread_count}\nExpected: {heap}\nActual: {algorithm.chunk_split(dataset_length, thread_count)}'