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