apache-hamilton 1.90.0.dev0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- apache_hamilton-1.90.0.dev0.dist-info/METADATA +407 -0
- apache_hamilton-1.90.0.dev0.dist-info/RECORD +151 -0
- apache_hamilton-1.90.0.dev0.dist-info/WHEEL +4 -0
- apache_hamilton-1.90.0.dev0.dist-info/entry_points.txt +9 -0
- apache_hamilton-1.90.0.dev0.dist-info/licenses/DISCLAIMER +10 -0
- apache_hamilton-1.90.0.dev0.dist-info/licenses/LICENSE +228 -0
- apache_hamilton-1.90.0.dev0.dist-info/licenses/NOTICE +5 -0
- hamilton/__init__.py +24 -0
- hamilton/ad_hoc_utils.py +132 -0
- hamilton/async_driver.py +465 -0
- hamilton/base.py +466 -0
- hamilton/caching/__init__.py +16 -0
- hamilton/caching/adapter.py +1475 -0
- hamilton/caching/cache_key.py +70 -0
- hamilton/caching/fingerprinting.py +287 -0
- hamilton/caching/stores/__init__.py +16 -0
- hamilton/caching/stores/base.py +242 -0
- hamilton/caching/stores/file.py +140 -0
- hamilton/caching/stores/memory.py +297 -0
- hamilton/caching/stores/sqlite.py +282 -0
- hamilton/caching/stores/utils.py +40 -0
- hamilton/cli/__init__.py +16 -0
- hamilton/cli/__main__.py +328 -0
- hamilton/cli/commands.py +126 -0
- hamilton/cli/logic.py +338 -0
- hamilton/common/__init__.py +76 -0
- hamilton/contrib/__init__.py +41 -0
- hamilton/data_quality/__init__.py +16 -0
- hamilton/data_quality/base.py +198 -0
- hamilton/data_quality/default_validators.py +560 -0
- hamilton/data_quality/pandera_validators.py +121 -0
- hamilton/dataflows/__init__.py +726 -0
- hamilton/dataflows/template/README.md +25 -0
- hamilton/dataflows/template/__init__.py +54 -0
- hamilton/dataflows/template/author.md +28 -0
- hamilton/dataflows/template/requirements.txt +0 -0
- hamilton/dataflows/template/tags.json +7 -0
- hamilton/dataflows/template/valid_configs.jsonl +1 -0
- hamilton/dev_utils/__init__.py +16 -0
- hamilton/dev_utils/deprecation.py +204 -0
- hamilton/driver.py +2112 -0
- hamilton/execution/__init__.py +16 -0
- hamilton/execution/debugging_utils.py +56 -0
- hamilton/execution/executors.py +502 -0
- hamilton/execution/graph_functions.py +421 -0
- hamilton/execution/grouping.py +430 -0
- hamilton/execution/state.py +539 -0
- hamilton/experimental/__init__.py +27 -0
- hamilton/experimental/databackend.py +61 -0
- hamilton/experimental/decorators/__init__.py +16 -0
- hamilton/experimental/decorators/parameterize_frame.py +233 -0
- hamilton/experimental/h_async.py +29 -0
- hamilton/experimental/h_cache.py +413 -0
- hamilton/experimental/h_dask.py +28 -0
- hamilton/experimental/h_databackends.py +174 -0
- hamilton/experimental/h_ray.py +28 -0
- hamilton/experimental/h_spark.py +32 -0
- hamilton/function_modifiers/README +40 -0
- hamilton/function_modifiers/__init__.py +121 -0
- hamilton/function_modifiers/adapters.py +900 -0
- hamilton/function_modifiers/base.py +859 -0
- hamilton/function_modifiers/configuration.py +310 -0
- hamilton/function_modifiers/delayed.py +202 -0
- hamilton/function_modifiers/dependencies.py +246 -0
- hamilton/function_modifiers/expanders.py +1230 -0
- hamilton/function_modifiers/macros.py +1634 -0
- hamilton/function_modifiers/metadata.py +434 -0
- hamilton/function_modifiers/recursive.py +908 -0
- hamilton/function_modifiers/validation.py +289 -0
- hamilton/function_modifiers_base.py +31 -0
- hamilton/graph.py +1153 -0
- hamilton/graph_types.py +264 -0
- hamilton/graph_utils.py +41 -0
- hamilton/htypes.py +450 -0
- hamilton/io/__init__.py +32 -0
- hamilton/io/data_adapters.py +216 -0
- hamilton/io/default_data_loaders.py +224 -0
- hamilton/io/materialization.py +500 -0
- hamilton/io/utils.py +158 -0
- hamilton/lifecycle/__init__.py +67 -0
- hamilton/lifecycle/api.py +833 -0
- hamilton/lifecycle/base.py +1130 -0
- hamilton/lifecycle/default.py +802 -0
- hamilton/log_setup.py +47 -0
- hamilton/models.py +92 -0
- hamilton/node.py +449 -0
- hamilton/plugins/README.md +48 -0
- hamilton/plugins/__init__.py +16 -0
- hamilton/plugins/dask_extensions.py +47 -0
- hamilton/plugins/dlt_extensions.py +161 -0
- hamilton/plugins/geopandas_extensions.py +49 -0
- hamilton/plugins/h_dask.py +331 -0
- hamilton/plugins/h_ddog.py +522 -0
- hamilton/plugins/h_diskcache.py +163 -0
- hamilton/plugins/h_experiments/__init__.py +22 -0
- hamilton/plugins/h_experiments/__main__.py +62 -0
- hamilton/plugins/h_experiments/cache.py +39 -0
- hamilton/plugins/h_experiments/data_model.py +68 -0
- hamilton/plugins/h_experiments/hook.py +219 -0
- hamilton/plugins/h_experiments/server.py +375 -0
- hamilton/plugins/h_kedro.py +152 -0
- hamilton/plugins/h_logging.py +454 -0
- hamilton/plugins/h_mcp/__init__.py +28 -0
- hamilton/plugins/h_mcp/__main__.py +33 -0
- hamilton/plugins/h_mcp/_helpers.py +129 -0
- hamilton/plugins/h_mcp/_templates.py +417 -0
- hamilton/plugins/h_mcp/server.py +328 -0
- hamilton/plugins/h_mlflow.py +335 -0
- hamilton/plugins/h_narwhals.py +134 -0
- hamilton/plugins/h_openlineage.py +400 -0
- hamilton/plugins/h_opentelemetry.py +167 -0
- hamilton/plugins/h_pandas.py +257 -0
- hamilton/plugins/h_pandera.py +117 -0
- hamilton/plugins/h_polars.py +304 -0
- hamilton/plugins/h_polars_lazyframe.py +282 -0
- hamilton/plugins/h_pyarrow.py +56 -0
- hamilton/plugins/h_pydantic.py +127 -0
- hamilton/plugins/h_ray.py +242 -0
- hamilton/plugins/h_rich.py +142 -0
- hamilton/plugins/h_schema.py +493 -0
- hamilton/plugins/h_slack.py +100 -0
- hamilton/plugins/h_spark.py +1380 -0
- hamilton/plugins/h_threadpool.py +125 -0
- hamilton/plugins/h_tqdm.py +122 -0
- hamilton/plugins/h_vaex.py +129 -0
- hamilton/plugins/huggingface_extensions.py +236 -0
- hamilton/plugins/ibis_extensions.py +93 -0
- hamilton/plugins/jupyter_magic.py +622 -0
- hamilton/plugins/kedro_extensions.py +117 -0
- hamilton/plugins/lightgbm_extensions.py +99 -0
- hamilton/plugins/matplotlib_extensions.py +108 -0
- hamilton/plugins/mlflow_extensions.py +216 -0
- hamilton/plugins/numpy_extensions.py +105 -0
- hamilton/plugins/pandas_extensions.py +1763 -0
- hamilton/plugins/plotly_extensions.py +150 -0
- hamilton/plugins/polars_extensions.py +71 -0
- hamilton/plugins/polars_implementations.py +25 -0
- hamilton/plugins/polars_lazyframe_extensions.py +302 -0
- hamilton/plugins/polars_post_1_0_0_extensions.py +877 -0
- hamilton/plugins/polars_pre_1_0_0_extension.py +836 -0
- hamilton/plugins/pydantic_extensions.py +98 -0
- hamilton/plugins/pyspark_pandas_extensions.py +47 -0
- hamilton/plugins/sklearn_plot_extensions.py +129 -0
- hamilton/plugins/spark_extensions.py +105 -0
- hamilton/plugins/vaex_extensions.py +51 -0
- hamilton/plugins/xgboost_extensions.py +91 -0
- hamilton/plugins/yaml_extensions.py +89 -0
- hamilton/registry.py +254 -0
- hamilton/settings.py +18 -0
- hamilton/telemetry.py +50 -0
- hamilton/version.py +18 -0
|
@@ -0,0 +1,257 @@
|
|
|
1
|
+
# Licensed to the Apache Software Foundation (ASF) under one
|
|
2
|
+
# or more contributor license agreements. See the NOTICE file
|
|
3
|
+
# distributed with this work for additional information
|
|
4
|
+
# regarding copyright ownership. The ASF licenses this file
|
|
5
|
+
# to you under the Apache License, Version 2.0 (the
|
|
6
|
+
# "License"); you may not use this file except in compliance
|
|
7
|
+
# with the License. You may obtain a copy of the License at
|
|
8
|
+
#
|
|
9
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
10
|
+
#
|
|
11
|
+
# Unless required by applicable law or agreed to in writing,
|
|
12
|
+
# software distributed under the License is distributed on an
|
|
13
|
+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
14
|
+
# KIND, either express or implied. See the License for the
|
|
15
|
+
# specific language governing permissions and limitations
|
|
16
|
+
# under the License.
|
|
17
|
+
|
|
18
|
+
import sys
|
|
19
|
+
from collections.abc import Callable, Collection
|
|
20
|
+
from types import ModuleType
|
|
21
|
+
from typing import Any, get_type_hints
|
|
22
|
+
|
|
23
|
+
_sys_version_info = sys.version_info
|
|
24
|
+
_version_tuple = (_sys_version_info.major, _sys_version_info.minor, _sys_version_info.micro)
|
|
25
|
+
|
|
26
|
+
if _version_tuple < (3, 11, 0):
|
|
27
|
+
pass
|
|
28
|
+
else:
|
|
29
|
+
pass
|
|
30
|
+
|
|
31
|
+
from hamilton import node, registry
|
|
32
|
+
from hamilton.function_modifiers.expanders import extract_columns
|
|
33
|
+
from hamilton.function_modifiers.recursive import (
|
|
34
|
+
_default_inject_parameter,
|
|
35
|
+
subdag,
|
|
36
|
+
with_columns_base,
|
|
37
|
+
)
|
|
38
|
+
from hamilton.plugins.pandas_extensions import DATAFRAME_TYPE
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class with_columns(with_columns_base):
|
|
42
|
+
"""Initializes a with_columns decorator for pandas. This allows you to efficiently run groups of map operations on a dataframe.
|
|
43
|
+
|
|
44
|
+
Here's an example of calling it -- if you've seen ``@subdag``, you should be familiar with
|
|
45
|
+
the concepts:
|
|
46
|
+
|
|
47
|
+
.. code-block:: python
|
|
48
|
+
|
|
49
|
+
# my_module.py
|
|
50
|
+
def a(a_from_df: pd.Series) -> pd.Series:
|
|
51
|
+
return _process(a)
|
|
52
|
+
|
|
53
|
+
def b(b_from_df: pd.Series) -> pd.Series:
|
|
54
|
+
return _process(b)
|
|
55
|
+
|
|
56
|
+
def a_b_average(a_from_df: pd.Series, b_from_df: pd.Series) -> pd.Series:
|
|
57
|
+
return (a_from_df + b_from_df) / 2
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
.. code-block:: python
|
|
61
|
+
|
|
62
|
+
# with_columns_module.py
|
|
63
|
+
def a_plus_b(a: pd.Series, b: pd.Series) -> pd.Series:
|
|
64
|
+
return a + b
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
# the with_columns call
|
|
68
|
+
@with_columns(
|
|
69
|
+
*[my_module], # Load from any module
|
|
70
|
+
*[a_plus_b], # or list operations directly
|
|
71
|
+
columns_to_pass=["a_from_df", "b_from_df"], # The columns to pass from the dataframe to
|
|
72
|
+
# the subdag
|
|
73
|
+
select=["a", "b", "a_plus_b", "a_b_average"], # The columns to select from the dataframe
|
|
74
|
+
)
|
|
75
|
+
def final_df(initial_df: pd.DataFrame) -> pd.DataFrame:
|
|
76
|
+
# process, or just return unprocessed
|
|
77
|
+
...
|
|
78
|
+
|
|
79
|
+
In this instance the ``initial_df`` would get two columns added: ``a_plus_b`` and ``a_b_average``.
|
|
80
|
+
|
|
81
|
+
The operations are applied in topological order. This allows you to
|
|
82
|
+
express the operations individually, making it easy to unit-test and reuse.
|
|
83
|
+
|
|
84
|
+
Note that the operation is "append", meaning that the columns that are selected are appended
|
|
85
|
+
onto the dataframe.
|
|
86
|
+
|
|
87
|
+
If the function takes multiple dataframes, the dataframe input to process will always be
|
|
88
|
+
the first argument. This will be passed to the subdag, transformed, and passed back to the function.
|
|
89
|
+
This follows the hamilton rule of reference by parameter name. To demonstarte this, in the code
|
|
90
|
+
above, the dataframe that is passed to the subdag is `initial_df`. That is transformed
|
|
91
|
+
by the subdag, and then returned as the final dataframe.
|
|
92
|
+
|
|
93
|
+
You can read it as:
|
|
94
|
+
|
|
95
|
+
"final_df is a function that transforms the upstream dataframe initial_df, running the transformations
|
|
96
|
+
from my_module. It starts with the columns a_from_df and b_from_df, and then adds the columns
|
|
97
|
+
a, b, and a_plus_b to the dataframe. It then returns the dataframe, and does some processing on it."
|
|
98
|
+
|
|
99
|
+
In case you need more flexibility you can alternatively use ``on_input``, for example,
|
|
100
|
+
|
|
101
|
+
.. code-block:: python
|
|
102
|
+
|
|
103
|
+
# with_columns_module.py
|
|
104
|
+
def a_from_df(initial_df: pd.Series) -> pd.Series:
|
|
105
|
+
return initial_df["a_from_df"] / 100
|
|
106
|
+
|
|
107
|
+
def b_from_df(initial_df: pd.Series) -> pd.Series:
|
|
108
|
+
return initial_df["b_from_df"] / 100
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
# the with_columns call
|
|
112
|
+
@with_columns(
|
|
113
|
+
*[my_module],
|
|
114
|
+
*[a_from_df],
|
|
115
|
+
on_input="initial_df",
|
|
116
|
+
select=["a_from_df", "b_from_df", "a", "b", "a_plus_b", "a_b_average"],
|
|
117
|
+
)
|
|
118
|
+
def final_df(initial_df: pd.DataFrame, ...) -> pd.DataFrame:
|
|
119
|
+
# process, or just return unprocessed
|
|
120
|
+
...
|
|
121
|
+
|
|
122
|
+
the above would output a dataframe where the two columns ``a_from_df`` and ``b_from_df`` get
|
|
123
|
+
overwritten.
|
|
124
|
+
"""
|
|
125
|
+
|
|
126
|
+
def __init__(
|
|
127
|
+
self,
|
|
128
|
+
*load_from: Callable | ModuleType,
|
|
129
|
+
columns_to_pass: list[str] = None,
|
|
130
|
+
pass_dataframe_as: str = None,
|
|
131
|
+
on_input: str = None,
|
|
132
|
+
select: list[str] = None,
|
|
133
|
+
namespace: str = None,
|
|
134
|
+
config_required: list[str] = None,
|
|
135
|
+
):
|
|
136
|
+
"""Instantiates a ``@with_columns`` decorator.
|
|
137
|
+
|
|
138
|
+
:param load_from: The functions or modules that will be used to generate the group of map operations.
|
|
139
|
+
:param columns_to_pass: The initial schema of the dataframe. This is used to determine which
|
|
140
|
+
upstream inputs should be taken from the dataframe, and which shouldn't. Note that, if this is
|
|
141
|
+
left empty (and external_inputs is as well), we will assume that all dependencies come
|
|
142
|
+
from the dataframe. This cannot be used in conjunction with on_input.
|
|
143
|
+
:param on_input: The name of the dataframe that we're modifying, as known to the subdag.
|
|
144
|
+
If you pass this in, you are responsible for extracting columns out. If not provided, you have
|
|
145
|
+
to pass columns_to_pass in, and we will extract the columns out on the first parameter for you.
|
|
146
|
+
:param select: The end nodes that represent columns to be appended to the original dataframe
|
|
147
|
+
via with_columns. Existing columns will be overridden. The selected nodes need to have the
|
|
148
|
+
corresponding column type, in this case pd.Series, to be appended to the original dataframe.
|
|
149
|
+
:param namespace: The namespace of the nodes, so they don't clash with the global namespace
|
|
150
|
+
and so this can be reused. If its left out, there will be no namespace (in which case you'll want
|
|
151
|
+
to be careful about repeating it/reusing the nodes in other parts of the DAG.)
|
|
152
|
+
:param config_required: the list of config keys that are required to resolve any functions. Pass in None\
|
|
153
|
+
if you want the functions/modules to have access to all possible config.
|
|
154
|
+
"""
|
|
155
|
+
|
|
156
|
+
if pass_dataframe_as is not None:
|
|
157
|
+
raise NotImplementedError(
|
|
158
|
+
"We currently do not support pass_dataframe_as for pandas. Please reach out if you need this "
|
|
159
|
+
"functionality."
|
|
160
|
+
)
|
|
161
|
+
|
|
162
|
+
super().__init__(
|
|
163
|
+
*load_from,
|
|
164
|
+
columns_to_pass=columns_to_pass,
|
|
165
|
+
on_input=on_input,
|
|
166
|
+
select=select,
|
|
167
|
+
namespace=namespace,
|
|
168
|
+
config_required=config_required,
|
|
169
|
+
dataframe_type=DATAFRAME_TYPE,
|
|
170
|
+
)
|
|
171
|
+
|
|
172
|
+
def _create_column_nodes(
|
|
173
|
+
self, fn: Callable, inject_parameter: str, params: dict[str, type[type]]
|
|
174
|
+
) -> list[node.Node]:
|
|
175
|
+
output_type = params[inject_parameter]
|
|
176
|
+
|
|
177
|
+
def temp_fn(**kwargs) -> Any:
|
|
178
|
+
return kwargs[inject_parameter]
|
|
179
|
+
|
|
180
|
+
# We recreate the df node to use extract columns
|
|
181
|
+
temp_node = node.Node(
|
|
182
|
+
name=inject_parameter,
|
|
183
|
+
typ=output_type,
|
|
184
|
+
callabl=temp_fn,
|
|
185
|
+
input_types={inject_parameter: output_type},
|
|
186
|
+
)
|
|
187
|
+
|
|
188
|
+
extract_columns_decorator = extract_columns(*self.initial_schema)
|
|
189
|
+
|
|
190
|
+
out_nodes = extract_columns_decorator.transform_node(temp_node, config={}, fn=temp_fn)
|
|
191
|
+
return out_nodes[1:]
|
|
192
|
+
|
|
193
|
+
def get_initial_nodes(
|
|
194
|
+
self, fn: Callable, params: dict[str, type[type]]
|
|
195
|
+
) -> tuple[str, Collection[node.Node]]:
|
|
196
|
+
"""Selects the correct dataframe and optionally extracts out columns."""
|
|
197
|
+
inject_parameter = _default_inject_parameter(fn=fn, target_dataframe=self.target_dataframe)
|
|
198
|
+
with_columns_base.validate_dataframe(
|
|
199
|
+
fn=fn,
|
|
200
|
+
inject_parameter=inject_parameter,
|
|
201
|
+
params=params,
|
|
202
|
+
required_type=self.dataframe_type,
|
|
203
|
+
)
|
|
204
|
+
|
|
205
|
+
initial_nodes = (
|
|
206
|
+
[]
|
|
207
|
+
if self.target_dataframe is not None
|
|
208
|
+
else self._create_column_nodes(fn=fn, inject_parameter=inject_parameter, params=params)
|
|
209
|
+
)
|
|
210
|
+
|
|
211
|
+
return inject_parameter, initial_nodes
|
|
212
|
+
|
|
213
|
+
def get_subdag_nodes(self, fn: Callable, config: dict[str, Any]) -> Collection[node.Node]:
|
|
214
|
+
return subdag.collect_nodes(config, self.subdag_functions)
|
|
215
|
+
|
|
216
|
+
def chain_subdag_nodes(
|
|
217
|
+
self, fn: Callable, inject_parameter: str, generated_nodes: Collection[node.Node]
|
|
218
|
+
) -> node.Node:
|
|
219
|
+
"Node that adds to / overrides columns for the original dataframe based on selected output."
|
|
220
|
+
# In case no node is selected we append all possible nodes that have a column type matching
|
|
221
|
+
# what the dataframe expects
|
|
222
|
+
if self.select is None:
|
|
223
|
+
self.select = [
|
|
224
|
+
sink_node.name
|
|
225
|
+
for sink_node in generated_nodes
|
|
226
|
+
if sink_node.type == registry.get_column_type_from_df_type(self.dataframe_type)
|
|
227
|
+
]
|
|
228
|
+
|
|
229
|
+
def new_callable(**kwargs) -> Any:
|
|
230
|
+
df = kwargs[inject_parameter]
|
|
231
|
+
columns_to_append = {}
|
|
232
|
+
for column in self.select:
|
|
233
|
+
columns_to_append[column] = kwargs[column]
|
|
234
|
+
|
|
235
|
+
return df.assign(**columns_to_append)
|
|
236
|
+
|
|
237
|
+
column_type = registry.get_column_type_from_df_type(self.dataframe_type)
|
|
238
|
+
input_map = {column: column_type for column in self.select}
|
|
239
|
+
input_map[inject_parameter] = self.dataframe_type
|
|
240
|
+
merge_node = node.Node(
|
|
241
|
+
name="_append",
|
|
242
|
+
typ=self.dataframe_type,
|
|
243
|
+
callabl=new_callable,
|
|
244
|
+
input_types=input_map,
|
|
245
|
+
)
|
|
246
|
+
output_nodes = generated_nodes + [merge_node]
|
|
247
|
+
return output_nodes, merge_node.name
|
|
248
|
+
|
|
249
|
+
def validate(self, fn: Callable):
|
|
250
|
+
inject_parameter = _default_inject_parameter(fn=fn, target_dataframe=self.target_dataframe)
|
|
251
|
+
params = get_type_hints(fn)
|
|
252
|
+
with_columns_base.validate_dataframe(
|
|
253
|
+
fn=fn,
|
|
254
|
+
inject_parameter=inject_parameter,
|
|
255
|
+
params=params,
|
|
256
|
+
required_type=self.dataframe_type,
|
|
257
|
+
)
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
# Licensed to the Apache Software Foundation (ASF) under one
|
|
2
|
+
# or more contributor license agreements. See the NOTICE file
|
|
3
|
+
# distributed with this work for additional information
|
|
4
|
+
# regarding copyright ownership. The ASF licenses this file
|
|
5
|
+
# to you under the Apache License, Version 2.0 (the
|
|
6
|
+
# "License"); you may not use this file except in compliance
|
|
7
|
+
# with the License. You may obtain a copy of the License at
|
|
8
|
+
#
|
|
9
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
10
|
+
#
|
|
11
|
+
# Unless required by applicable law or agreed to in writing,
|
|
12
|
+
# software distributed under the License is distributed on an
|
|
13
|
+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
14
|
+
# KIND, either express or implied. See the License for the
|
|
15
|
+
# specific language governing permissions and limitations
|
|
16
|
+
# under the License.
|
|
17
|
+
|
|
18
|
+
import typing
|
|
19
|
+
|
|
20
|
+
import pandera
|
|
21
|
+
from pandera import typing as pa_typing
|
|
22
|
+
|
|
23
|
+
from hamilton import node
|
|
24
|
+
from hamilton.data_quality import base as dq_base
|
|
25
|
+
from hamilton.function_modifiers import InvalidDecoratorException
|
|
26
|
+
from hamilton.function_modifiers import base as fm_base
|
|
27
|
+
from hamilton.function_modifiers import check_output as base_check_output
|
|
28
|
+
from hamilton.function_modifiers.validation import BaseDataValidationDecorator
|
|
29
|
+
from hamilton.htypes import custom_subclass_check
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class check_output(BaseDataValidationDecorator):
|
|
33
|
+
def __init__(
|
|
34
|
+
self,
|
|
35
|
+
importance: str = dq_base.DataValidationLevel.WARN.value,
|
|
36
|
+
target: fm_base.TargetType = None,
|
|
37
|
+
):
|
|
38
|
+
"""Specific output-checker for pandera schemas. This decorator utilizes the output type of the function, which has
|
|
39
|
+
to be of type pandera.typing.pandas.DataFrame or pandera.typing.pandas.Series, with an annotation argument.
|
|
40
|
+
|
|
41
|
+
:param schema: The schema to use for validation. If this is not provided, then the output type of the function is used.
|
|
42
|
+
:param importance: Importance level (either "warn" or "fail") -- see documentation for check_output for more details.
|
|
43
|
+
:param target: The target of the decorator -- see documentation for check_output for more details.
|
|
44
|
+
|
|
45
|
+
Let's look at equivalent examples to demonstrate:
|
|
46
|
+
|
|
47
|
+
.. code-block:: python
|
|
48
|
+
:name: "@check_output using output type"
|
|
49
|
+
|
|
50
|
+
import pandera as pa
|
|
51
|
+
import pandas as pd
|
|
52
|
+
from hamilton.plugins import h_pandera
|
|
53
|
+
from pandera.typing.pandas import DataFrame
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class MySchema(pa.DataFrameModel):
|
|
57
|
+
a: int
|
|
58
|
+
b: float
|
|
59
|
+
c: str = pa.Field(nullable=True) # For example, allow None values
|
|
60
|
+
d: float # US dollars
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
@h_pandera.check_output()
|
|
64
|
+
def foo() -> DataFrame[MySchema]:
|
|
65
|
+
return pd.DataFrame() # will fail
|
|
66
|
+
|
|
67
|
+
.. code-block:: python
|
|
68
|
+
:name: "@check_output with passed in type"
|
|
69
|
+
|
|
70
|
+
from hamilton import function_modifiers
|
|
71
|
+
|
|
72
|
+
schema = pa.DataFrameSchema({
|
|
73
|
+
"a": pa.Column(pa.Int),
|
|
74
|
+
"b": pa.Column(pa.Float),
|
|
75
|
+
"c": pa.Column(pa.String, nullable=True),
|
|
76
|
+
"d": pa.Column(pa.Float),
|
|
77
|
+
})
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@function_modifiers.check_output(schema=schema)
|
|
81
|
+
def foo() -> pd.DataFrame:
|
|
82
|
+
return pd.DataFrame() # will fail
|
|
83
|
+
|
|
84
|
+
These two are functionally equivalent. Note that we do not (yet) support modification of the output.
|
|
85
|
+
"""
|
|
86
|
+
super(check_output, self).__init__(target)
|
|
87
|
+
self.importance = importance
|
|
88
|
+
self.target = target
|
|
89
|
+
|
|
90
|
+
def get_validators(self, node_to_validate: node.Node) -> list[dq_base.DataValidator]:
|
|
91
|
+
"""Gets validators for the node. Delegates to the standard check_output(schema=...) decorator.
|
|
92
|
+
|
|
93
|
+
:param node_to_validate: Node to validate
|
|
94
|
+
:return: List of validators
|
|
95
|
+
"""
|
|
96
|
+
output_type = node_to_validate.type
|
|
97
|
+
schema = None
|
|
98
|
+
origin = typing.get_origin(output_type)
|
|
99
|
+
args = typing.get_args(output_type)
|
|
100
|
+
if custom_subclass_check(origin, pa_typing.DataFrame) and len(args) == 1:
|
|
101
|
+
schema = output_type.__args__[0] # TODO -- determine if it can ever have multiple...
|
|
102
|
+
if not issubclass(schema, pandera.DataFrameModel):
|
|
103
|
+
schema = None
|
|
104
|
+
|
|
105
|
+
if schema is None:
|
|
106
|
+
raise InvalidDecoratorException(
|
|
107
|
+
f"Output type {output_type} is not a valid pandera schema. "
|
|
108
|
+
f"Note that we currently only support pandera dataframes annotated with "
|
|
109
|
+
f"subclasses of pandera.DataFrameModel. Please reach out/open an issue "
|
|
110
|
+
f"if you want more complete integration."
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
# We can just delegate to teh standard check_output, which has pandera associated with schema...
|
|
114
|
+
# This is a clever way of reusing as much code as possible
|
|
115
|
+
return base_check_output(
|
|
116
|
+
importance=self.importance, schema=schema, target_=self.target
|
|
117
|
+
).get_validators(node_to_validate)
|
|
@@ -0,0 +1,304 @@
|
|
|
1
|
+
# Licensed to the Apache Software Foundation (ASF) under one
|
|
2
|
+
# or more contributor license agreements. See the NOTICE file
|
|
3
|
+
# distributed with this work for additional information
|
|
4
|
+
# regarding copyright ownership. The ASF licenses this file
|
|
5
|
+
# to you under the Apache License, Version 2.0 (the
|
|
6
|
+
# "License"); you may not use this file except in compliance
|
|
7
|
+
# with the License. You may obtain a copy of the License at
|
|
8
|
+
#
|
|
9
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
10
|
+
#
|
|
11
|
+
# Unless required by applicable law or agreed to in writing,
|
|
12
|
+
# software distributed under the License is distributed on an
|
|
13
|
+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
14
|
+
# KIND, either express or implied. See the License for the
|
|
15
|
+
# specific language governing permissions and limitations
|
|
16
|
+
# under the License.
|
|
17
|
+
|
|
18
|
+
import sys
|
|
19
|
+
from collections.abc import Callable, Collection
|
|
20
|
+
from types import ModuleType
|
|
21
|
+
from typing import Any, get_type_hints
|
|
22
|
+
|
|
23
|
+
import polars as pl
|
|
24
|
+
|
|
25
|
+
_sys_version_info = sys.version_info
|
|
26
|
+
_version_tuple = (_sys_version_info.major, _sys_version_info.minor, _sys_version_info.micro)
|
|
27
|
+
|
|
28
|
+
if _version_tuple < (3, 11, 0):
|
|
29
|
+
pass
|
|
30
|
+
else:
|
|
31
|
+
pass
|
|
32
|
+
|
|
33
|
+
# Copied this over from function_graph
|
|
34
|
+
# TODO -- determine the best place to put this code
|
|
35
|
+
from hamilton import base, node, registry
|
|
36
|
+
from hamilton.function_modifiers.expanders import extract_columns
|
|
37
|
+
from hamilton.function_modifiers.recursive import (
|
|
38
|
+
_default_inject_parameter,
|
|
39
|
+
subdag,
|
|
40
|
+
with_columns_base,
|
|
41
|
+
)
|
|
42
|
+
from hamilton.plugins.polars_extensions import DATAFRAME_TYPE
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class PolarsDataFrameResult(base.ResultMixin):
|
|
46
|
+
"""A ResultBuilder that produces a polars dataframe.
|
|
47
|
+
|
|
48
|
+
Use this when you want to create a polars dataframe from the outputs. Caveat: you need to ensure that the length
|
|
49
|
+
of the outputs is the same, otherwise you will get an error; mixed outputs aren't that well handled.
|
|
50
|
+
|
|
51
|
+
To use:
|
|
52
|
+
|
|
53
|
+
.. code-block:: python
|
|
54
|
+
|
|
55
|
+
from hamilton import base, driver
|
|
56
|
+
from hamilton.plugins import polars_extensions
|
|
57
|
+
|
|
58
|
+
polars_builder = polars_extensions.PolarsDataFrameResult()
|
|
59
|
+
adapter = base.SimplePythonGraphAdapter(polars_builder)
|
|
60
|
+
dr = driver.Driver(config, *modules, adapter=adapter)
|
|
61
|
+
df = dr.execute([...], inputs=...) # returns polars dataframe
|
|
62
|
+
|
|
63
|
+
Note: this is just a first attempt at something for Polars. Think it should handle more? Come chat/open a PR!
|
|
64
|
+
"""
|
|
65
|
+
|
|
66
|
+
def build_result(self, **outputs: dict[str, pl.Series | pl.DataFrame | Any]) -> pl.DataFrame:
|
|
67
|
+
"""This is the method that Hamilton will call to build the final result. It will pass in the results
|
|
68
|
+
of the requested outputs that you passed in to the execute() method.
|
|
69
|
+
|
|
70
|
+
Note: this function could do smarter things; looking for contributions here!
|
|
71
|
+
|
|
72
|
+
:param outputs: The results of the requested outputs.
|
|
73
|
+
:return: a polars DataFrame.
|
|
74
|
+
"""
|
|
75
|
+
if len(outputs) == 1:
|
|
76
|
+
(value,) = outputs.values() # this works because it's length 1.
|
|
77
|
+
if isinstance(value, pl.DataFrame): # it's a dataframe
|
|
78
|
+
return value
|
|
79
|
+
if isinstance(value, pl.LazyFrame): # it's a lazyframe
|
|
80
|
+
return value.collect()
|
|
81
|
+
elif not isinstance(value, pl.Series): # it's a single scalar/object
|
|
82
|
+
key, value = outputs.popitem()
|
|
83
|
+
return pl.DataFrame({key: [value]})
|
|
84
|
+
else: # it's a series
|
|
85
|
+
return pl.DataFrame(outputs)
|
|
86
|
+
# TODO: check for length of outputs and determine what should
|
|
87
|
+
# happen for mixed outputs that include scalars for example.
|
|
88
|
+
return pl.DataFrame(outputs)
|
|
89
|
+
|
|
90
|
+
def output_type(self) -> type:
|
|
91
|
+
return pl.DataFrame
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
# Do we need this here?
|
|
95
|
+
class with_columns(with_columns_base):
|
|
96
|
+
"""Initializes a with_columns decorator for polars.
|
|
97
|
+
|
|
98
|
+
This allows you to efficiently run groups of map operations on a dataframe. We support
|
|
99
|
+
both eager and lazy mode in polars. In case of using eager mode the type should be
|
|
100
|
+
pl.DataFrame and the subsequent operations run on columns with type pl.Series.
|
|
101
|
+
|
|
102
|
+
Here's an example of calling in eager mode -- if you've seen ``@subdag``, you should be familiar with
|
|
103
|
+
the concepts:
|
|
104
|
+
|
|
105
|
+
.. code-block:: python
|
|
106
|
+
|
|
107
|
+
# my_module.py
|
|
108
|
+
def a_b_average(a: pl.Series, b: pl.Series) -> pl.Series:
|
|
109
|
+
return (a + b) / 2
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
.. code-block:: python
|
|
113
|
+
|
|
114
|
+
# with_columns_module.py
|
|
115
|
+
def a_plus_b(a: pl.Series, b: pl.Series) -> pl.Series:
|
|
116
|
+
return a + b
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
# the with_columns call
|
|
120
|
+
@with_columns(
|
|
121
|
+
*[my_module], # Load from any module
|
|
122
|
+
*[a_plus_b], # or list operations directly
|
|
123
|
+
columns_to_pass=["a", "b"], # The columns to pass from the dataframe to
|
|
124
|
+
# the subdag
|
|
125
|
+
select=["a_plus_b", "a_b_average"], # The columns to append to the dataframe
|
|
126
|
+
)
|
|
127
|
+
def final_df(initial_df: pl.DataFrame) -> pl.DataFrame:
|
|
128
|
+
# process, or just return unprocessed
|
|
129
|
+
...
|
|
130
|
+
|
|
131
|
+
In this instance the ``initial_df`` would get two columns added: ``a_plus_b`` and ``a_b_average``.
|
|
132
|
+
|
|
133
|
+
Note that the operation is "append", meaning that the columns that are selected are appended
|
|
134
|
+
onto the dataframe.
|
|
135
|
+
|
|
136
|
+
If the function takes multiple dataframes, the dataframe input to process will always be
|
|
137
|
+
the first argument. This will be passed to the subdag, transformed, and passed back to the function.
|
|
138
|
+
This follows the hamilton rule of reference by parameter name. To demonstarte this, in the code
|
|
139
|
+
above, the dataframe that is passed to the subdag is `initial_df`. That is transformed
|
|
140
|
+
by the subdag, and then returned as the final dataframe.
|
|
141
|
+
|
|
142
|
+
You can read it as:
|
|
143
|
+
|
|
144
|
+
"final_df is a function that transforms the upstream dataframe initial_df, running the transformations
|
|
145
|
+
from my_module. It starts with the columns a_from_df and b_from_df, and then adds the columns
|
|
146
|
+
a, b, and a_plus_b to the dataframe. It then returns the dataframe, and does some processing on it."
|
|
147
|
+
|
|
148
|
+
In case you need more flexibility you can alternatively use ``on_input``, for example,
|
|
149
|
+
|
|
150
|
+
.. code-block:: python
|
|
151
|
+
|
|
152
|
+
# with_columns_module.py
|
|
153
|
+
def a_from_df() -> pl.Expr:
|
|
154
|
+
return pl.col(a).alias("a") / 100
|
|
155
|
+
|
|
156
|
+
def b_from_df() -> pl.Expr:
|
|
157
|
+
return pl.col(b).alias("b") / 100
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
# the with_columns call
|
|
161
|
+
@with_columns(
|
|
162
|
+
*[my_module],
|
|
163
|
+
on_input="initial_df",
|
|
164
|
+
select=["a_from_df", "b_from_df", "a_plus_b", "a_b_average"],
|
|
165
|
+
)
|
|
166
|
+
def final_df(initial_df: pl.DataFrame) -> pl.DataFrame:
|
|
167
|
+
# process, or just return unprocessed
|
|
168
|
+
...
|
|
169
|
+
|
|
170
|
+
the above would output a dataframe where the two columns ``a`` and ``b`` get
|
|
171
|
+
overwritten.
|
|
172
|
+
"""
|
|
173
|
+
|
|
174
|
+
def __init__(
|
|
175
|
+
self,
|
|
176
|
+
*load_from: Callable | ModuleType,
|
|
177
|
+
columns_to_pass: list[str] = None,
|
|
178
|
+
pass_dataframe_as: str = None,
|
|
179
|
+
on_input: str = None,
|
|
180
|
+
select: list[str] = None,
|
|
181
|
+
namespace: str = None,
|
|
182
|
+
config_required: list[str] = None,
|
|
183
|
+
):
|
|
184
|
+
"""Instantiates a ``@with_columns`` decorator.
|
|
185
|
+
|
|
186
|
+
:param load_from: The functions or modules that will be used to generate the group of map operations.
|
|
187
|
+
:param columns_to_pass: The initial schema of the dataframe. This is used to determine which
|
|
188
|
+
upstream inputs should be taken from the dataframe, and which shouldn't. Note that, if this is
|
|
189
|
+
left empty (and external_inputs is as well), we will assume that all dependencies come
|
|
190
|
+
from the dataframe. This cannot be used in conjunction with on_input.
|
|
191
|
+
:param on_input: The name of the dataframe that we're modifying, as known to the subdag.
|
|
192
|
+
If you pass this in, you are responsible for extracting columns out. If not provided, you have
|
|
193
|
+
to pass columns_to_pass in, and we will extract the columns out on the first parameter for you.
|
|
194
|
+
:param select: The end nodes that represent columns to be appended to the original dataframe
|
|
195
|
+
via with_columns. Existing columns will be overridden. The selected nodes need to have the
|
|
196
|
+
corresponding column type, in this case pl.Series, to be appended to the original dataframe.
|
|
197
|
+
:param namespace: The namespace of the nodes, so they don't clash with the global namespace
|
|
198
|
+
and so this can be reused. If its left out, there will be no namespace (in which case you'll want
|
|
199
|
+
to be careful about repeating it/reusing the nodes in other parts of the DAG.)
|
|
200
|
+
:param config_required: the list of config keys that are required to resolve any functions. Pass in None\
|
|
201
|
+
if you want the functions/modules to have access to all possible config.
|
|
202
|
+
"""
|
|
203
|
+
|
|
204
|
+
if pass_dataframe_as is not None:
|
|
205
|
+
raise NotImplementedError(
|
|
206
|
+
"We currently do not support pass_dataframe_as for pandas. Please reach out if you need this "
|
|
207
|
+
"functionality."
|
|
208
|
+
)
|
|
209
|
+
|
|
210
|
+
super().__init__(
|
|
211
|
+
*load_from,
|
|
212
|
+
columns_to_pass=columns_to_pass,
|
|
213
|
+
on_input=on_input,
|
|
214
|
+
select=select,
|
|
215
|
+
namespace=namespace,
|
|
216
|
+
config_required=config_required,
|
|
217
|
+
dataframe_type=DATAFRAME_TYPE,
|
|
218
|
+
)
|
|
219
|
+
|
|
220
|
+
def _create_column_nodes(
|
|
221
|
+
self, fn: Callable, inject_parameter: str, params: dict[str, type[type]]
|
|
222
|
+
) -> list[node.Node]:
|
|
223
|
+
output_type = params[inject_parameter]
|
|
224
|
+
|
|
225
|
+
def temp_fn(**kwargs) -> Any:
|
|
226
|
+
return kwargs[inject_parameter]
|
|
227
|
+
|
|
228
|
+
# We recreate the df node to use extract columns
|
|
229
|
+
temp_node = node.Node(
|
|
230
|
+
name=inject_parameter,
|
|
231
|
+
typ=output_type,
|
|
232
|
+
callabl=temp_fn,
|
|
233
|
+
input_types={inject_parameter: output_type},
|
|
234
|
+
)
|
|
235
|
+
|
|
236
|
+
extract_columns_decorator = extract_columns(*self.initial_schema)
|
|
237
|
+
|
|
238
|
+
out_nodes = extract_columns_decorator.transform_node(temp_node, config={}, fn=temp_fn)
|
|
239
|
+
return out_nodes[1:]
|
|
240
|
+
|
|
241
|
+
def get_initial_nodes(
|
|
242
|
+
self, fn: Callable, params: dict[str, type[type]]
|
|
243
|
+
) -> tuple[str, Collection[node.Node]]:
|
|
244
|
+
"""Selects the correct dataframe and optionally extracts out columns."""
|
|
245
|
+
inject_parameter = _default_inject_parameter(fn=fn, target_dataframe=self.target_dataframe)
|
|
246
|
+
with_columns_base.validate_dataframe(
|
|
247
|
+
fn=fn,
|
|
248
|
+
inject_parameter=inject_parameter,
|
|
249
|
+
params=params,
|
|
250
|
+
required_type=self.dataframe_type,
|
|
251
|
+
)
|
|
252
|
+
|
|
253
|
+
initial_nodes = (
|
|
254
|
+
[]
|
|
255
|
+
if self.target_dataframe is not None
|
|
256
|
+
else self._create_column_nodes(fn=fn, inject_parameter=inject_parameter, params=params)
|
|
257
|
+
)
|
|
258
|
+
|
|
259
|
+
return inject_parameter, initial_nodes
|
|
260
|
+
|
|
261
|
+
def get_subdag_nodes(self, fn: Callable, config: dict[str, Any]) -> Collection[node.Node]:
|
|
262
|
+
return subdag.collect_nodes(config, self.subdag_functions)
|
|
263
|
+
|
|
264
|
+
def chain_subdag_nodes(
|
|
265
|
+
self, fn: Callable, inject_parameter: str, generated_nodes: Collection[node.Node]
|
|
266
|
+
) -> node.Node:
|
|
267
|
+
"Node that adds to / overrides columns for the original dataframe based on selected output."
|
|
268
|
+
|
|
269
|
+
if self.select is None:
|
|
270
|
+
self.select = [
|
|
271
|
+
sink_node.name
|
|
272
|
+
for sink_node in generated_nodes
|
|
273
|
+
if sink_node.type == registry.get_column_type_from_df_type(self.dataframe_type)
|
|
274
|
+
]
|
|
275
|
+
|
|
276
|
+
def new_callable(**kwargs) -> Any:
|
|
277
|
+
df = kwargs[inject_parameter]
|
|
278
|
+
columns_to_append = {}
|
|
279
|
+
for column in self.select:
|
|
280
|
+
columns_to_append[column] = kwargs[column]
|
|
281
|
+
|
|
282
|
+
return df.with_columns(**columns_to_append)
|
|
283
|
+
|
|
284
|
+
column_type = registry.get_column_type_from_df_type(self.dataframe_type)
|
|
285
|
+
input_map = {column: column_type for column in self.select}
|
|
286
|
+
input_map[inject_parameter] = self.dataframe_type
|
|
287
|
+
merge_node = node.Node(
|
|
288
|
+
name="_append",
|
|
289
|
+
typ=self.dataframe_type,
|
|
290
|
+
callabl=new_callable,
|
|
291
|
+
input_types=input_map,
|
|
292
|
+
)
|
|
293
|
+
output_nodes = generated_nodes + [merge_node]
|
|
294
|
+
return output_nodes, merge_node.name
|
|
295
|
+
|
|
296
|
+
def validate(self, fn: Callable):
|
|
297
|
+
inject_parameter = _default_inject_parameter(fn=fn, target_dataframe=self.target_dataframe)
|
|
298
|
+
params = get_type_hints(fn)
|
|
299
|
+
with_columns_base.validate_dataframe(
|
|
300
|
+
fn=fn,
|
|
301
|
+
inject_parameter=inject_parameter,
|
|
302
|
+
params=params,
|
|
303
|
+
required_type=self.dataframe_type,
|
|
304
|
+
)
|