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

hamilton/plugins/h_threadpool.py

@@ -0,0 +1,125 @@
+# 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 collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import Any
+
+from hamilton import lifecycle, node
+from hamilton.lifecycle import base
+
+
+def _new_fn(fn: Callable, **fn_kwargs) -> Any:
+    """Function that runs in the thread.
+
+    It can recursively check for Futures because we don't have to worry about
+    process serialization.
+    :param fn: Function to run
+    :param fn_kwargs: Keyword arguments to pass to the function
+    """
+    for k, v in fn_kwargs.items():
+        if isinstance(v, Future):
+            while isinstance(v, Future):
+                v = v.result()  # this blocks until the future is resolved
+            fn_kwargs[k] = v
+    # execute the function once all the futures are resolved
+    return fn(**fn_kwargs)
+
+
+class FutureAdapter(base.BaseDoRemoteExecute, lifecycle.ResultBuilder):
+    """Adapter that lazily submits each function for execution to a ThreadpoolExecutor.
+
+    This adapter has similar behavior to the async Hamilton driver which allows for parallel execution of functions.
+
+    This adapter works because we don't have to worry about object serialization.
+
+    Caveats:
+    - DAGs with lots of CPU intense functions will limit usefulness of this adapter, unless they release the GIL.
+    - DAGs with lots of I/O bound work will benefit from this adapter, e.g. making API calls.
+    - The max parallelism is limited by the number of threads in the ThreadPoolExecutor.
+
+    Unsupported behavior:
+    - The FutureAdapter does not support DAGs with Parallelizable & Collect functions. This is due to laziness
+    rather than anything inherently technical. If you'd like this feature, please open an issue on the Hamilton
+    repository.
+
+    """
+
+    def __init__(
+        self,
+        max_workers: int = None,
+        thread_name_prefix: str = "",
+        result_builder: lifecycle.ResultBuilder = None,
+    ):
+        """Constructor.
+        :param max_workers: The maximum number of threads that can be used to execute the given calls.
+        :param thread_name_prefix: An optional name prefix to give our threads.
+        :param result_builder: Optional. Result builder to use for building the result.
+        """
+        self.executor = ThreadPoolExecutor(
+            max_workers=max_workers, thread_name_prefix=thread_name_prefix
+        )
+        self.result_builder = result_builder
+
+    def input_types(self) -> list[type[type]]:
+        """Gives the applicable types to this result builder.
+        This is optional for backwards compatibility, but is recommended.
+
+        :return: A list of types that this can apply to.
+        """
+        # since this wraps a potential result builder, expose the input types of the wrapped
+        # result builder doesn't make sense.
+        return [Any]
+
+    def output_type(self) -> type:
+        """Returns the output type of this result builder
+        :return: the type that this creates
+        """
+        if self.result_builder:
+            return self.result_builder.output_type()
+        return Any
+
+    def do_remote_execute(
+        self,
+        *,
+        execute_lifecycle_for_node: Callable,
+        node: node.Node,
+        **kwargs: dict[str, Any],
+    ) -> Any:
+        """Function that submits the passed in function to the ThreadPoolExecutor to be executed
+        after wrapping it with the _new_fn function.
+
+        :param node: Node that is being executed
+        :param execute_lifecycle_for_node: Function executing lifecycle_hooks and lifecycle_methods
+        :param kwargs: Keyword arguments that are being passed into the function
+        """
+        return self.executor.submit(_new_fn, execute_lifecycle_for_node, **kwargs)
+
+    def build_result(self, **outputs: Any) -> Any:
+        """Given a set of outputs, build the result.
+
+        This function will block until all futures are resolved.
+
+        :param outputs: the outputs from the execution of the graph.
+        :return: the result of the execution of the graph.
+        """
+        for k, v in outputs.items():
+            if isinstance(v, Future):
+                outputs[k] = v.result()
+        if self.result_builder:
+            return self.result_builder.build_result(**outputs)
+        return outputs

hamilton/plugins/h_tqdm.py

@@ -0,0 +1,122 @@
+# 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 collections.abc import Collection
+from typing import Any
+
+import tqdm
+
+from hamilton import graph_types
+from hamilton.lifecycle import GraphExecutionHook, NodeExecutionHook
+
+
+class ProgressBar(
+    GraphExecutionHook,
+    NodeExecutionHook,
+):
+    """An adapter that uses tqdm to show progress bars for the graph execution.
+
+    Note: you need to have tqdm installed for this to work.
+    If you don't have it installed, you can install it with `pip install tqdm`
+    (or `pip install sf-hamilton[tqdm]` -- use quotes if you're using zsh).
+
+    .. code-block:: python
+
+        from hamilton.plugins import h_tqdm
+
+        dr = (
+            driver.Builder()
+            .with_config({})
+            .with_modules(some_modules)
+            .with_adapters(h_tqdm.ProgressBar(desc="DAG-NAME"))
+            .build()
+        )
+        # and then when you call .execute() or .materialize() you'll get a progress bar!
+    """
+
+    def __init__(self, desc: str = "Graph execution", max_node_name_width: int = 50, **kwargs):
+        """Create a new Progress Bar adapter.
+
+        :param desc: The description to show in the progress bar. E.g. DAG Name is a good choice.
+        :param kwargs: Additional kwargs to pass to TQDM. See TQDM docs for more info.
+        :param node_name_target_width: the target width for the node name so that the progress bar is consistent. If this is None, it will take the longest, until it hits max_node_name_width.
+
+        """
+        self.desc = desc
+        self.kwargs = kwargs
+        self.node_name_target_width = (
+            None  # what we target padding for -- starts at None as we adjust.
+        )
+        self.max_node_name_width = max_node_name_width  # what we cap the padding at.
+        self.progress_bar = None
+
+    def _get_node_name_display(self, node_name: str) -> str:
+        """Gives the node name display given a max width and a node name. Max width could be DAG-dependent."""
+        out = (
+            node_name
+            if len(node_name) <= self.node_name_target_width
+            else node_name[: self.node_name_target_width - 3] + "..."
+        )
+        if len(out) < self.node_name_target_width:
+            out += " " * (self.node_name_target_width - len(out))
+        return out
+
+    def run_before_graph_execution(
+        self,
+        *,
+        graph: graph_types.HamiltonGraph,
+        final_vars: list[str],
+        inputs: dict[str, Any],
+        overrides: dict[str, Any],
+        execution_path: Collection[str],
+        **future_kwargs: Any,
+    ):
+        total_node_to_execute = len(execution_path)
+        max_node_name_length = min(
+            max([len(node) for node in execution_path]), self.max_node_name_width
+        )
+        if self.node_name_target_width is None:
+            self.node_name_target_width = max_node_name_length
+        self.progress_bar = tqdm.tqdm(
+            desc=self.desc, unit="funcs", total=total_node_to_execute, **self.kwargs
+        )
+
+    def run_before_node_execution(
+        self,
+        *,
+        node_name: str,
+        node_tags: dict[str, Any],
+        node_kwargs: dict[str, Any],
+        node_return_type: type,
+        task_id: str | None,
+        **future_kwargs: Any,
+    ):
+        name_display = self._get_node_name_display(node_name)
+        self.progress_bar.set_description_str(f"{self.desc} -> {name_display}")
+
+    def run_after_node_execution(self, **future_kwargs):
+        self.progress_bar.update(1)
+
+    def run_after_graph_execution(self, **future_kwargs):
+        name_part = "Execution Complete!"
+        if len(name_part) > self.node_name_target_width:
+            padding = ""
+        else:
+            padding = " " * (self.node_name_target_width - len(name_part))
+        self.progress_bar.set_description_str(f"{self.desc} -> {name_part + padding}")
+        self.progress_bar.set_postfix({})
+        self.progress_bar.close()

hamilton/plugins/h_vaex.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.
+
+from typing import Any
+
+import numpy as np
+import pandas as pd
+
+from hamilton import base
+
+try:
+    import vaex
+except ImportError as e:
+    raise NotImplementedError("Vaex is not installed.") from e
+
+
+class VaexDataFrameResult(base.ResultMixin):
+    """A ResultBuilder that produces a Vaex dataframe.
+
+    Use this when you want to create a Vaex dataframe from the outputs.
+    Caveat: you need to ensure that the length
+    of the outputs is the same (except scalars), otherwise you will get an error;
+    mixed outputs aren't that well handled.
+
+    To use:
+
+    .. code-block:: python
+
+        from hamilton import base, driver
+        from hamilton.plugins import h_vaex
+
+        vaex_builder = h_vaex.VaexDataFrameResult()
+        adapter = base.SimplePythonGraphAdapter(vaex_builder)
+        dr = driver.Driver(config, *modules, adapter=adapter)
+        df = dr.execute([...], inputs=...)  # returns vaex dataframe
+
+    Note: this is just a first attempt at something for Vaex.
+    Think it should handle more? Come chat/open a PR!
+    """
+
+    def build_result(
+        self,
+        **outputs: dict[str, vaex.expression.Expression | vaex.dataframe.DataFrame | Any],
+    ):
+        """This is the method that Hamilton will call to build the final result.
+        It will pass in the results of the requested outputs that
+        you passed in to the execute() method.
+
+        :param outputs: The results of the requested outputs.
+        :return: a Vaex DataFrame.
+        """
+
+        # We split all outputs into DataFrames, arrays and scalars
+        dfs: list[vaex.dataframe.DataFrame] = []  # Vaex DataFrames from outputs
+        arrays: dict[str, np.ndarray] = dict()  # arrays from outputs
+        scalars: dict[str, Any] = dict()  # scalars from outputs
+
+        for name, value in outputs.items():
+            if isinstance(value, vaex.dataframe.DataFrame):
+                dfs.append(value)
+            elif isinstance(value, vaex.expression.Expression):
+                nparray = value.to_numpy()
+                if nparray.ndim == 0:  # value is scalar
+                    scalars[name] = nparray.item()
+                elif nparray.shape == (1,):  # value is scalar
+                    scalars[name] = nparray[0]
+                else:  # value is array
+                    arrays[name] = nparray
+            elif isinstance(value, np.ndarray):
+                if value.ndim == 0:  # value is scalar
+                    scalars[name] = value.item()
+                elif value.shape == (1,):  # value is scalar
+                    scalars[name] = value[0]
+                else:  # value is array
+                    arrays[name] = value
+            elif pd.api.types.is_scalar(value):  # value is scalar
+                scalars[name] = value
+            else:
+                value_type = str(type(value))
+                message = f"VaexDataFrameResult doesn't support {value_type}"
+                raise NotImplementedError(message)
+
+        df = None
+
+        if arrays:
+            # Check if all arrays have correct and identical shapes.
+            first_expression_shape = next(arrays.values().__iter__()).shape
+            if len(first_expression_shape) > 1:
+                raise NotImplementedError(
+                    "VaexDataFrameResult supports only one-dimensional Expression results"
+                )
+            for _name, a in arrays.items():
+                if a.shape != first_expression_shape:
+                    raise NotImplementedError(
+                        "VaexDataFrameResult supports Expression results with same dimension only"
+                    )
+
+            # All scalars become arrays with the same shape as other arrays.
+            for name, scalar in scalars.items():
+                arrays[name] = np.full(first_expression_shape, scalar)
+
+            df = vaex.from_arrays(**arrays)
+
+        elif scalars:
+            # There are not arrays in outputs,
+            # so we construct Vaex DataFrame with one row consisting of scalars.
+            df = vaex.from_arrays(**{name: np.array([value]) for name, value in scalars.items()})
+
+        if df:
+            dfs.append(df)
+
+        return vaex.concat(dfs)
+
+    def output_type(self) -> type:
+        return vaex.dataframe.DataFrame

hamilton/plugins/huggingface_extensions.py

@@ -0,0 +1,236 @@
+# 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, Mapping, Sequence
+from os import PathLike
+from typing import (
+    Any,
+    BinaryIO,
+)
+
+try:
+    from datasets import (
+        Dataset,
+        DatasetDict,
+        DownloadConfig,
+        DownloadMode,
+        Features,
+        IterableDataset,
+        IterableDatasetDict,
+        VerificationMode,
+        Version,
+        load_dataset,
+    )
+    from datasets.formatting.formatting import LazyBatch
+except ImportError as e:
+    raise NotImplementedError("huggingface datasets library is not installed.") from e
+
+try:
+    import lancedb
+    from lancedb import table  # noqa: F401
+except ImportError:
+    lancedb = None
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+COLUMN_FRIENDLY_DF_TYPE = False
+
+HF_types = (DatasetDict, Dataset, IterableDatasetDict, IterableDataset)
+
+
+@dataclasses.dataclass
+class HuggingFaceDSLoader(DataLoader):
+    """Data loader for hugging face datasets. Uses load_data method."""
+
+    path: str
+    dataset_name: str | None = None  # this can't be `name` because it clashes with `.name()`
+    data_dir: str | None = None
+    data_files: str | Sequence[str] | Mapping[str, str | Sequence[str]] | None = None
+    split: str | None = None
+    cache_dir: str | None = None
+    features: Features | None = None
+    download_config: DownloadConfig | None = None
+    download_mode: DownloadMode | str | None = None
+    verification_mode: VerificationMode | str | None = None
+    ignore_verifications = "deprecated"
+    keep_in_memory: bool | None = None
+    save_infos: bool = False
+    revision: str | Version | None = None
+    token: bool | str | None = None
+    use_auth_token = "deprecated"
+    task = "deprecated"
+    streaming: bool = False
+    num_proc: int | None = None
+    storage_options: dict | None = None
+    config_kwargs: dict | None = None
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return list(HF_types)
+
+    def _get_loading_kwargs(self) -> dict:
+        # Puts kwargs in a dict
+        kwargs = dataclasses.asdict(self)
+        # we send path separately
+        del kwargs["path"]
+        config_kwargs: dict | None = kwargs.pop("config_kwargs", None)
+        if config_kwargs:
+            # add config kwargs as needed.
+            kwargs.update(config_kwargs)
+
+        # need to pass in name
+        kwargs["name"] = kwargs.pop("dataset_name", None)
+
+        return kwargs
+
+    def load_data(self, type_: type) -> tuple[HF_types, dict[str, Any]]:
+        """Loads the data set given the path and class values."""
+        ds = load_dataset(self.path, **self._get_loading_kwargs())
+        is_dataset = isinstance(ds, Dataset)
+        f_meta = {"path": self.path}
+        ds_meta = {"rows": ds.num_rows, "columns": ds.column_names}
+        if is_dataset:
+            ds_meta["size_in_bytes"] = ds.size_in_bytes
+            ds_meta["features"] = ds.features.to_dict()
+        return ds, {"file_metadata": f_meta, "dataset_metadata": ds_meta}
+
+    @classmethod
+    def name(cls) -> str:
+        return "hf_dataset"
+
+
+@dataclasses.dataclass
+class HuggingFaceDSParquetSaver(DataSaver):
+    """Saves a Huggingface dataset to parquet."""
+
+    path_or_buf: PathLike | BinaryIO
+    batch_size: int | None = None
+    parquet_writer_kwargs: dict | None = None
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return list(HF_types)
+
+    @classmethod
+    def applies_to(cls, type_: type[type]) -> bool:
+        return type_ in HF_types
+
+    def _get_saving_kwargs(self) -> dict:
+        # Puts kwargs in a dict
+        kwargs = dataclasses.asdict(self)
+        # we put path_or_buff as a positional argument
+        del kwargs["path_or_buf"]
+        parquet_writer_kwargs: dict | None = kwargs.pop("parquet_writer_kwargs", None)
+        if parquet_writer_kwargs:
+            # add config kwargs as needed.
+            kwargs.update(parquet_writer_kwargs)
+
+        return kwargs
+
+    def save_data(self, ds: HF_types) -> dict[str, Any]:
+        """Saves the data to parquet."""
+        is_dataset = isinstance(ds, Dataset)
+        ds.to_parquet(self.path_or_buf, **self._get_saving_kwargs())
+        ds_meta = {
+            "rows": ds.num_rows,
+            "columns": ds.column_names,
+        }
+        if is_dataset:
+            ds_meta.update({"size_in_bytes": ds.size_in_bytes, "features": ds.features.to_dict()})
+        if isinstance(self.path_or_buf, BinaryIO):
+            f_meta = {}
+        else:
+            f_meta = (utils.get_file_metadata(self.path_or_buf),)
+        return {"file_metadata": f_meta, "dataset_metadata": ds_meta}
+
+    @classmethod
+    def name(cls) -> str:
+        return "parquet"
+
+
+# we do this here just in case lancedb is not installed.
+if lancedb is not None:
+
+    def _batch_write(
+        dataset_batch: LazyBatch, db: lancedb.DBConnection, table_name: str, columns: str
+    ) -> None:
+        """Helper function to batch write to lancedb."""
+        if columns is None:
+            data = dataset_batch.pa_table
+        else:
+            data = dataset_batch.pa_table.select(columns)
+        try:
+            db.create_table(table_name, data)
+        except (OSError, ValueError):
+            tbl = db.open_table(table_name)
+            tbl.add(data)
+        return None
+
+    @dataclasses.dataclass
+    class HuggingFaceDSLanceDBSaver(DataSaver):
+        """Data saver that saves Huggingface datasets to lancedb."""
+
+        db_client: lancedb.DBConnection
+        table_name: str
+        columns_to_write: list[str] = None  # None means all.
+        write_batch_size: int = 100
+
+        @classmethod
+        def applicable_types(cls) -> Collection[type]:
+            return list(HF_types)
+
+        def save_data(self, ds: HF_types) -> dict[str, Any]:
+            """This batches writes to lancedb."""
+            ds.map(
+                _batch_write,
+                batched=True,
+                batch_size=self.write_batch_size,
+                fn_kwargs={
+                    "db": self.db_client,
+                    "table_name": self.table_name,
+                    "columns": self.columns_to_write,
+                },
+                desc=f"writing to lancedb table {self.table_name}",
+            )
+            is_dataset = isinstance(ds, Dataset)
+            ds_meta = {
+                "rows": ds.num_rows,
+                "columns": ds.column_names,
+            }
+            if is_dataset:
+                ds_meta.update(
+                    {"size_in_bytes": ds.size_in_bytes, "features": ds.features.to_dict()}
+                )
+            return {"db_meta": {"table_name": self.table_name}, "dataset_metadata": ds_meta}
+
+        @classmethod
+        def name(cls) -> str:
+            return "lancedb"
+
+
+def register_data_loaders_savers():
+    loaders = [HuggingFaceDSLoader, HuggingFaceDSParquetSaver]
+    if lancedb:
+        loaders.append(HuggingFaceDSLanceDBSaver)
+    for loader in loaders:
+        registry.register_adapter(loader)
+
+
+register_data_loaders_savers()