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.
Files changed (151) hide show
  1. apache_hamilton-1.90.0.dev0.dist-info/METADATA +407 -0
  2. apache_hamilton-1.90.0.dev0.dist-info/RECORD +151 -0
  3. apache_hamilton-1.90.0.dev0.dist-info/WHEEL +4 -0
  4. apache_hamilton-1.90.0.dev0.dist-info/entry_points.txt +9 -0
  5. apache_hamilton-1.90.0.dev0.dist-info/licenses/DISCLAIMER +10 -0
  6. apache_hamilton-1.90.0.dev0.dist-info/licenses/LICENSE +228 -0
  7. apache_hamilton-1.90.0.dev0.dist-info/licenses/NOTICE +5 -0
  8. hamilton/__init__.py +24 -0
  9. hamilton/ad_hoc_utils.py +132 -0
  10. hamilton/async_driver.py +465 -0
  11. hamilton/base.py +466 -0
  12. hamilton/caching/__init__.py +16 -0
  13. hamilton/caching/adapter.py +1475 -0
  14. hamilton/caching/cache_key.py +70 -0
  15. hamilton/caching/fingerprinting.py +287 -0
  16. hamilton/caching/stores/__init__.py +16 -0
  17. hamilton/caching/stores/base.py +242 -0
  18. hamilton/caching/stores/file.py +140 -0
  19. hamilton/caching/stores/memory.py +297 -0
  20. hamilton/caching/stores/sqlite.py +282 -0
  21. hamilton/caching/stores/utils.py +40 -0
  22. hamilton/cli/__init__.py +16 -0
  23. hamilton/cli/__main__.py +328 -0
  24. hamilton/cli/commands.py +126 -0
  25. hamilton/cli/logic.py +338 -0
  26. hamilton/common/__init__.py +76 -0
  27. hamilton/contrib/__init__.py +41 -0
  28. hamilton/data_quality/__init__.py +16 -0
  29. hamilton/data_quality/base.py +198 -0
  30. hamilton/data_quality/default_validators.py +560 -0
  31. hamilton/data_quality/pandera_validators.py +121 -0
  32. hamilton/dataflows/__init__.py +726 -0
  33. hamilton/dataflows/template/README.md +25 -0
  34. hamilton/dataflows/template/__init__.py +54 -0
  35. hamilton/dataflows/template/author.md +28 -0
  36. hamilton/dataflows/template/requirements.txt +0 -0
  37. hamilton/dataflows/template/tags.json +7 -0
  38. hamilton/dataflows/template/valid_configs.jsonl +1 -0
  39. hamilton/dev_utils/__init__.py +16 -0
  40. hamilton/dev_utils/deprecation.py +204 -0
  41. hamilton/driver.py +2112 -0
  42. hamilton/execution/__init__.py +16 -0
  43. hamilton/execution/debugging_utils.py +56 -0
  44. hamilton/execution/executors.py +502 -0
  45. hamilton/execution/graph_functions.py +421 -0
  46. hamilton/execution/grouping.py +430 -0
  47. hamilton/execution/state.py +539 -0
  48. hamilton/experimental/__init__.py +27 -0
  49. hamilton/experimental/databackend.py +61 -0
  50. hamilton/experimental/decorators/__init__.py +16 -0
  51. hamilton/experimental/decorators/parameterize_frame.py +233 -0
  52. hamilton/experimental/h_async.py +29 -0
  53. hamilton/experimental/h_cache.py +413 -0
  54. hamilton/experimental/h_dask.py +28 -0
  55. hamilton/experimental/h_databackends.py +174 -0
  56. hamilton/experimental/h_ray.py +28 -0
  57. hamilton/experimental/h_spark.py +32 -0
  58. hamilton/function_modifiers/README +40 -0
  59. hamilton/function_modifiers/__init__.py +121 -0
  60. hamilton/function_modifiers/adapters.py +900 -0
  61. hamilton/function_modifiers/base.py +859 -0
  62. hamilton/function_modifiers/configuration.py +310 -0
  63. hamilton/function_modifiers/delayed.py +202 -0
  64. hamilton/function_modifiers/dependencies.py +246 -0
  65. hamilton/function_modifiers/expanders.py +1230 -0
  66. hamilton/function_modifiers/macros.py +1634 -0
  67. hamilton/function_modifiers/metadata.py +434 -0
  68. hamilton/function_modifiers/recursive.py +908 -0
  69. hamilton/function_modifiers/validation.py +289 -0
  70. hamilton/function_modifiers_base.py +31 -0
  71. hamilton/graph.py +1153 -0
  72. hamilton/graph_types.py +264 -0
  73. hamilton/graph_utils.py +41 -0
  74. hamilton/htypes.py +450 -0
  75. hamilton/io/__init__.py +32 -0
  76. hamilton/io/data_adapters.py +216 -0
  77. hamilton/io/default_data_loaders.py +224 -0
  78. hamilton/io/materialization.py +500 -0
  79. hamilton/io/utils.py +158 -0
  80. hamilton/lifecycle/__init__.py +67 -0
  81. hamilton/lifecycle/api.py +833 -0
  82. hamilton/lifecycle/base.py +1130 -0
  83. hamilton/lifecycle/default.py +802 -0
  84. hamilton/log_setup.py +47 -0
  85. hamilton/models.py +92 -0
  86. hamilton/node.py +449 -0
  87. hamilton/plugins/README.md +48 -0
  88. hamilton/plugins/__init__.py +16 -0
  89. hamilton/plugins/dask_extensions.py +47 -0
  90. hamilton/plugins/dlt_extensions.py +161 -0
  91. hamilton/plugins/geopandas_extensions.py +49 -0
  92. hamilton/plugins/h_dask.py +331 -0
  93. hamilton/plugins/h_ddog.py +522 -0
  94. hamilton/plugins/h_diskcache.py +163 -0
  95. hamilton/plugins/h_experiments/__init__.py +22 -0
  96. hamilton/plugins/h_experiments/__main__.py +62 -0
  97. hamilton/plugins/h_experiments/cache.py +39 -0
  98. hamilton/plugins/h_experiments/data_model.py +68 -0
  99. hamilton/plugins/h_experiments/hook.py +219 -0
  100. hamilton/plugins/h_experiments/server.py +375 -0
  101. hamilton/plugins/h_kedro.py +152 -0
  102. hamilton/plugins/h_logging.py +454 -0
  103. hamilton/plugins/h_mcp/__init__.py +28 -0
  104. hamilton/plugins/h_mcp/__main__.py +33 -0
  105. hamilton/plugins/h_mcp/_helpers.py +129 -0
  106. hamilton/plugins/h_mcp/_templates.py +417 -0
  107. hamilton/plugins/h_mcp/server.py +328 -0
  108. hamilton/plugins/h_mlflow.py +335 -0
  109. hamilton/plugins/h_narwhals.py +134 -0
  110. hamilton/plugins/h_openlineage.py +400 -0
  111. hamilton/plugins/h_opentelemetry.py +167 -0
  112. hamilton/plugins/h_pandas.py +257 -0
  113. hamilton/plugins/h_pandera.py +117 -0
  114. hamilton/plugins/h_polars.py +304 -0
  115. hamilton/plugins/h_polars_lazyframe.py +282 -0
  116. hamilton/plugins/h_pyarrow.py +56 -0
  117. hamilton/plugins/h_pydantic.py +127 -0
  118. hamilton/plugins/h_ray.py +242 -0
  119. hamilton/plugins/h_rich.py +142 -0
  120. hamilton/plugins/h_schema.py +493 -0
  121. hamilton/plugins/h_slack.py +100 -0
  122. hamilton/plugins/h_spark.py +1380 -0
  123. hamilton/plugins/h_threadpool.py +125 -0
  124. hamilton/plugins/h_tqdm.py +122 -0
  125. hamilton/plugins/h_vaex.py +129 -0
  126. hamilton/plugins/huggingface_extensions.py +236 -0
  127. hamilton/plugins/ibis_extensions.py +93 -0
  128. hamilton/plugins/jupyter_magic.py +622 -0
  129. hamilton/plugins/kedro_extensions.py +117 -0
  130. hamilton/plugins/lightgbm_extensions.py +99 -0
  131. hamilton/plugins/matplotlib_extensions.py +108 -0
  132. hamilton/plugins/mlflow_extensions.py +216 -0
  133. hamilton/plugins/numpy_extensions.py +105 -0
  134. hamilton/plugins/pandas_extensions.py +1763 -0
  135. hamilton/plugins/plotly_extensions.py +150 -0
  136. hamilton/plugins/polars_extensions.py +71 -0
  137. hamilton/plugins/polars_implementations.py +25 -0
  138. hamilton/plugins/polars_lazyframe_extensions.py +302 -0
  139. hamilton/plugins/polars_post_1_0_0_extensions.py +877 -0
  140. hamilton/plugins/polars_pre_1_0_0_extension.py +836 -0
  141. hamilton/plugins/pydantic_extensions.py +98 -0
  142. hamilton/plugins/pyspark_pandas_extensions.py +47 -0
  143. hamilton/plugins/sklearn_plot_extensions.py +129 -0
  144. hamilton/plugins/spark_extensions.py +105 -0
  145. hamilton/plugins/vaex_extensions.py +51 -0
  146. hamilton/plugins/xgboost_extensions.py +91 -0
  147. hamilton/plugins/yaml_extensions.py +89 -0
  148. hamilton/registry.py +254 -0
  149. hamilton/settings.py +18 -0
  150. hamilton/telemetry.py +50 -0
  151. 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
+ )