from torch.utils.data.datapipes.datapipe import DataChunk
from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper

__all__ = ["DataChunkDF", ]


class DataChunkDF(DataChunk):
    """
        DataChunkDF iterating over individual items inside of DataFrame containers,
        to access DataFrames user `raw_iterator`
    """

    def __iter__(self):
        for df in self.items:
            for record in df_wrapper.iterate(df):
                yield record

    def __len__(self):
        total_len = 0
        for df in self.items:
            total_len += df_wrapper.get_len(df)
        return total_len