
data_type: The type of the data to read.

Raises:ImportError: If pyarrow or fastparquet is not installed.

Returns:

The pandas dataframe or series.

"""

if self.artifact_store.exists(self.parquet_path):

if self.pyarrow_exists:

with self.artifact_store.open(

self.parquet_path, mode="rb"

) as f:

df = pd.read_parquet(f)

else:

raise ImportError(

"You have an old version of a `PandasMaterializer` "

"data artifact stored in the artifact store "

"as a `.parquet` file, which requires `pyarrow` "

"for reading, You can install `pyarrow` by running "

"'`pip install pyarrow fastparquet`'."

else:

with self.artifact_store.open(self.csv_path, mode="rb") as f:

df = pd.read_csv(f, index_col=0, parse_dates=True)

# validate the type of the data.

def is_dataframe_or_series(

df: Union[pd.DataFrame, pd.Series],

) -> Union[pd.DataFrame, pd.Series]:

"""Checks if the data is a `pd.DataFrame` or `pd.Series`.

Args:

df: The data to check.

Returns:

The data if it is a `pd.DataFrame` or `pd.Series`.

"""

if issubclass(data_type, pd.Series):

# Taking the first column if it is a series as the assumption

# is that there will only be one

assert len(df.columns) == 1

df = df[df.columns[0]]

return df

else:

return df

return is_dataframe_or_series(df)

def save(self, df: Union[pd.DataFrame, pd.Series]) -> None:

"""Writes a pandas dataframe or series to the specified filename.

Args:

df: The pandas dataframe or series to write.

"""

if isinstance(df, pd.Series):

df = df.to_frame(name="series")

if self.pyarrow_exists:

with self.artifact_store.open(self.parquet_path, mode="wb") as f:

df.to_parquet(f, compression=COMPRESSION_TYPE)

else:

with self.artifact_store.open(self.csv_path, mode="wb") as f:

df.to_csv(f, index=True)

Code example

Let's see how materialization works with a basic example. Let's say you have a custom class called MyObject that flows between two steps in a pipeline:

import logging

from zenml import step, pipeline

class MyObj:

def __init__(self, name: str):

self.name = name

@step