bikes.io.datasets

  1"""Read/Write datasets from/to external sources/destinations."""
  2
  3# %% IMPORTS
  4
  5import abc
  6import typing as T
  7
  8import mlflow.data.pandas_dataset as lineage
  9import pandas as pd
 10import pydantic as pdt
 11
 12# %% TYPINGS
 13
 14Lineage: T.TypeAlias = lineage.PandasDataset
 15
 16# %% READERS
 17
 18
 19class Reader(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 20    """Base class for a dataset reader.
 21
 22    Use a reader to load a dataset in memory.
 23    e.g., to read file, database, cloud storage, ...
 24
 25    Parameters:
 26        limit (int, optional): maximum number of rows to read. Defaults to None.
 27    """
 28
 29    KIND: str
 30
 31    limit: int | None = None
 32
 33    @abc.abstractmethod
 34    def read(self) -> pd.DataFrame:
 35        """Read a dataframe from a dataset.
 36
 37        Returns:
 38            pd.DataFrame: dataframe representation.
 39        """
 40
 41    @abc.abstractmethod
 42    def lineage(
 43        self,
 44        name: str,
 45        data: pd.DataFrame,
 46        targets: str | None = None,
 47        predictions: str | None = None,
 48    ) -> Lineage:
 49        """Generate lineage information.
 50
 51        Args:
 52            name (str): dataset name.
 53            data (pd.DataFrame): reader dataframe.
 54            targets (str | None): name of the target column.
 55            predictions (str | None): name of the prediction column.
 56
 57        Returns:
 58            Lineage: lineage information.
 59        """
 60
 61
 62class ParquetReader(Reader):
 63    """Read a dataframe from a parquet file.
 64
 65    Parameters:
 66        path (str): local path to the dataset.
 67    """
 68
 69    KIND: T.Literal["ParquetReader"] = "ParquetReader"
 70
 71    path: str
 72
 73    @T.override
 74    def read(self) -> pd.DataFrame:
 75        # can't limit rows at read time
 76        data = pd.read_parquet(self.path)
 77        if self.limit is not None:
 78            data = data.head(self.limit)
 79        return data
 80
 81    @T.override
 82    def lineage(
 83        self,
 84        name: str,
 85        data: pd.DataFrame,
 86        targets: str | None = None,
 87        predictions: str | None = None,
 88    ) -> Lineage:
 89        return lineage.from_pandas(
 90            df=data, name=name, source=self.path, targets=targets, predictions=predictions
 91        )
 92
 93
 94ReaderKind = ParquetReader
 95
 96# %% WRITERS
 97
 98
 99class Writer(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
100    """Base class for a dataset writer.
101
102    Use a writer to save a dataset from memory.
103    e.g., to write file, database, cloud storage, ...
104    """
105
106    KIND: str
107
108    @abc.abstractmethod
109    def write(self, data: pd.DataFrame) -> None:
110        """Write a dataframe to a dataset.
111
112        Args:
113            data (pd.DataFrame): dataframe representation.
114        """
115
116
117class ParquetWriter(Writer):
118    """Writer a dataframe to a parquet file.
119
120    Parameters:
121        path (str): local or S3 path to the dataset.
122    """
123
124    KIND: T.Literal["ParquetWriter"] = "ParquetWriter"
125
126    path: str
127
128    @T.override
129    def write(self, data: pd.DataFrame) -> None:
130        pd.DataFrame.to_parquet(data, self.path)
131
132
133WriterKind = ParquetWriter