apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/plugins/pydantic_extensions.py

@@ -0,0 +1,98 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+from typing import Any
+
+from hamilton.data_quality import base, default_validators
+from hamilton.htypes import custom_subclass_check
+
+try:
+    import pydantic  # noqa: F401
+except ModuleNotFoundError as e:
+    raise NotImplementedError(
+        "Cannot import `pydantic` from `pydantic_validators`. Run pip install 'sf-hamilton[pydantic]' if needed."
+    ) from e
+
+try:
+    from pydantic import BaseModel, TypeAdapter, ValidationError
+except ImportError as e:
+    raise NotImplementedError(
+        "`pydantic>=2.0` required to use `pydantic_validators`. Run pip install 'sf-hamilton[pydantic]' if needed."
+    ) from e
+
+
+COLUMN_FRIENDLY_DF_TYPE = False
+
+
+class PydanticModelValidator(base.BaseDefaultValidator):
+    """Pydantic model compatibility validator (requires ``pydantic>=2.0``)
+
+    Note that this validator uses pydantic's strict mode, which does not allow for
+    coercion of data. This means that if an object does not exactly match the reference
+    type, it will fail validation, regardless of whether it could be coerced into the
+    correct type.
+
+    :param model: Pydantic model to validate against
+    :param importance: Importance of the validator, possible values "warn" and "fail"
+    :param arbitrary_types_allowed: Whether arbitrary types are allowed in the model
+    """
+
+    def __init__(self, model: type[BaseModel], importance: str):
+        super(PydanticModelValidator, self).__init__(importance)
+        self.model = model
+        self._model_adapter = TypeAdapter(model)
+
+    @classmethod
+    def applies_to(cls, datatype: type[type]) -> bool:
+        # In addition to checking for a subclass of BaseModel, we also check for dict
+        # as this is the standard 'de-serialized' format of pydantic models in python
+        return custom_subclass_check(datatype, BaseModel) or custom_subclass_check(datatype, dict)
+
+    def description(self) -> str:
+        return "Validates that the returned object is compatible with the specified pydantic model"
+
+    def validate(self, data: Any) -> base.ValidationResult:
+        try:
+            # Currently, validate can not alter the output data, so we must use
+            # strict=True. The downside to this is that data that could be coerced
+            # into the correct type will fail validation.
+            self._model_adapter.validate_python(data, strict=True)
+        except ValidationError as e:
+            return base.ValidationResult(
+                passes=False, message=str(e), diagnostics={"model_errors": e.errors()}
+            )
+        return base.ValidationResult(
+            passes=True,
+            message=f"Data passes pydantic check for model {str(self.model)}",
+        )
+
+    @classmethod
+    def arg(cls) -> str:
+        return "model"
+
+    @classmethod
+    def name(cls) -> str:
+        return "pydantic_validator"
+
+
+def register_validators():
+    """Utility method to append pydantic validators as needed"""
+    validators = [PydanticModelValidator]
+    default_validators.AVAILABLE_DEFAULT_VALIDATORS.extend(validators)
+
+
+register_validators()

hamilton/plugins/pyspark_pandas_extensions.py

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+from typing import Any
+
+try:
+    import pyspark.pandas as ps
+except ImportError as e:
+    raise NotImplementedError("Pyspark is not installed.") from e
+
+from hamilton import registry
+
+DATAFRAME_TYPE = ps.DataFrame
+COLUMN_TYPE = ps.Series
+
+
+@registry.get_column.register(ps.DataFrame)
+def get_column_pyspark_pandas(df: ps.DataFrame, column_name: str) -> ps.Series:
+    return df[column_name]
+
+
+@registry.fill_with_scalar.register(ps.DataFrame)
+def fill_with_scalar_pyspark_pandas(df: ps.DataFrame, column_name: str, value: Any) -> ps.DataFrame:
+    df[column_name] = value
+    return df
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("pyspark_pandas", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+register_types()

hamilton/plugins/sklearn_plot_extensions.py

@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+import dataclasses
+from collections.abc import Collection
+from os import PathLike
+from typing import Any, Union
+
+try:
+    import sklearn.inspection
+    import sklearn.metrics
+    import sklearn.model_selection
+    from matplotlib import pyplot
+except ImportError as e:
+    raise NotImplementedError("scikit-learn is not installed.") from e
+
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataSaver
+
+display_classes = [
+    "ConfusionMatrixDisplay",
+    "DetCurveDisplay",
+    "PrecisionRecallDisplay",
+    "PredictionErrorDisplay",
+    "RocCurveDisplay",
+    "CalibrationDisplay",
+    "DecisionBoundaryDisplay",
+    "LearningCurveDisplay",
+    "PartialDependenceDisplay",
+    "ValidationCurveDisplay",
+]
+SKLEARN_PLOT_TYPES = []
+for class_name in display_classes:
+    # get the attribute via string from sklearn.metrics; if not found return None
+    class_metrics = getattr(sklearn.metrics, class_name, None)
+    class_inspection = getattr(sklearn.inspection, class_name, None)
+    class_model_selection = getattr(sklearn.model_selection, class_name, None)
+    if class_metrics:
+        SKLEARN_PLOT_TYPES.append(class_metrics)
+    if class_inspection:
+        SKLEARN_PLOT_TYPES.append(class_inspection)
+    if class_model_selection:
+        SKLEARN_PLOT_TYPES.append(class_model_selection)
+
+SKLEARN_PLOT_TYPES.append(pyplot.Figure)
+SKLEARN_PLOT_TYPES_ANNOTATION = Union[tuple(SKLEARN_PLOT_TYPES)]
+
+
+@dataclasses.dataclass
+class SklearnPlotSaver(DataSaver):
+    path: str | PathLike
+    # kwargs
+    dpi: float = 200
+    format: str = "png"
+    metadata: dict | None = None
+    bbox_inches: str = None
+    pad_inches: float = 0.1
+    backend: str | None = None
+    papertype: str = None
+    transparent: bool = None
+    bbox_extra_artists: list | None = None
+    pil_kwargs: dict | None = None
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return SKLEARN_PLOT_TYPES
+
+    def _get_saving_kwargs(self) -> dict[str, Any]:
+        kwargs = {}
+        if self.dpi is not None:
+            kwargs["dpi"] = self.dpi
+        if self.format is not None:
+            kwargs["format"] = self.format
+        if self.metadata is not None:
+            kwargs["metadata"] = self.metadata
+        if self.bbox_inches is not None:
+            kwargs["bbox_inches"] = self.bbox_inches
+        if self.pad_inches is not None:
+            kwargs["pad_inches"] = self.pad_inches
+        if self.backend is not None:
+            kwargs["backend"] = self.backend
+        if self.papertype is not None:
+            kwargs["papertype"] = self.papertype
+        if self.transparent is not None:
+            kwargs["transparent"] = self.transparent
+        if self.bbox_extra_artists is not None:
+            kwargs["bbox_extra_artists"] = self.bbox_extra_artists
+        if self.pil_kwargs is not None:
+            kwargs["pil_kwargs"] = self.pil_kwargs
+        return kwargs
+
+    def save_data(self, data: SKLEARN_PLOT_TYPES_ANNOTATION) -> dict[str, Any]:
+        if isinstance(data, pyplot.Figure):
+            figure = data
+        else:
+            data.plot()
+            figure = data.figure_
+        figure.savefig(self.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "png"
+
+
+def register_data_loaders():
+    """Function to register the data loaders for this extension."""
+    registry.register_adapter(SklearnPlotSaver)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/spark_extensions.py

@@ -0,0 +1,105 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+import abc
+import dataclasses
+from collections.abc import Collection
+from typing import Any
+
+try:
+    import pyspark.sql as ps
+except ImportError as e:
+    raise NotImplementedError("Pyspark is not installed.") from e
+
+from pandas import DataFrame
+from pyspark.sql import SparkSession
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader
+
+
+@dataclasses.dataclass
+class SparkDataFrameDataLoader(DataLoader):
+    """Base class for data loaders that load pyspark dataframes.
+    We are not yet including data savers, but that will be added to this most likely..
+    """
+
+    spark: SparkSession
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [ps.DataFrame]
+
+    @abc.abstractmethod
+    def load_data(self, type_: type[DataFrame]) -> tuple[ps.DataFrame, dict[str, Any]]:
+        pass
+
+
+@dataclasses.dataclass
+class CSVDataLoader(SparkDataFrameDataLoader):
+    path: str  # It supports multiple but for now we're going to have a single one
+    # We can always make that a list of strings, or make a multiple reader (.multicsv)
+    header: bool = True
+    sep: str = ","
+
+    def load_data(self, type_: type[DataFrame]) -> tuple[ps.DataFrame, dict[str, Any]]:
+        return (
+            self.spark.read.csv(self.path, header=self.header, sep=self.sep, inferSchema=True),
+            utils.get_file_metadata(self.path),
+        )
+
+    @classmethod
+    def name(cls) -> str:
+        return "csv"
+
+
+@dataclasses.dataclass
+class ParquetDataLoader(SparkDataFrameDataLoader):
+    path: str  # It supports multiple but for now we're going to have a single one
+
+    # We can always make that a list of strings, or make a multiple reader (.multicsv)
+
+    def load_data(self, type_: type[DataFrame]) -> tuple[ps.DataFrame, dict[str, Any]]:
+        return self.spark.read.parquet(self.path), utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "parquet"
+
+
+def register_data_loaders():
+    """Function to register the data loaders for this extension."""
+    for loader in [CSVDataLoader, ParquetDataLoader]:
+        registry.register_adapter(loader)
+
+
+COLUMN_FRIENDLY_DF_TYPE = False
+
+register_data_loaders()
+
+
+DATAFRAME_TYPE = ps.DataFrame
+COLUMN_TYPE = None
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("spark", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+register_types()

hamilton/plugins/vaex_extensions.py

@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+from typing import Any
+
+import numpy as np
+
+from hamilton import registry
+
+try:
+    import vaex
+except ImportError as e:
+    raise NotImplementedError("Vaex is not installed.") from e
+
+DATAFRAME_TYPE = vaex.dataframe.DataFrame
+COLUMN_TYPE = vaex.expression.Expression
+
+
+@registry.get_column.register(vaex.dataframe.DataFrame)
+def get_column_vaex(df: vaex.dataframe.DataFrame, column_name: str) -> vaex.expression.Expression:
+    return df[column_name]
+
+
+@registry.fill_with_scalar.register(vaex.dataframe.DataFrame)
+def fill_with_scalar_vaex(
+    df: vaex.dataframe.DataFrame, column_name: str, scalar_value: Any
+) -> vaex.dataframe.DataFrame:
+    df[column_name] = np.full((df.shape[0],), scalar_value)
+    return df
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("vaex", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+register_types()

hamilton/plugins/xgboost_extensions.py

@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+import dataclasses
+from collections.abc import Collection
+from os import PathLike
+from typing import Any, Union
+
+try:
+    import xgboost
+except ImportError as e:
+    raise NotImplementedError("XGBoost is not installed.") from e
+
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+XGBOOST_MODEL_TYPES = [xgboost.XGBModel, xgboost.Booster]
+XGBOOST_MODEL_TYPES_ANNOTATION = Union[xgboost.XGBModel, xgboost.Booster]
+
+
+@dataclasses.dataclass
+class XGBoostJsonWriter(DataSaver):
+    """Write XGBoost models and boosters to json format
+    See differences with pickle format: https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html
+    """
+
+    path: str | PathLike
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return XGBOOST_MODEL_TYPES
+
+    def save_data(self, data: XGBOOST_MODEL_TYPES_ANNOTATION) -> dict[str, Any]:
+        data.save_model(self.path)
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def name(cls) -> str:
+        return "json"
+
+
+@dataclasses.dataclass
+class XGBoostJsonReader(DataLoader):
+    """Load XGBoost models and boosters to json format
+    See differences with pickle format: https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html
+    """
+
+    path: str | bytearray | PathLike
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return XGBOOST_MODEL_TYPES
+
+    def load_data(self, type_: type) -> tuple[XGBOOST_MODEL_TYPES_ANNOTATION, dict[str, Any]]:
+        model = type_()
+        model.load_model(self.path)
+        metadata = utils.get_file_metadata(self.path)
+        return model, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "json"
+
+
+def register_data_loaders():
+    for loader in [
+        XGBoostJsonReader,
+        XGBoostJsonWriter,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/yaml_extensions.py

@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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
+#
+#   http://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.
+
+try:
+    import yaml
+except ImportError as e:
+    raise NotImplementedError("yaml is not installed and is needed for yaml hamilton plugin") from e
+
+import dataclasses
+import pathlib
+from collections.abc import Collection
+from typing import Any, Union
+
+from hamilton import registry
+from hamilton.io.data_adapters import DataLoader, DataSaver
+from hamilton.io.utils import get_file_metadata
+
+PrimitiveTypes = str, int, float, bool, dict, list
+AcceptedTypes = Union[PrimitiveTypes]
+
+
+@dataclasses.dataclass
+class YAMLDataLoader(DataLoader):
+    path: str | pathlib.Path
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [*PrimitiveTypes]
+
+    @classmethod
+    def name(cls) -> str:
+        return "yaml"
+
+    def load_data(self, type_: type) -> tuple[AcceptedTypes, dict[str, Any]]:
+        path = self.path
+        if isinstance(self.path, str):
+            path = pathlib.Path(self.path)
+
+        with path.open(mode="r") as f:
+            return yaml.safe_load(f), get_file_metadata(path)
+
+
+@dataclasses.dataclass
+class YAMLDataSaver(DataSaver):
+    path: str | pathlib.Path
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [*PrimitiveTypes]
+
+    @classmethod
+    def name(cls) -> str:
+        return "yaml"
+
+    def save_data(self, data: AcceptedTypes) -> dict[str, Any]:
+        path = self.path
+        if isinstance(path, str):
+            path = pathlib.Path(path)
+        with path.open("w") as f:
+            yaml.dump(data, f)
+        return get_file_metadata(self.path)
+
+
+COLUMN_FRIENDLY_DF_TYPE = False
+
+
+def register_data_loaders():
+    for materializer in [
+        YAMLDataLoader,
+        YAMLDataSaver,
+    ]:
+        registry.register_adapter(materializer)
+
+
+register_data_loaders()