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    backend: T.Literal["pyarrow", "numpy_nullable"] = "pyarrow"
 73
 74    @T.override
 75    def read(self) -> pd.DataFrame:
 76        # can't limit rows at read time
 77        data = pd.read_parquet(self.path, dtype_backend="pyarrow")
 78        if self.limit is not None:
 79            data = data.head(self.limit)
 80        return data
 81
 82    @T.override
 83    def lineage(
 84        self,
 85        name: str,
 86        data: pd.DataFrame,
 87        targets: str | None = None,
 88        predictions: str | None = None,
 89    ) -> Lineage:
 90        return lineage.from_pandas(
 91            df=data,
 92            name=name,
 93            source=self.path,
 94            targets=targets,
 95            predictions=predictions,
 96        )
 97
 98
 99ReaderKind = ParquetReader
100
101# %% WRITERS
102
103
104class Writer(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
105    """Base class for a dataset writer.
106
107    Use a writer to save a dataset from memory.
108    e.g., to write file, database, cloud storage, ...
109    """
110
111    KIND: str
112
113    @abc.abstractmethod
114    def write(self, data: pd.DataFrame) -> None:
115        """Write a dataframe to a dataset.
116
117        Args:
118            data (pd.DataFrame): dataframe representation.
119        """
120
121
122class ParquetWriter(Writer):
123    """Writer a dataframe to a parquet file.
124
125    Parameters:
126        path (str): local or S3 path to the dataset.
127    """
128
129    KIND: T.Literal["ParquetWriter"] = "ParquetWriter"
130
131    path: str
132
133    @T.override
134    def write(self, data: pd.DataFrame) -> None:
135        pd.DataFrame.to_parquet(data, self.path)
136
137
138WriterKind = ParquetWriter