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

hamilton/plugins/dask_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 dask.dataframe as dd
+except ImportError as e:
+    raise NotImplementedError("Dask is not installed.") from e
+
+from hamilton import registry
+
+DATAFRAME_TYPE = dd.DataFrame
+COLUMN_TYPE = dd.Series
+
+
+@registry.get_column.register(dd.DataFrame)
+def get_column_dask(df: dd.DataFrame, column_name: str) -> dd.Series:
+    return df[column_name]
+
+
+@registry.fill_with_scalar.register(dd.DataFrame)
+def fill_with_scalar_dask(df: dd.DataFrame, column_name: str, value: Any) -> dd.DataFrame:
+    df[column_name] = value
+    return df
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("dask", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+register_types()

hamilton/plugins/dlt_extensions.py

@@ -0,0 +1,161 @@
+# 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, Iterable, Sequence
+from typing import Any, Literal
+
+try:
+    import dlt
+    from dlt.common.destination.capabilities import TLoaderFileFormat
+    from dlt.common.schema import Schema, TColumnSchema
+
+    # importing TDestinationReferenceArg fails if Destination isn't imported
+    from dlt.extract.resource import DltResource
+
+except ImportError as e:
+    # raise import error first
+    raise ImportError(f"Failed to import the DLT library. {e}") from e
+except Exception as e:
+    # raise import error with custom message
+    raise ImportError("Failed to import the DLT library.") from e
+
+import pandas as pd
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+DATAFRAME_TYPES = [Iterable, pd.DataFrame]
+
+# TODO add types for other Dataframe libraries
+try:
+    import pyarrow as pa
+
+    DATAFRAME_TYPES.extend([pa.Table, pa.RecordBatch])
+except ModuleNotFoundError:
+    pass
+
+# convert to tuple to dynamically define type `Union[DATAFRAME_TYPES]`
+DATAFRAME_TYPES = tuple(DATAFRAME_TYPES)
+COLUMN_FRIENDLY_DF_TYPE = False
+
+
+@dataclasses.dataclass
+class DltResourceLoader(DataLoader):
+    resource: DltResource
+
+    @classmethod
+    def name(cls) -> str:
+        return "dlt"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [pd.DataFrame]
+
+    def load_data(self, type_: type) -> tuple[pd.DataFrame, dict[str, Any]]:
+        """Creates a pipeline and conduct `extract` and `normalize` steps.
+        Then, "load packages" are read with pandas
+        """
+        pipeline = dlt.pipeline(
+            pipeline_name="Hamilton-DltResourceLoader", destination="filesystem"
+        )
+        pipeline.extract(self.resource, loader_file_format="parquet")
+        normalize_info = pipeline.normalize()
+
+        partition_file_paths = []
+        package = normalize_info.load_packages[0]
+        for job in package.jobs["new_jobs"]:
+            if job.job_file_info.table_name == self.resource.name:
+                partition_file_paths.append(job.file_path)
+
+        # TODO use pyarrow directly to support different dataframe libraries
+        # ref: https://github.com/dlt-hub/verified-sources/blob/master/sources/filesystem/readers.py
+        # ref: https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetDataset.html#pyarrow.parquet.ParquetDataset
+        df = pd.concat([pd.read_parquet(f) for f in partition_file_paths], ignore_index=True)
+
+        # delete the pipeline
+        pipeline.drop()
+
+        metadata = utils.get_dataframe_metadata(df)
+        return df, metadata
+
+
+# TODO handle behavior with `combine=`, currently only supports materializing a single node
+@dataclasses.dataclass
+class DltDestinationSaver(DataSaver):
+    """Materialize results using a dlt pipeline with the specified destination.
+
+    In reference to an Extract, Transform, Load (ETL) pipeline, here, the Hamilton
+    dataflow is responsible for Transform, and `DltDestination` for Load.
+    """
+
+    pipeline: dlt.Pipeline
+    table_name: str
+    primary_key: str | None = None
+    write_disposition: Literal["skip", "append", "replace", "merge"] | None = None
+    columns: Sequence[TColumnSchema] | None = None
+    schema: Schema | None = None
+    loader_file_format: TLoaderFileFormat | None = None
+
+    @classmethod
+    def name(cls) -> str:
+        return "dlt"
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return DATAFRAME_TYPES
+
+    def _get_kwargs(self) -> dict:
+        kwargs = {}
+        fields_to_skip = ["pipeline"]
+        for field in dataclasses.fields(self):
+            field_value = getattr(self, field.name)
+            if field.name in fields_to_skip:
+                continue
+
+            if field_value != field.default:
+                kwargs[field.name] = field_value
+
+        return kwargs
+
+    # TODO get pyarrow table from polars, dask, etc.
+    def save_data(self, data) -> dict[str, Any]:
+        """
+        ref: https://dlthub.com/docs/dlt-ecosystem/verified-sources/arrow-pandas
+        """
+        if isinstance(data, dict):
+            raise NotImplementedError(
+                "DltDestinationSaver received data of type `dict`."
+                "Currently, it doesn't support specifying `combine=base.DictResult()`"
+            )
+
+        load_info = self.pipeline.run(data, **self._get_kwargs())
+        # follows the pattern of metadata output found in hamilton.io.utils
+        return {"dlt_metadata": load_info.asdict()}
+
+
+def register_data_loaders():
+    """Function to register the data loaders for this extension."""
+    for loader in [
+        DltDestinationSaver,
+        DltResourceLoader,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()

hamilton/plugins/geopandas_extensions.py

@@ -0,0 +1,49 @@
+# 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 geopandas as gpd
+except ImportError as e:
+    raise NotImplementedError("geopandas is not installed.") from e
+
+from hamilton import registry
+
+DATAFRAME_TYPE = gpd.GeoDataFrame
+COLUMN_TYPE = gpd.GeoSeries
+
+
+@registry.get_column.register(gpd.GeoDataFrame)
+def get_column_geopandas(df: gpd.GeoDataFrame, column_name: str) -> gpd.GeoSeries:
+    return df[column_name]
+
+
+@registry.fill_with_scalar.register(gpd.GeoDataFrame)
+def fill_with_scalar_geopandas(
+    df: gpd.GeoDataFrame, column_name: str, value: Any
+) -> gpd.GeoDataFrame:
+    df[column_name] = value
+    return df
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("geopandas", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+register_types()

hamilton/plugins/h_dask.py

@@ -0,0 +1,331 @@
+# 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 typing
+
+import dask.array
+import dask.dataframe
+import numpy as np
+import pandas as pd
+from dask import compute
+from dask.base import tokenize
+from dask.delayed import Delayed, delayed
+from dask.distributed import Client as DaskClient
+
+from hamilton import base, htypes, node
+from hamilton.execution import executors
+
+logger = logging.getLogger(__name__)
+
+try:
+    from dask.dataframe.dask_expr import Scalar as dask_scalar
+except ImportError:
+    # this is for older versions of dask
+    from dask.dataframe.core import Scalar as dask_scalar
+
+
+class DaskGraphAdapter(base.HamiltonGraphAdapter):
+    """Class representing what's required to make Hamilton run on Dask.
+
+    This walks the graph and translates it to run onto `Dask <https://dask.org/>`__.
+
+    Use `pip install sf-hamilton[dask]` to get the dependencies required to run this.
+
+    Try this adapter when:
+
+        1. Dask is a good choice to scale computation when you really can't do things in memory anymore with pandas. \
+        For most simple pandas operations, you should not have to do anything to scale! You just need to load in \
+        data via dask rather than pandas.
+        2. Dask can help scale to larger data sets if running on a cluster -- you'll just have to switch to\
+        natively using their object types if that's the case (set use_delayed=False, and compute_at_end=False).
+        3. Use this adapter if you want to utilize multiple cores on a single machine, or you want to scale to large \
+        data set sizes with a Dask cluster that you can connect to.
+        4. The ONLY CAVEAT really is whether you use `delayed` or `dask datatypes` (or both).
+
+    Please read the following notes about its limitations.
+
+    Notes on scaling:
+    -----------------
+      - Multi-core on single machine ✅
+      - Distributed computation on a Dask cluster ✅
+      - Scales to any size of data supported by Dask ✅; assuming you load it appropriately via Dask loaders.
+      - Works best with Pandas 2.0+ and pyarrow backend.
+
+    Function return object types supported:
+    ---------------------------------------
+      - Works for any python object that can be serialized by the Dask framework. ✅
+
+    Pandas?
+    -------
+    Dask implements a good subset of the Pandas API:
+      - You might be able to get away with scaling without having to change your code at all!
+      - See https://docs.dask.org/en/latest/dataframe-api.html for Pandas supported APIs.
+      - If it is not supported by their API, you have to then read up and think about how to structure you hamilton\
+      function computation -- https://docs.dask.org/en/latest/dataframe.html
+      - if paired with DaskDataFrameResult & use_delayed=False & compute_at_end=False, it will help you produce a \
+      dask dataframe as a result that you can then convert back to pandas if you want.
+
+    Loading Data:
+    -------------
+      - see https://docs.dask.org/en/latest/best-practices.html#load-data-with-dask.
+      - we recommend creating a python module specifically encapsulating functions that help you load data.
+
+    CAVEATS with use_delayed=True:
+    ------------------------------
+      - If using `use_delayed=True` serialization costs can outweigh the benefits of parallelism, so you should \
+       benchmark your code to see if it's worth it.
+      - With this adapter & use_delayed=True, it can naively wrap all your functions with `delayed`, which will mean \
+       they will be executed and scheduled across the dask workers. This is a good choice if your computation is slow, \
+       or Hamilton graph is highly parallelizable.
+
+    DISCLAIMER -- this class is experimental, so signature changes are a possibility! But we'll aim to be backwards
+    compatible where possible.
+    """
+
+    def __init__(
+        self,
+        dask_client: DaskClient,
+        result_builder: base.ResultMixin = None,
+        visualize_kwargs: dict = None,
+        use_delayed: bool = True,
+        compute_at_end: bool = True,
+    ):
+        """Constructor
+
+        You have the ability to pass in a ResultMixin object to the constructor to control the return type that gets\
+        produced by running on Dask.
+
+        :param dask_client: the dask client -- we don't do anything with it, but thought that it would be useful\
+            to wire through here.
+        :param result_builder: The function that will build the result. Optional, defaults to pandas dataframe.
+        :param visualize_kwargs: Arguments to visualize the graph using dask's internals.\
+            **None**, means no visualization.\
+            **Dict**, means visualize -- see https://docs.dask.org/en/latest/api.html?highlight=visualize#dask.visualize\
+            for what to pass in.
+        :param use_delayed: Default is True for backwards compatibility. Whether to use dask.delayed to wrap every
+            function. Note: it is probably not necessary to mix this with using dask objects, e.g. dataframes/series.
+            They are by nature lazily computed and operate over the dask data types, so you don't need to wrap them
+            with delayed. Use delayed if you want to farm out computation.
+        :param compute_at_end: Default is True for backwards compatibility. Whether to compute() at the end.
+            That is, should `.compute()` be called in the result builder to quick off computation.
+        """
+        self.client = dask_client
+        self.result_builder = result_builder or base.PandasDataFrameResult()
+        self.visualize_kwargs = visualize_kwargs
+        self.use_delayed = use_delayed
+        self.compute_at_end = compute_at_end
+
+    @staticmethod
+    def check_input_type(node_type: type, input_value: typing.Any) -> bool:
+        # NOTE: the type of dask Delayed is unknown until they are computed
+        if isinstance(input_value, Delayed):
+            return True
+        elif node_type == pd.Series and isinstance(input_value, dask.dataframe.Series):
+            return True
+        elif node_type == np.array and isinstance(input_value, dask.array.Array):
+            return True
+        return htypes.check_input_type(node_type, input_value)
+
+    @staticmethod
+    def check_node_type_equivalence(node_type: type, input_type: type) -> bool:
+        if node_type == dask.array.Array and input_type == pd.Series:
+            return True
+        elif node_type == dask.dataframe.Series and input_type == pd.Series:
+            return True
+        return node_type == input_type
+
+    def execute_node(self, node: node.Node, kwargs: dict[str, typing.Any]) -> typing.Any:
+        """Function that is called as we walk the graph to determine how to execute a hamilton function.
+
+        :param node: the node from the graph.
+        :param kwargs: the arguments that should be passed to it.
+        :return: returns a dask delayed object.
+        """
+        if not self.use_delayed:
+            return node.callable(**kwargs)
+        # we want to ensure the name in dask corresponds to the node name, and not the wrapped
+        # function name that hamilton might have wrapped it with.
+        hash = tokenize(kwargs)  # this is what the dask docs recommend.
+        name = node.name + hash
+        dask_key_name = str(node.name) + "_" + hash
+        return delayed(node.callable, name=name)(
+            **kwargs,
+            dask_key_name=dask_key_name,  # this is what shows up in the dask console
+        )
+
+    def build_result(self, **outputs: dict[str, typing.Any]) -> typing.Any:
+        """Builds the result and brings it back to this running process.
+
+        :param outputs: the dictionary of key -> Union[delayed object reference | value]
+        :return: The type of object returned by self.result_builder. Note the following behaviors:
+          - if you use_delayed=True, then the result will be a delayed object.
+          - if you use_delayed=True & computed_at_end=True, then the result will be the return type
+          of self.result_builder.
+          - if you use_delayed=False & computed_at_end=True, this will only work if the self.result_builder
+          returns a dask type, as we will try to compute it.
+          - if you use_delayed=False & computed_at_end=False, this will return the result of self.result_builder.
+        """
+        if logger.isEnabledFor(logging.DEBUG):
+            for k, v in outputs.items():
+                logger.debug(f"Got column {k}, with type [{type(v)}].")
+        if self.use_delayed:
+            delayed_result = delayed(self.result_builder.build_result)(**outputs)
+        else:
+            delayed_result = self.result_builder.build_result(**outputs)
+        if self.visualize_kwargs is not None:
+            delayed_result.visualize(**self.visualize_kwargs)
+        if self.compute_at_end:
+            (result,) = compute(delayed_result)
+            return result
+        else:
+            return delayed_result
+
+
+class DaskDataFrameResult(base.ResultMixin):
+    @staticmethod
+    def build_result(**outputs: dict[str, typing.Any]) -> typing.Any:
+        """Builds a dask dataframe from the outputs.
+
+        This has some assumptions:
+         1. the order specified in the output will mirror the order of "joins" here.
+         2. it tries to massage types into dask types where it can
+         3. otherwise it duplicates any "scalars/objects" using the first valid input with an index as the
+            template. It assumes a single partition.
+        """
+
+        def get_output_name(output_name: str, column_name: str) -> str:
+            """Add function prefix to columns.
+            Note this means that they stop being valid python identifiers due to the `.` in the string.
+            """
+            return f"{output_name}.{column_name}"
+
+        if len(outputs) == 0:
+            raise ValueError("No outputs were specified. Cannot build a dataframe.")
+        if logger.isEnabledFor(logging.DEBUG):
+            for k, v in outputs.items():
+                logger.debug(f"Got column {k}, with type [{type(v)}].")
+
+        length = 0
+        index = None
+        massaged_outputs = {}
+        columns_expected = []
+        for k, v in outputs.items():
+            if isinstance(v, (dask.dataframe.Series, dask.dataframe.DataFrame)):
+                if length == 0:
+                    length = len(v)
+                    index = v.index
+                massaged_outputs[k] = v
+                if isinstance(v, dask.dataframe.Series):
+                    columns_expected.append(k)
+                else:
+                    columns_expected.extend([get_output_name(k, v_col) for v_col in v.columns])
+            elif isinstance(v, (pd.Series, pd.DataFrame)):
+                converted = dask.dataframe.from_pandas(v, npartitions=1)
+                massaged_outputs[k] = converted
+                if length == 0:
+                    length = len(converted)
+                    index = converted.index
+                if isinstance(v, pd.Series):
+                    columns_expected.append(k)
+                else:
+                    columns_expected.extend([get_output_name(k, v_col) for v_col in v.columns])
+            elif isinstance(v, (np.ndarray, np.generic)):
+                massaged_outputs[k] = dask.dataframe.from_array(v)
+                columns_expected.append(k)
+            elif isinstance(v, (list, tuple)):
+                massaged_outputs[k] = dask.dataframe.from_array(dask.array.from_array(v))
+                columns_expected.append(k)
+            elif isinstance(v, (dask_scalar,)):
+                scalar = v.compute()
+                if length == 0:
+                    massaged_outputs[k] = dask.dataframe.from_pandas(
+                        pd.DataFrame([scalar], index=[0]), npartitions=1
+                    )
+                else:
+                    massaged_outputs[k] = dask.dataframe.from_pandas(
+                        pd.DataFrame([scalar] * length, index=index), npartitions=1
+                    )
+                columns_expected.append(k)
+            elif isinstance(v, (int, float, str, bool, object)):
+                scalar = v
+                if length == 0:
+                    massaged_outputs[k] = dask.dataframe.from_pandas(
+                        pd.DataFrame([scalar], index=[0]), npartitions=1
+                    )
+                else:
+                    massaged_outputs[k] = dask.dataframe.from_pandas(
+                        pd.DataFrame([scalar] * length, index=index), npartitions=1
+                    )
+                columns_expected.append(k)
+            else:
+                raise ValueError(
+                    f"Unknown type {type(v)} for output {k}. "
+                    f"Do not know how to handle making a dataframe from this."
+                )
+
+        # assumption is that everything here is a dask series or dataframe
+        # we assume that we do column concatenation and that it's an outer join (TBD: make this configurable)
+        _df = dask.dataframe.concat([o for o in massaged_outputs.values()], axis=1, join="outer")
+        _df.columns = columns_expected
+        return _df
+
+
+# TODO: add ResultMixins for dask types
+
+
+class DaskExecutor(executors.TaskExecutor):
+    """A DaskExecutor for task-based execution on dask in the new Hamilton execution API."""
+
+    def __init__(self, *, client: DaskClient):
+        """Initializes the DaskExecutor. Note this currently takes in the client -- we will likely
+        add the ability to make it take in parameters to instantiate/tear down a client on its own.
+        This just allows full flexibility for now.
+
+        """
+        self.client = client
+
+    def init(self):
+        """No-op -- client already passed in by the user."""
+        pass
+
+    def finalize(self):
+        """No-op -- client already passed in by the user, who is responsible for shutting it
+        down."""
+        pass
+
+    def submit_task(self, task: executors.TaskImplementation) -> executors.TaskFuture:
+        """Submits a task using the dask futures API. Note that we are not using dask delayed --
+        as the idea is that tasks are potentially dynamic, meaning that we have to resolve some
+        before we create others. That makes the delayed API a little messier -- we would have to
+        call .compute() at certain steps. We *may* consider doing this, but for now, we are just
+        utilizing the futures API, and grouping it into tasks.
+
+        :param task: Task to execute (contains all arguments necessary)
+        :return: The future for the task
+        """
+
+        return executors.TaskFutureWrappingPythonFuture(
+            self.client.submit(executors.base_execute_task, task)
+        )
+
+    def can_submit_task(self) -> bool:
+        """For now we always can -- it will block on the dask side.
+
+        :return: True
+        """
+        return True