Class BatchWrapper

dataflow/wrapper/batch_wrapper.py:20–102

通用的批处理 Wrapper。 在静态检查/IDE 里，BatchWrapper.run 的签名和 operator.run 完全一致。 运行时，会把 operator.run 的 __doc__ 和 __signature__ 也拷过来， 这样 help(bw.run) 时能看到原 operator 的文档。

        ...

class BatchWrapper(Generic[P, R]):
    """
    通用的批处理 Wrapper。

    在静态检查/IDE 里，BatchWrapper.run 的签名和 operator.run 完全一致。
    运行时，会把 operator.run 的 __doc__ 和 __signature__ 也拷过来，
    这样 help(bw.run) 时能看到原 operator 的文档。
    """
    def __init__(self, operator: HasRun[P, R], start_batch: int = 0, batch_size: int = 32, batch_cache: bool = False) -> None:
        self._operator = operator
        self._logger = get_logger()
        self._batch_size = batch_size
        self._batch_cache = batch_cache
        self.start_batch = start_batch

        # 动态拷贝 operator.run 的 __doc__ 和 inspect.signature
        orig = operator.run
        sig = inspect.signature(orig)
        # wrapped = wraps(orig)(self.run)        # 先 wrap docstring, __name__…
        # wrapped.__signature__ = sig            # 再贴上准确的 signature
        # 把它绑到实例上覆盖掉 class method，这样 help(instance.run) 能看到原文档
        # object.__setattr__(self, "run", wrapped)

    def run(self, *args: P.args, **kwargs: P.kwargs) -> List[R]:
        # —— 1. 提取 storage —— 
        # only support FileStorage for now
        if args:
            storage: FileStorage = args[0]    # type: ignore[assignment]
            rest_args = args[1:]
            rest_kwargs = kwargs
        else:
            storage: FileStorage = kwargs.get("storage")      # type: ignore[assignment]
            if storage is None:
                raise ValueError(
                    f"A DataFlowStorage is required for {self._operator!r}.run()"
                )
            rest_kwargs = {k: v for k, v in kwargs.items() if k != "storage"}
            rest_args = ()
            
        # prepare a dummy storage，for batch processing
        if self._batch_cache: # if we need to cache each batch result
            self._dummy_storage = DummyStorage(
                cache_path=storage.cache_path,
                file_name_prefix=storage.file_name_prefix,
                cache_type=storage.cache_type
            )
        else:  # if we don't need to cache each batch result
            self._dummy_storage = DummyStorage()

        # —— 2. 读出全量数据并按 batch_size 切分 —— 
        whole_dataframe = storage.read()
        num_batches = (len(whole_dataframe) + self._batch_size - 1) // self._batch_size  # Calculate number of batches

        output_dataframe = pd.DataFrame()
        self._logger.info(f"Total {len(whole_dataframe)} items, will process in {num_batches} batches of size {self._batch_size}.")
        # for batch_num in tqdm(range(num_batches)):
        for batch_num in tqdm(range(self.start_batch, num_batches)):       # 手动设置断点
            start_index = batch_num * self._batch_size