blob: 6aa1ef3cde6cfacb2cc6d89c2f9b27607b177fe0 [file] [log] [blame]
import pandas as pd
from hamilton.function_modifiers import ResolveAt, extract_columns, parameterize, resolve, source
@extract_columns("a", "b", "c", "d", "e")
def df() -> pd.DataFrame:
"""Produces a dataframe with columns a, b, c, d, e consisting entirely of 1s"""
return pd.DataFrame(
{
"a": [1],
"b": [1],
"c": [1],
"d": [1],
"e": [1],
}
)
@resolve(
when=ResolveAt.CONFIG_AVAILABLE,
decorate_with=lambda columns_to_sum_map: parameterize(
**{
key: {"col_1": source(value[0]), "col_2": source(value[1])}
for key, value in columns_to_sum_map.items()
}
),
)
def generic_summation(col_1: pd.Series, col_2: pd.Series) -> pd.Series:
return col_1 + col_2
# TODO -- add grouping then we can test everything