content: #  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Materializer for Pandas.
  
  This materializer handles pandas DataFrame and Series objects.
  
  Special features:
      - Handles pandas DataFrames and Series with various data types
      - Provides helpful error messages for custom data type errors
      - Warns when custom data types are detected that might need additional libraries
  
  Environment Variables:
      ZENML_PANDAS_SAMPLE_ROWS: Controls the number of sample rows to include in
          visualizations. Defaults to 10 if not set.
  """
  
  import os
  from typing import Any, ClassVar, Dict, Optional, Tuple, Type, Union
  
  import pandas as pd
  
  from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
  from zenml.enums import ArtifactType, VisualizationType
  from zenml.logger import get_logger
  from zenml.materializers.base_materializer import BaseMaterializer
  from zenml.metadata.metadata_types import DType, MetadataType
  
  logger = get_logger(__name__)
  
  PARQUET_FILENAME = "df.parquet.gzip"
  COMPRESSION_TYPE = "gzip"
  
  CSV_FILENAME = "df.csv"
  
  # Default number of sample rows to display in visualizations
  DEFAULT_SAMPLE_ROWS = 10
  
  # List of standard pandas/numpy dtype prefixes for type checking
  STANDARD_DTYPE_PREFIXES = [
      "int",
      "float",
      "bool",
      "datetime",
      "timedelta",
      "object",
      "category",
      "string",
      "complex",
  ]
  
  
  def is_standard_dtype(dtype_str: str) -> bool:
      """Check if a dtype string represents a standard pandas/numpy dtype.
  
      Args:
          dtype_str: String representation of the dtype
  
      Returns:
          bool: True if it's a standard dtype, False otherwise
      """
      dtype_str = dtype_str.lower()
      return any(prefix in dtype_str for prefix in STANDARD_DTYPE_PREFIXES)
  
  
  class PandasMaterializer(BaseMaterializer):
      """Materializer to read data to and from pandas."""
  
      ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = (
          pd.DataFrame,
          pd.Series,
      )
      ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.DATA
  
      def __init__(
          self, uri: str, artifact_store: Optional[BaseArtifactStore] = None
      ):
          """Define `self.data_path`.
  
          Args:
              uri: The URI where the artifact data is stored.
              artifact_store: The artifact store where the artifact data is stored.
          """
          super().__init__(uri, artifact_store)
          try:
              import pyarrow  # type: ignore # noqa
  
              self.pyarrow_exists = True
          except ImportError:
              self.pyarrow_exists = False
              logger.warning(
                  "By default, the `PandasMaterializer` stores data as a "
                  "`.csv` file. If you want to store data more efficiently, "
                  "you can install `pyarrow` by running "
                  "'`pip install pyarrow`'. This will allow `PandasMaterializer` "
                  "to automatically store the data as a `.parquet` file instead."
              )
          finally:
              self.parquet_path = os.path.join(self.uri, PARQUET_FILENAME)
              self.csv_path = os.path.join(self.uri, CSV_FILENAME)
  
      def load(self, data_type: Type[Any]) -> Union[pd.DataFrame, pd.Series]:
          """Reads `pd.DataFrame` or `pd.Series` from a `.parquet` or `.csv` file.
  
          Args:
              data_type: The type of the data to read.
  
          Raises:
              ImportError: If pyarrow or fastparquet is not installed.
              TypeError: Raised if there is an error when reading parquet files.
              zenml_type_error: If the data type is a custom data type.
  
          Returns:
              The pandas dataframe or series.
          """
          try:
              # First try normal loading
              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)
          except TypeError as e:
              # Check for common data type error patterns
              error_str = str(e).lower()
              is_dtype_error = (
                  "not understood" in error_str
                  or "no type" in error_str
                  or "cannot deserialize" in error_str
                  or "data type" in error_str
              )
  
              if is_dtype_error:
                  # If the error is due to a custom data type, raise a ZenML TypeError
                  # This is to avoid the original error from being swallowed
                  # and to provide a more helpful error message
                  zenml_type_error = TypeError(
                      "Encountered an error with custom data types. This may be due to "
                      "missing libraries that were used when the data was originally created. "
                      "For example, you might need to install libraries like 'geopandas' for "
                      "GeoPandas data types, 'pandas-gbq' for BigQuery data types, or "
                      "'pyarrow' for Arrow data types. Make sure to import these libraries "
                      "in your step code as well as adding them to your step requirements, "
                      "even if you're not directly using them in your code. Pandas needs "
                      "these libraries to be imported to properly load the custom data types. "
                      "Try installing any packages that were used in previous pipeline steps "
                      "but might not be available in the current environment."
                  )
                  raise zenml_type_error from e
              # We don't know how to handle this error, so re-raise the original error
              raise e
  
          # 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 its 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)
  
      def save_visualizations(
          self, df: Union[pd.DataFrame, pd.Series]
      ) -> Dict[str, VisualizationType]:
          """Save visualizations of the given pandas dataframe or series.
  
          Creates two visualizations:
          1. A statistical description of the data (using df.describe())
          2. A sample of the data (first N rows controlled by ZENML_PANDAS_SAMPLE_ROWS)
  
          Note:
              The number of sample rows shown can be controlled with the
              ZENML_PANDAS_SAMPLE_ROWS environment variable.
  
          Args:
              df: The pandas dataframe or series to visualize.
  
          Returns:
              A dictionary of visualization URIs and their types.
          """
          visualizations = {}
          describe_uri = os.path.join(self.uri, "describe.csv")
          describe_uri = describe_uri.replace("\\", "/")
          with self.artifact_store.open(describe_uri, mode="wb") as f:
              df.describe().to_csv(f)
          visualizations[describe_uri] = VisualizationType.CSV
  
          # Get the number of sample rows from environment variable or use default
          sample_rows = int(
              os.environ.get("ZENML_PANDAS_SAMPLE_ROWS", DEFAULT_SAMPLE_ROWS)
          )
  
          # Add our sample visualization (with configurable number of rows)
          if isinstance(df, pd.Series):
              sample_df = df.head(sample_rows).to_frame()
          else:
              sample_df = df.head(sample_rows)
  
          sample_uri = os.path.join(self.uri, "sample.csv")
          sample_uri = sample_uri.replace("\\", "/")
          with self.artifact_store.open(sample_uri, mode="wb") as f:
              sample_df.to_csv(f)
  
          visualizations[sample_uri] = VisualizationType.CSV
  
          return visualizations
  
      def extract_metadata(
          self, df: Union[pd.DataFrame, pd.Series]
      ) -> Dict[str, "MetadataType"]:
          """Extract metadata from the given pandas dataframe or series.
  
          Args:
              df: The pandas dataframe or series to extract metadata from.
  
          Returns:
              The extracted metadata as a dictionary.
          """
          # Store whether it's a Series for later reference
          is_series = isinstance(df, pd.Series)
  
          # Store original shape before conversion
          original_shape = df.shape
  
          # Convert Series to DataFrame for consistent handling of dtypes
          if is_series:
              series_obj = df  # Keep original Series for some calculations
              df = df.to_frame(name="series")
  
          pandas_metadata: Dict[str, "MetadataType"] = {"shape": original_shape}
  
          # Add information about custom data types to metadata
          custom_types = {}
          try:
              for col, dtype in df.dtypes.items():
                  dtype_str = str(dtype)
                  if not is_standard_dtype(dtype_str):
                      col_name = "series" if is_series else str(col)
                      custom_types[col_name] = dtype_str
                      # Try to get module information if available
                      try:
                          module_name = dtype.type.__module__
                          custom_types[f"{col_name}_module"] = module_name
                      except (AttributeError, TypeError):
                          pass
  
              if custom_types:
                  pandas_metadata["custom_types"] = custom_types
          except Exception as e:
              logger.debug(f"Error extracting custom type metadata: {e}")
  
          if is_series:
              # For Series, use the original series object for statistics
              pandas_metadata["dtype"] = DType(series_obj.dtype.type)
              pandas_metadata["mean"] = float(series_obj.mean().item())
              pandas_metadata["std"] = float(series_obj.std().item())
              pandas_metadata["min"] = float(series_obj.min().item())
              pandas_metadata["max"] = float(series_obj.max().item())
  
          else:
              pandas_metadata["dtype"] = {
                  str(key): DType(value.type) for key, value in df.dtypes.items()
              }
              for stat_name, stat in {
                  "mean": df.mean,
                  "std": df.std,
                  "min": df.min,
                  "max": df.max,
              }.items():
                  pandas_metadata[stat_name] = {
                      str(key): float(value)
                      for key, value in stat(numeric_only=True).to_dict().items()
                  }
  
          return pandas_metadata
  
      def get_item_count(self, data: Any) -> Optional[int]:
          """Get the number of items for the given data.
  
          Args:
              data: The data to get the number of items for.
  
          Returns:
              The number of items for the given data.
          """
          if isinstance(data, pd.Series):
              return len(data)
          elif isinstance(data, pd.DataFrame):
              return int(data.shape[1])
          return None
  
      def load_item(self, data_type: Type[Any], index: int) -> Any:
          """Load a specific item of the data.
  
          Args:
              data_type: The type of the data to load.
              index: The index of the item to load.
  
          Returns:
              The loaded item.
          """
          data = self.load(data_type)
          item = data[index]
  
          if not isinstance(item, data_type):
              try:
                  item = data_type(item)
              except Exception as e:
                  # We only log an error here, potentially pydantic can handle the
                  # conversion when validating step function inputs.
                  logger.error(
                      "Failed to convert item `%s` to expected type `%s`. This "
                      "is most likely due to a mismatching type annotation on "
                      "your step function input. Error: %s",
                      item,
                      data_type.__name__,
                      e,
                  )
  
          return item
  