(self, *args: P.args, **kwargs: P.kwargs)
| 41 | # object.__setattr__(self, "run", wrapped) |
| 42 | |
| 43 | def run(self, *args: P.args, **kwargs: P.kwargs) -> List[R]: |
| 44 | # —— 1. 提取 storage —— |
| 45 | # only support FileStorage for now |
| 46 | if args: |
| 47 | storage: FileStorage = args[0] # type: ignore[assignment] |
| 48 | rest_args = args[1:] |
| 49 | rest_kwargs = kwargs |
| 50 | else: |
| 51 | storage: FileStorage = kwargs.get("storage") # type: ignore[assignment] |
| 52 | if storage is None: |
| 53 | raise ValueError( |
| 54 | f"A DataFlowStorage is required for {self._operator!r}.run()" |
| 55 | ) |
| 56 | rest_kwargs = {k: v for k, v in kwargs.items() if k != "storage"} |
| 57 | rest_args = () |
| 58 | |
| 59 | # prepare a dummy storage，for batch processing |
| 60 | if self._batch_cache: # if we need to cache each batch result |
| 61 | self._dummy_storage = DummyStorage( |
| 62 | cache_path=storage.cache_path, |
| 63 | file_name_prefix=storage.file_name_prefix, |
| 64 | cache_type=storage.cache_type |
| 65 | ) |
| 66 | else: # if we don't need to cache each batch result |
| 67 | self._dummy_storage = DummyStorage() |
| 68 | |
| 69 | # —— 2. 读出全量数据并按 batch_size 切分 —— |
| 70 | whole_dataframe = storage.read() |
| 71 | num_batches = (len(whole_dataframe) + self._batch_size - 1) // self._batch_size # Calculate number of batches |
| 72 | |
| 73 | output_dataframe = pd.DataFrame() |
| 74 | self._logger.info(f"Total {len(whole_dataframe)} items, will process in {num_batches} batches of size {self._batch_size}.") |
| 75 | # for batch_num in tqdm(range(num_batches)): |
| 76 | for batch_num in tqdm(range(self.start_batch, num_batches)): # 手动设置断点 |
| 77 | start_index = batch_num * self._batch_size |
| 78 | end_index = min((batch_num + 1) * self._batch_size, len(whole_dataframe)) |
| 79 | batch_df = whole_dataframe.iloc[start_index:end_index] |
| 80 | # Clear and write the current batch |
| 81 | self._dummy_storage.set_data(batch_df) |
| 82 | self._dummy_storage.set_file_name_prefix( |
| 83 | f"{storage.file_name_prefix}_step{storage.operator_step}_batch{batch_num}" |
| 84 | ) |
| 85 | # Run the operator with the dummy storage |
| 86 | self._logger.info(f"Running batch with {len(batch_df)} items...") |
| 87 | self._operator.run(self._dummy_storage, *rest_args, **rest_kwargs) |
| 88 | |
| 89 | res: pd.DataFrame = self._dummy_storage.read() |
| 90 | output_dataframe = pd.concat([output_dataframe, res], axis=0) |
| 91 | |
| 92 | # Find columns in res that are not in whole_dataframe |
| 93 | # new_cols = [c for c in res.columns if c not in whole_dataframe.columns] |
| 94 | |
| 95 | # Create new columns in whole_dataframe with NaN values |
| 96 | # for c in new_cols: |
| 97 | # whole_dataframe[c] = pd.NA |
| 98 | |
| 99 | # Write the values from res back to whole_dataframe |
| 100 | # whole_dataframe.loc[res.index, res.columns] = res |
nothing calls this directly
no test coverage detected