1# mypy: allow-untyped-defs 2from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper 3from torch.utils.data.datapipes.datapipe import DataChunk 4 5 6__all__ = ["DataChunkDF"] 7 8 9class DataChunkDF(DataChunk): 10 """DataChunkDF iterating over individual items inside of DataFrame containers, to access DataFrames user `raw_iterator`.""" 11 12 def __iter__(self): 13 for df in self.items: 14 yield from df_wrapper.iterate(df) 15 16 def __len__(self): 17 total_len = 0 18 for df in self.items: 19 total_len += df_wrapper.get_len(df) 20 return total_len 21