GRID/src/data/loading/components/iterators.py

import os
import random
from abc import ABC, abstractmethod
from typing import Callable, Dict, List


from src.utils.decorators import retry
from src.utils.file_utils import open_pyarrow_file

# We suppress the tensorflow warnings. Needs to happend before the tf import
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

import tensorflow as tf

tf.config.set_visible_devices([], "GPU")  # Disable all for tensorflow
# if GPU version of TF installed,
# it will automatically occupy the full GPU memory


class RawDataIterator(ABC):
    """the abstract class for raw data iterator (e.g., parquet, avro, etc.)

    Parameters
    ----------
    list_of_file_paths : List[str]
        the list of file paths to read from
    """

    def __init__(
        self,
        **kwargs,
    ):
        self.list_of_file_paths = None
        self.should_shuffle_rows = None

    def update_list_of_file_paths(self, list_of_file_paths: List[str]):
        self.list_of_file_paths = list_of_file_paths

    @abstractmethod
    def get_file_suffix(self) -> str:
        raise NotImplementedError("Must be implemented in child classes")

    @abstractmethod
    def iterrows(self):
        raise NotImplementedError("Must be implemented in child classes")

    @abstractmethod
    def shuffle(self, seed=42):
        raise NotImplementedError("Must be implemented in child classes")

    @abstractmethod
    def iter_batches(self, batch_size: int):
        raise NotImplementedError("Must be implemented in child classes")

    @retry()
    def _get_next_example(self, dataset_iterator):
        try:
            return next(dataset_iterator)
        except StopIteration:
            return None

class ParquetDataIterator(RawDataIterator):
    """Data iterator class for parquet files

    Parameters
    ----------
    list_of_file_paths : List[str]
        the list of file paths to read from
    """

    def __init__(self, buffer_size=1000, features_to_consider=[], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.buffer_size = buffer_size
        self.features_to_consider = features_to_consider

    def iterrows(self):
        assert self.list_of_file_paths is not None, "list_of_file_paths is not set"

        for batch in self.iter_batches(batch_size=self.buffer_size):
            for row in batch.to_pylist():
                yield row

    def iter_batches(self, batch_size: int) -> Dict[str, tf.Tensor]:  # type: ignore
        assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
        for file_path in self.list_of_file_paths:
            with open_pyarrow_file(file_path) as f:
                parquet_file = pq.ParquetFile(f)

                for batch in parquet_file.iter_batches(
                    columns=self.features_to_consider
                    if len(self.features_to_consider)
                    else None,
                    batch_size=batch_size,
                ):
                    yield batch

    def shuffle(self, seed=42) -> RawDataIterator:
        random.seed(seed)
        random.shuffle(self.list_of_file_paths)  # type: ignore
        return self

    def get_file_suffix(self) -> str:
        return "parquet"

class TFRecordIterator(RawDataIterator):
    """Data iterator class for tfrecord files
    Parameters
    ----------
    use_ragged_tensor: bool
        Whether to use ragged tensors.
    batch_tf_processing_functions: List[Callable]
        A list of tensorflow functions to apply to the batches.
    should_drop_last_batch: bool
        Whether to drop the last batch if it is not a multiple of the batch size.
    """

    def __init__(
        self,
        use_ragged_tensor: bool = False,
        batch_tf_processing_functions: List[Callable] = [],
        should_drop_last_batch: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.feature_description = None
        self.use_ragged_tensor = use_ragged_tensor
        self.batch_tf_processing_functions = batch_tf_processing_functions
        self.should_drop_last_batch = should_drop_last_batch

    def initialize_feature_description(self, raw_dataset: tf.data.TFRecordDataset):
        """
        If the feature description is not set, infer the feature description from the first record in the dataset.
        """
        if self.feature_description is None:
            sample_record: tf.Tensor = next(iter(raw_dataset))  # type: ignore

            self.feature_description = self.infer_feature_type(
                self.parse_tfrecord(sample_record).features.feature  # type: ignore
            )

    def iterrows(self):
        assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
        raw_dataset = tf.data.TFRecordDataset(
            [self.list_of_file_paths], compression_type="GZIP"
        )
        if self.should_shuffle_rows:
            # the buffer here is the number of records to shuffle
            # the larger the buffer, the more memory it will use
            # too large might cause OOM
            raw_dataset = raw_dataset.shuffle(buffer_size=128)

        self.initialize_feature_description(raw_dataset)
        # We create an iterator and manually iterate to allow for retrying the
        # "next" operation in case of a failure.
        dataset_iterator = iter(raw_dataset)
        curr_example = self._get_next_example(dataset_iterator)
        while curr_example:
            example = tf.io.parse_single_example(curr_example, self.feature_description)
            yield example
            curr_example = self._get_next_example(dataset_iterator)

    def iter_batches(self, batch_size: int) -> Dict[str, tf.Tensor]:  # type: ignore
        assert self.list_of_file_paths is not None, "list_of_file_paths is not set"
        raw_dataset = tf.data.TFRecordDataset(
            self.list_of_file_paths, compression_type="GZIP"
        )

        if self.should_shuffle_rows:
            # the buffer here is the number of records to shuffle
            # the larger the buffer, the more memory it will use
            # too large might cause OOM
            raw_dataset = raw_dataset.shuffle(buffer_size=128)

        self.initialize_feature_description(raw_dataset)
        # to avoid the issues with tf record warnings, we drop the last instances
        batch_function = (
            raw_dataset.ragged_batch if self.use_ragged_tensor else raw_dataset.batch
        )

        batched_dataset = batch_function(
            batch_size, drop_remainder=self.should_drop_last_batch
        ).prefetch(buffer_size=tf.data.AUTOTUNE)

        for batch_tf_processing_function in self.batch_tf_processing_functions:
            batched_dataset = batched_dataset.map(batch_tf_processing_function)

        dataset_iterator = iter(batched_dataset)

        curr_batch = self._get_next_example(dataset_iterator)

        while curr_batch is not None:
            example = tf.io.parse_example(curr_batch, self.feature_description)
            yield example
            curr_batch = self._get_next_example(dataset_iterator)

    # dynamic inferring the feature description of tfrecord files
    def infer_feature_type(self, example_proto: tf.Tensor) -> dict:
        feature_description = {}
        tf_feature_type = (
            tf.io.RaggedFeature if self.use_ragged_tensor else tf.io.VarLenFeature
        )
        for key, value in example_proto.items():  # type: ignore
            if isinstance(value, tf.train.Feature):
                if value.HasField("bytes_list"):
                    feature_description[key] = tf_feature_type(tf.string)
                elif value.HasField("float_list"):
                    feature_description[key] = tf_feature_type(tf.float32)
                elif value.HasField("int64_list"):
                    feature_description[key] = tf_feature_type(tf.int64)
                else:
                    raise ValueError("Unknown feature type")
        return feature_description

    # parsing the tfrecord files from bytes
    def parse_tfrecord(self, record: tf.Tensor) -> tf.Tensor:
        example = tf.train.Example()
        example.ParseFromString(record.numpy())  # type: ignore
        return example

    def shuffle(self, seed=42) -> RawDataIterator:
        # TODO(lneves): Unify the shuffle method for all iterators
        # Currently this one shuffles only files, parquet shuffles rows.
        random.seed(seed)
        random.shuffle(self.list_of_file_paths)  # type: ignore
        return self

    def get_file_suffix(self) -> str:
        return "tfrecord.gz"