Class WrapWorkFunctionForFileSystem

pandarallel/core.py:48–110

class WrapWorkFunctionForFileSystem:
    def __init__(
        self,
        work_function: Callable[
            [Any, Callable, tuple, Dict[str, Any], Dict[str, Any]], Any
        ],
    ) -> None:
        self.work_function = work_function

    def __call__(
        self,
        input_file_path: Path,
        output_file_path: Path,
        progress_bars_type: ProgressBarsType,
        worker_index: int,
        master_workers_queue: multiprocessing.Queue,
        dilled_user_defined_function: bytes,
        user_defined_function_args: tuple,
        user_defined_function_kwargs: Dict[str, Any],
        extra: Dict[str, Any],
    ) -> None:
        try:
            # Load dataframe from input file
            with input_file_path.open("rb") as file_descriptor:
                data = pickle.load(file_descriptor)

            # Delete input file since we don't need it any more. It will free some RAM
            # since the input file is stored into Shared Memory.
            input_file_path.unlink()

            data_size = len(data)
            user_defined_function: Callable = dill.loads(dilled_user_defined_function)

            progress_wrapped_user_defined_function = progress_wrapper(
                user_defined_function, master_workers_queue, worker_index, data_size
            )

            used_user_defined_function = (
                progress_wrapped_user_defined_function
                if progress_bars_type
                in (
                    ProgressBarsType.InUserDefinedFunction,
                    ProgressBarsType.InUserDefinedFunctionMultiplyByNumberOfColumns,
                )
                else user_defined_function
            )

            result = self.work_function(
                data,
                used_user_defined_function,
                user_defined_function_args,
                user_defined_function_kwargs,
                extra,
            )

            with output_file_path.open("wb") as file_descriptor:
                pickle.dump(result, file_descriptor)