diff --git a/pandarallel/core.py b/pandarallel/core.py index 8eaaf12..f5e6786 100644 --- a/pandarallel/core.py +++ b/pandarallel/core.py @@ -6,6 +6,7 @@ from pathlib import Path from tempfile import NamedTemporaryFile from typing import Any, Callable, Dict, Iterator, Optional, Tuple, Type, cast +from functools import partial import dill import pandas as pd @@ -205,6 +206,7 @@ def parallelize_with_memory_file_system( nb_requested_workers: int, data_type: Type[DataType], progress_bars_type: ProgressBarsType, + single_bar ): def closure( data: Any, @@ -239,7 +241,7 @@ def closure( show_progress_bars = progress_bars_type != ProgressBarsType.No - progress_bars = get_progress_bars(progresses_length, show_progress_bars) + progress_bars = get_progress_bars(progresses_length, show_progress_bars, single_bar) progresses = [0] * nb_workers workers_status = [WorkerStatus.Running] * nb_workers @@ -344,6 +346,7 @@ def parallelize_with_pipe( nb_requested_workers: int, data_type: Type[DataType], progress_bars_type: ProgressBarsType, + single_bar ): def closure( data: Any, @@ -380,7 +383,7 @@ def closure( show_progress_bars = progress_bars_type != ProgressBarsType.No - progress_bars = get_progress_bars(progresses_length, show_progress_bars) + progress_bars = get_progress_bars(progresses_length, show_progress_bars, single_bar) progresses = [0] * nb_workers workers_status = [WorkerStatus.Running] * nb_workers @@ -444,6 +447,7 @@ def initialize( shm_size_mb=None, nb_workers=NB_PHYSICAL_CORES, progress_bar=False, + single_progress_bar=False, verbose=2, use_memory_fs: Optional[bool] = None, ) -> None: @@ -459,6 +463,7 @@ def initialize( if use_memory_fs else parallelize_with_pipe ) + parallelize = partial(parallelize, single_bar=single_progress_bar) if use_memory_fs and not is_memory_fs_available: raise SystemError("Memory file system is not available") diff --git a/pandarallel/progress_bars.py b/pandarallel/progress_bars.py index dfe9190..c7ff979 100644 --- a/pandarallel/progress_bars.py +++ b/pandarallel/progress_bars.py @@ -54,8 +54,11 @@ def is_notebook_lab() -> bool: class ProgressBarsConsole(ProgressBars): - def __init__(self, maxs: List[int], show: bool) -> None: + def __init__(self, maxs: List[int], show: bool, single_bar=True) -> None: self.__show = show + self.__single_bar = single_bar + if self.__single_bar: + maxs = [sum(maxs)] self.__bars = [[0, max] for max in maxs] self.__width = self.__get_width() @@ -107,6 +110,8 @@ def update(self, values: List[int]) -> None: if not self.__show: return + if self.__single_bar: + values = [sum(values)] for index, value in enumerate(values): self.__bars[index][0] = value @@ -118,7 +123,7 @@ def update(self, values: List[int]) -> None: class ProgressBarsNotebookLab(ProgressBars): - def __init__(self, maxs: List[int], show: bool) -> None: + def __init__(self, maxs: List[int], show: bool, single_bar=True) -> None: """Initialization. Positional argument: maxs - List containing the max value of each progress bar @@ -131,6 +136,9 @@ def __init__(self, maxs: List[int], show: bool) -> None: from IPython.display import display from ipywidgets import HBox, IntProgress, Label, VBox + self.__single_bar = single_bar + if self.__single_bar: + maxs = [sum(maxs)] self.__bars = [ HBox( [ @@ -150,7 +158,8 @@ def update(self, values: List[int]) -> None: """ if not self.__show: return - + if self.__single_bar: + values = [sum(values)] for index, value in enumerate(values): bar, label = self.__bars[index].children @@ -166,18 +175,18 @@ def set_error(self, index: int) -> None: """Set a bar on error""" if not self.__show: return - + if self.__single_bar: index = 0 bar, _ = self.__bars[index].children bar.bar_style = "danger" def get_progress_bars( - maxs: List[int], show + maxs: List[int], show, single_bar ) -> Union[ProgressBarsNotebookLab, ProgressBarsConsole]: return ( - ProgressBarsNotebookLab(maxs, show) + ProgressBarsNotebookLab(maxs, show, single_bar) if is_notebook_lab() - else ProgressBarsConsole(maxs, show) + else ProgressBarsConsole(maxs, show, single_bar) )