Batched storage will cache the dataframe after reading for the first time

dataflow/pipeline/Pipeline.py

@@ -543,6 +543,9 @@ def _compiled_forward(self, resume_step: int=0, batch_size: int|None=None, resum
             batch_size (int|None): if set, run the pipeline in batch mode with this batch size
             resume_from_last (bool): if True, resume from the last successful step and batch
         """
+        if not self.compiled:
+            raise RuntimeError("Pipeline is not compiled yet. Please call `compile()` before running the pipeline.")
+
         if resume_step > 0 and resume_from_last:
             raise ValueError("Cannot set both `resume_step` and `resume_from_last` to True.")
 

dataflow/utils/storage.py

@@ -940,6 +940,7 @@ def __init__(
         super().__init__(first_entry_file_name, cache_path, file_name_prefix, cache_type)
         self.batch_size = batch_size
         self.batch_step = 0
+        self._dataframe_buffer = {}
         if cache_type not in ["jsonl", "csv"]:
             raise ValueError(f"BatchedFileStorage only supports 'jsonl' and 'csv' cache types, got: {cache_type}")
 
@@ -1004,8 +1005,11 @@ def read(self, output_type: Literal["dataframe", "dict"]="dataframe") -> Any:
                 local_cache = file_path.split(".")[-1]
         else:
             local_cache = self.cache_type
-        # TODO Code below may be a bottleneck for large files, consider optimizing later
-        dataframe = self._load_local_file(file_path, local_cache)
+        if self._dataframe_buffer.get(self.operator_step) is not None:
+            dataframe = self._dataframe_buffer[self.operator_step].copy()
+        else:
+            dataframe = self._load_local_file(file_path, local_cache)
+            self._dataframe_buffer[self.operator_step] = dataframe.copy()
         self.record_count = len(dataframe)
         # 读出当前批次数据
         dataframe = dataframe.iloc[