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

hamilton/experimental/decorators/parameterize_frame.py

@@ -0,0 +1,233 @@
+# 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 pandas as pd
+
+from hamilton.function_modifiers import (
+    UpstreamDependency,
+    base,
+    parameterize_extract_columns,
+    source,
+    value,
+)
+from hamilton.function_modifiers.expanders import ParameterizedExtract
+
+
+def _get_dep_type(dep_type: str) -> UpstreamDependency:
+    """Converts dependency type to the type known by function_modifier"""
+    if dep_type == "out":
+        return None
+    if dep_type == "value":
+        return value
+    if dep_type == "source":
+        return source
+    raise ValueError(f"Invalid dep type: {dep_type}")
+
+
+def _get_index_levels(index: pd.MultiIndex) -> list[list]:
+    out = [[] for _ in index[0]]
+    for specific_index in index:
+        for i, key in enumerate(specific_index):
+            out[i].append(key)
+    return out
+
+
+def _validate_df_parameterization(parameterization: pd.DataFrame):
+    # TODO -- validate that its a multi-index
+    columns = _get_index_levels(parameterization.columns)
+    if (not len(columns) == 2) or "out" not in columns[1]:
+        raise base.InvalidDecoratorException(
+            "Decorator must have a double-index -- first index should be a "
+            "list of {output, source, value} strs. Second must be a list of "
+            "arguments in your function."
+        )
+
+
+def _convert_params_from_df(parameterization: pd.DataFrame) -> list[ParameterizedExtract]:
+    _validate_df_parameterization(parameterization)
+    args, dep_types = _get_index_levels(parameterization.columns)
+    dep_types_converted = [_get_dep_type(val) for val in dep_types]
+    out = []
+    for _, column_set in parameterization.iterrows():
+        parameterization = {
+            arg: dep_type(col_value)
+            for arg, col_value, dep_type in zip(args, column_set, dep_types_converted, strict=False)
+            if dep_type is not None
+        }
+        extracted_columns = [
+            col for col, dep_type in zip(column_set, dep_types, strict=False) if dep_type == "out"
+        ]
+        out.append(ParameterizedExtract(tuple(extracted_columns), parameterization))
+    return out
+
+
+class parameterize_frame(parameterize_extract_columns):
+    """EXPERIMENTAL! Instantiates a parameterize_extract decorator using a dataframe to specify a set of extracts + \
+    parameterizations.
+
+    This is an experimental decorator and the API may change in the future; please provide feedback \
+    whether this API does or does not work for you.
+
+    :param parameterization: Parameterization dataframe. See below.
+
+    This is of a specific shape:
+
+        1. Index - Level 0: list of parameter names
+        2. Index - Level 1: types of things to inject, either:
+
+            - "out" (meaning this is an output),
+            - "value" (meaning this is a literal value)
+            - "source" (meaning this node comes from an upstream value)
+
+        3. Contents:
+
+          - Each row corresponds to the index. Each of these corresponds to an output node from this.
+
+
+    Note your function has to take in the column-names and output a dataframe with those names -- \
+    we will likely change it so that's not the case, and it can just use the position of the columns.
+
+    Example usage:
+
+    .. code-block:: python
+
+        from hamilton.experimental.decorators.parameterize_frame import parameterize_frame
+        df = pd.DataFrame(
+        [
+           ["outseries1a", "outseries2a", "inseries1a", "inseries2a", 5.0],
+           ["outseries1b", "outseries2b", "inseries1b", "inseries2b", 0.2],
+        ],
+        # specify column names corresponding to function arguments and
+        # if outputting multiple columns, output dataframe columns.
+        columns=[
+           ["output1", "output2", "input1", "input2", "input3"],
+           ["out", "out", "source", "source", "value"],
+        ])
+
+        @parameterize_frame(df)
+        def my_func(
+            input1: pd.Series, input2: pd.Series, input3: float
+        ) -> pd.DataFrame:
+           ...
+
+    """
+
+    def __init__(self, parameterization: pd.DataFrame):
+        super(parameterize_frame, self).__init__(*_convert_params_from_df(parameterization))
+
+
+# Examples below
+if __name__ == "__main__":
+    df = pd.DataFrame(
+        [
+            ["outseries1a", "outseries2a", "inseries1a", "inseries2a", 5.0],
+            ["outseries1b", "outseries2b", "inseries1b", "inseries2b", 0.2],
+            # ...
+        ],
+        # Have to switch as indices have to be unique
+        columns=[
+            [
+                "output1",
+                "output2",
+                "input1",
+                "input2",
+                "input3",
+            ],
+            # configure whether column is source or value and also whether it's input ("source", "value") or output ("out")
+            ["out", "out", "source", "source", "value"],
+        ],
+    )  # specify column names (corresponding to function arguments and (if outputting multiple columns) output dataframe columns)
+
+    @parameterize_frame(df)
+    def my_func(input1: pd.Series, input2: pd.Series, input3: float) -> pd.DataFrame:
+        return pd.DataFrame(
+            [input1 * input2 * input3, input1 + input2 + input3]
+        )  # if there's a single column it could maybe just return a series instead and pick up the name from the first column of the dataframe
+
+    @parameterize_extract_columns(
+        ParameterizedExtract(
+            ("outseries1a", "outseries2a"),
+            {"input1": source("inseries1a"), "input2": source("inseries2a"), "input3": value(5.0)},
+        ),
+        ParameterizedExtract(
+            ("outseries1b", "outseries2b"),
+            {"input1": source("inseries1b"), "input2": source("inseries2b"), "input3": value(0.2)},
+        ),
+    )
+    def my_func_parameterized_extract(
+        input1: pd.Series, input2: pd.Series, input3: float
+    ) -> pd.DataFrame:
+        print("running my_func_parameterized_extract")
+        return pd.concat([input1 * input2 * input3, input1 + input2 + input3], axis=1)
+
+    my_func_parameterized_extract.decorated = "false"
+
+    # Test by running the @parameterized_extract decorator
+    from hamilton.ad_hoc_utils import create_temporary_module
+    from hamilton.driver import Driver
+
+    dr = Driver({}, create_temporary_module(my_func_parameterized_extract))
+    dr.visualize_execution(
+        final_vars=["outseries1a", "outseries1b", "outseries2a", "outseries2b"],
+        output_file_path="./out1.pdf",
+        render_kwargs={},
+        inputs={
+            "inseries1a": pd.Series([1, 2]),
+            "inseries1b": pd.Series([2, 3]),
+            "inseries2a": pd.Series([3, 4]),
+            "inseries2b": pd.Series([4, 5]),
+        },
+    )
+
+    df_1 = dr.execute(
+        final_vars=["outseries1a", "outseries1b", "outseries2a", "outseries2b"],
+        # final_vars=["outseries1a", "outseries2a"],
+        inputs={
+            "inseries1a": pd.Series([1, 2]),
+            "inseries1b": pd.Series([2, 3]),
+            "inseries2a": pd.Series([3, 4]),
+            "inseries2b": pd.Series([4, 5]),
+        },
+    )
+    print(df_1)
+
+    # Test by running the @parameterized_extract decorator
+    dr = Driver({}, create_temporary_module(my_func))
+    dr.visualize_execution(
+        final_vars=["outseries1a", "outseries1b", "outseries2a", "outseries2b"],
+        output_file_path="./out2.pdf",
+        render_kwargs={},
+        inputs={
+            "inseries1a": pd.Series([1, 2]),
+            "inseries1b": pd.Series([2, 3]),
+            "inseries2a": pd.Series([3, 4]),
+            "inseries2b": pd.Series([4, 5]),
+        },
+    )
+
+    df_2 = dr.execute(
+        final_vars=["outseries1a", "outseries1b", "outseries2a", "outseries2b"],
+        # final_vars=["outseries1a", "outseries2a"],
+        inputs={
+            "inseries1a": pd.Series([1, 2]),
+            "inseries1b": pd.Series([2, 3]),
+            "inseries2a": pd.Series([3, 4]),
+            "inseries2b": pd.Series([4, 5]),
+        },
+    )
+    print(df_2)

hamilton/experimental/h_async.py

@@ -0,0 +1,29 @@
+# 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 logging
+
+import hamilton.async_driver
+
+logger = logging.getLogger(__name__)
+
+logger.warning(
+    "This module is deprecated and will be removed in Hamilton 2.0 "
+    "Please use `hamilton.async_driver` instead. "
+)
+
+AsyncDriver = hamilton.async_driver.AsyncDriver

hamilton/experimental/h_cache.py

@@ -0,0 +1,413 @@
+# 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 json
+import logging
+import os
+import pickle
+from collections.abc import Callable
+from functools import singledispatch
+from typing import Any
+
+import typing_inspect
+
+from hamilton.base import SimplePythonGraphAdapter
+from hamilton.node import Node
+
+logger = logging.getLogger(__name__)
+
+
+logger.warning(
+    "The module `hamilton.experimental.h_cache` and the class `CachingGraphAdapter `"
+    "are deprecated and will be removed in Hamilton 2.0. "
+    "Consider enabling the core caching feature via `Builder.with_cache()`. "
+    "This might not be 1-to-1 replacement, so please reach out if there are missing features. "
+    "See https://hamilton.apache.org/concepts/caching/ to learn more."
+)
+
+
+"""
+Base SERDE functions.
+
+Basic format is:
+@singledispatch
+def write_<format>(data: object, filepath: str, name: str) -> None:
+
+@singledispatch
+def read_<format>(data: object, filepath: str) -> Any:
+
+Functions should register themselves with the appropriate type.
+"""
+
+
+@singledispatch
+def write_feather(data: object, filepath: str, name: str) -> None:
+    """Writes data to a feather file."""
+    raise NotImplementedError(f"No feather writer for type {type(data)} registered.")
+
+
+@singledispatch
+def read_feather(data: object, filepath: str) -> Any:
+    """Reads from a feather file"""
+    raise NotImplementedError(f"No feather reader for type {type(data)} registered.")
+
+
+@singledispatch
+def write_parquet(data: object, filepath: str, name: str) -> None:
+    """Writes data to a parquet file."""
+    raise NotImplementedError(f"No parquet writer for type {type(data)} registered.")
+
+
+@singledispatch
+def read_parquet(data: object, filepath: str) -> Any:
+    """Reads from a parquet file"""
+    raise NotImplementedError(f"No parquet reader for type {type(data)} registered.")
+
+
+@singledispatch
+def write_json(data: object, filepath: str, name: str) -> None:
+    """Writes data to a json file."""
+    raise NotImplementedError(f"No json writer for type {type(data)} registered.")
+
+
+@singledispatch
+def read_json(data: object, filepath: str) -> Any:
+    """Reads from a json file"""
+    raise NotImplementedError(f"No json reader for type {type(data)} registered.")
+
+
+@singledispatch
+def write_pickle(data: Any, filepath: str, name: str) -> None:
+    """Writes data to a pickle file."""
+    raise NotImplementedError(f"No object writer for type {type(data)} registered.")
+
+
+@singledispatch
+def read_pickle(data: Any, filtepath: str) -> object:
+    """Reads from a pickle file"""
+    raise NotImplementedError(f"No object reader for type {type(data)} registered.")
+
+
+try:
+    import pandas as pd  # conditional import to avoid pandas dependency
+
+    @write_json.register(pd.DataFrame)
+    def write_json_pd1(data: pd.DataFrame, filepath: str, name: str) -> None:
+        """Writes a dataframe to a feather file."""
+        return data.to_json(filepath)
+
+    @write_json.register(pd.Series)
+    def write_json_pd2(data: pd.Series, filepath: str, name: str) -> None:
+        """Writes a series to a feather file."""
+        _df = data.to_frame(name=name)
+        return _df.to_json(filepath)
+
+    @read_json.register(pd.Series)
+    def read_json_pd1(data: pd.Series, filepath: str) -> pd.Series:
+        """Reads a series from a feather file."""
+        _df = pd.read_json(filepath)
+        return _df[_df.columns[0]]
+
+    @read_json.register(pd.DataFrame)
+    def read_json_pd2(data: pd.DataFrame, filepath: str) -> pd.DataFrame:
+        """Reads a dataframe from a feather file."""
+        return pd.read_json(filepath)
+
+    try:
+        import pyarrow  # noqa: F401  # conditional import to avoid pyarrow dependency
+
+        @write_feather.register(pd.DataFrame)
+        def write_feather_pd1(data: pd.DataFrame, filepath: str, name: str) -> None:
+            """Writes a dataframe to a feather file."""
+            data.to_feather(filepath)
+
+        @write_feather.register(pd.Series)
+        def write_feather_pd2(data: pd.Series, filepath: str, name: str) -> None:
+            """Writes a series to a feather file."""
+            data.to_frame(name=name).to_feather(filepath)
+
+        @read_feather.register(pd.DataFrame)
+        def read_feather_pd1(data: pd.DataFrame, filepath: str) -> pd.DataFrame:
+            """Reads a dataframe from a feather file."""
+            return pd.read_feather(filepath)
+
+        @read_feather.register(pd.Series)
+        def read_feather_pd2(data: pd.Series, filepath: str) -> pd.Series:
+            """Reads a series from a feather file."""
+            _df = pd.read_feather(filepath)
+            return _df[_df.columns[0]]
+
+        @write_parquet.register(pd.DataFrame)
+        def write_parquet_pd1(data: pd.DataFrame, filepath: str, name: str) -> None:
+            """Writes a dataframe to a parquet file."""
+            data.to_parquet(filepath)
+
+        @write_parquet.register(pd.Series)
+        def write_parquet_pd2(data: pd.Series, filepath: str, name: str) -> None:
+            """Writes a data frame to a parquet file."""
+            data.to_frame(name=name).to_parquet(filepath)
+
+        @read_parquet.register(pd.DataFrame)
+        def read_parquet_pd1(data: pd.DataFrame, filepath: str) -> pd.DataFrame:
+            """Reads a dataframe from a parquet file."""
+            return pd.read_parquet(filepath)
+
+        @read_parquet.register(pd.Series)
+        def read_parquet_pd2(data: pd.Series, filepath: str) -> pd.Series:
+            """Reads a series from a parquet file."""
+            _df = pd.read_parquet(filepath)
+            return _df[_df.columns[0]]
+
+    except ImportError:
+        pass
+
+
+except ImportError:
+    pass
+
+
+@write_json.register(dict)
+def write_json_dict(data: dict, filepath: str, name: str) -> None:
+    """Writes a dictionary to a JSON file."""
+    if isinstance(data, dict):
+        with open(filepath, "w", encoding="utf8") as file:
+            json.dump(data, file)
+    else:
+        raise ValueError(f"Expected a dict, got {type(data)}")
+
+
+@read_json.register(dict)
+def read_json_dict(data: dict, filepath: str) -> dict:
+    """Reads a dictionary from a JSON file."""
+    with open(filepath, "r", encoding="utf8") as file:
+        return json.load(file)
+
+
+@write_pickle.register(object)
+def write_pickle_object(data: object, filepath: str, name: str) -> None:
+    if isinstance(data, object):
+        with open(filepath, "wb") as file:
+            pickle.dump(data, file)
+    else:
+        raise ValueError(f"Expected an object, got {type(data)}")
+
+
+@read_pickle.register(object)
+def read_pickle_object(data: object, filepath: str) -> object:
+    """Reads a pickle file"""
+    with open(filepath, "rb") as file:
+        return pickle.load(file)
+
+
+class CachingGraphAdapter(SimplePythonGraphAdapter):
+    """Caching adapter.
+
+    Any node with tag "cache" will be cached (or loaded from cache) in the format defined by the
+    tag's value. There are a handful of formats supported, and other formats' readers and writers
+    can be provided to the constructor.
+
+    Values are loaded from cache if the node's file exists, unless one of these is true:
+     * node is explicitly forced to be computed with a constructor argument,
+     * any of its (potentially transitive) dependencies that are configured to be cached
+       was nevertheless computed (either forced or missing cached file).
+
+    Custom Serializers
+    ------------------
+
+    One can provide custom readers and writers for any format by passing them to the constructor.
+    These readers and writers will override the default ones. If you don't want to override, but
+    rather extend the default ones, you can do so by registering them with the `register` method
+    on the appropriate function.
+
+    Writer functions need to have the following signature:
+    `def write_<format>(data: Any, filepath: str, name: str) -> None: ...`
+    where `data` is the data to be written, `filepath` is the path to the file to be written to,
+    and `name` is the name of the node that is being written.
+
+    Reader functions need to have the following signature:
+    `def read_<format>(data: Any, filepath: str) -> Any: ...`
+    where `data` is an EMPTY OBJECT of the type you wish to instantiate, and `filepath` is the
+    path to the file to be read from.
+
+    For example, if you want to extend JSON reader/writer to work with your custom type `T`,
+    you can do the following:
+
+    .. code-block:: python
+
+        @write_json.register(T)
+        def write_json_pd1(data: T, filepath: str, name: str) -> None: ...
+
+
+        @read_json.register(T)
+        def read_json_dict(data: T, filepath: str) -> T: ...
+
+    Usage
+    -----
+
+    This is a simple example of the usage of `CachingGraphAdapter`.
+
+    First, let's define some nodes in `nodes.py`:
+
+    .. code-block:: python
+
+        import pandas as pd
+        from hamilton.function_modifiers import tag
+
+
+        def data_a() -> pd.DataFrame: ...
+
+
+        @tag(cache="parquet")
+        def data_b() -> pd.DataFrame: ...
+
+
+        def transformed(data_a: pd.DataFrame, data_b: pd.DataFrame) -> pd.DataFrame: ...
+
+    Notice that `data_b` is configured to be cached in a parquet file.
+
+    We then simply initialize the driver with a caching adapter:
+
+    .. code-block:: python
+
+        from hamilton import base
+        from hamilton.driver import Driver
+        from hamilton.experimental import h_cache
+
+        import nodes
+
+        adapter = h_cache.CachingGraphAdapter(cache_path, base.PandasDataFrameResult())
+        dr = Driver(config, nodes, adapter=adapter)
+        result = dr.execute(["transformed"])
+
+        # Because `data_b` has been cached now, only `data_a` and `transformed` nodes
+        # will actually run.
+        result = dr.execute(["transformed"])
+    """
+
+    def __init__(
+        self,
+        cache_path: str,
+        *args,
+        force_compute: set[str] | None = None,
+        writers: dict[str, Callable[[Any, str, str], None]] | None = None,
+        readers: dict[str, Callable[[Any, str], Any]] | None = None,
+        **kwargs,
+    ):
+        """Constructs the adapter.
+
+        :param cache_path: Path to the directory where cached files are stored.
+        :param force_compute: Set of nodes that should be forced to compute even if cache exists.
+        :param writers: A dictionary of writers for custom formats.
+        :param readers: A dictionary of readers for custom formats.
+        """
+
+        super().__init__(*args, **kwargs)
+        self.cache_path = cache_path
+        self.force_compute = force_compute if force_compute is not None else {}
+        self.computed_nodes = set()
+
+        self.writers = writers or {}
+        self.readers = readers or {}
+
+        self._init_default_readers_writers()
+
+    def _init_default_readers_writers(self):
+        if "json" not in self.writers:
+            self.writers["json"] = write_json
+        if "json" not in self.readers:
+            self.readers["json"] = read_json
+
+        if "feather" not in self.writers:
+            self.writers["feather"] = write_feather
+        if "feather" not in self.readers:
+            self.readers["feather"] = read_feather
+
+        if "parquet" not in self.writers:
+            self.writers["parquet"] = write_parquet
+        if "parquet" not in self.readers:
+            self.readers["parquet"] = read_parquet
+
+        if "pickle" not in self.writers:
+            self.writers["pickle"] = write_pickle
+        if "pickle" not in self.readers:
+            self.readers["pickle"] = read_pickle
+
+    def _check_format(self, fmt):
+        if fmt not in self.writers:
+            raise ValueError(f"invalid cache format: {fmt}")
+
+    def _write_cache(self, fmt: str, data: Any, filepath: str, node_name: str) -> None:
+        self._check_format(fmt)
+        self.writers[fmt](data, filepath, node_name)
+
+    def _read_cache(self, fmt: str, expected_type: Any, filepath: str) -> None:
+        self._check_format(fmt)
+        return self.readers[fmt](expected_type, filepath)
+
+    def _get_empty_expected_type(self, expected_type: type) -> Any:
+        if typing_inspect.is_generic_type(expected_type):
+            return typing_inspect.get_origin(expected_type)()
+        return expected_type()  # This ASSUMES that we can just do `str()`, `pd.DataFrame()`, etc.
+
+    def execute_node(self, node: Node, kwargs: dict[str, Any]) -> Any:
+        """Executes nodes conditionally according to caching rules.
+
+        This node is executed if at least one of these is true:
+
+        * no cache is present,
+        * it is explicitly forced by passing it to the adapter in ``force_compute``,
+        * at least one of its upstream nodes that had a @cache annotation was computed,
+          either due to lack of cache or being explicitly forced.
+
+        """
+        cache_format = node.tags.get("cache")
+        implicitly_forced = any(dep.name in self.computed_nodes for dep in node.dependencies)
+        if cache_format is not None:
+            filepath = f"{self.cache_path}/{node.name}.{cache_format}"
+            explicitly_forced = node.name in self.force_compute
+            if explicitly_forced or implicitly_forced or not os.path.exists(filepath):
+                result = node.callable(**kwargs)
+                logger.debug(
+                    "Writing cache for %s to %s with type %s to %s",
+                    node.name,
+                    filepath,
+                    type(result),
+                    cache_format,
+                )
+                self._write_cache(cache_format, result, filepath, node.name)
+                self.computed_nodes.add(node.name)
+                return result
+            empty_expected_type = self._get_empty_expected_type(node.type)
+            logger.debug(
+                "Reading cache for %s from %s with type %s to %s",
+                node.name,
+                filepath,
+                type(empty_expected_type),
+                cache_format,
+            )
+            return self._read_cache(cache_format, empty_expected_type, filepath)
+
+        if implicitly_forced:
+            # For purposes of caching, we only mark it as computed if any cached input was computed.
+            # Otherwise, dependants would always be recomputed if they have a non-cached dependency.
+            self.computed_nodes.add(node.name)
+        return node.callable(**kwargs)
+
+    def build_result(self, **outputs: dict[str, Any]) -> Any:
+        """Clears the computed nodes information and delegates to the super class."""
+        self.computed_nodes = set()
+        return super().build_result(**outputs)