import dataclasses
import datetime
from typing import Optional

import pandas as pd
import pyarrow.parquet as pq

@dataclasses.dataclass(frozen=True)
class LoadTable:
    table_name: str
    columns: Optional[list[str]]

    def __call__(self, dates: list[datetime.date], filters: list) -> pd.DataFrame:
        files = [
            f"{self.table_name}/{date:%Y-%m-%d}.parquet"
            for date in dates
        ]
        return pq.ParquetDataset(files, filters=filters).read(columns=self.columns).to_pandas()


@dataclasses.dataclass(frozen=True)
class MergeAsOf:
    on: str
    by: str
    tolerance: Optional[pd.Timedelta]

    def __call__(self, left, right):
        pd.merge_asof(
            left, right,
            on=self.on,
            by=<http://self.by|self.by>,
            tolerance=self.tolerance,
            allow_exact_matches=True
        )


dag = Dag()
dates = dag.add_input("dates")
pokemon_names = dag.add_input("pokemon_names")
feature_1 = dag.add_node(LoadTable("pokemon-colour-features", ['red', "green", "blue"])).map(
    dates, pokemon_names
)
feature_2 = dag.add_node(LoadTable("pokemon-power-features", ['water_power', "fire_power"])).map(
    dates, pokemon_names
)
both_features = dag.add_node(MergeAsOf(on="timestamp", by="pokemon_name"))(
    feature_1, feature_2
)

dag.set_output(both_features)


features = dag.execute(
    dates=pd.date_range(pd.to_datetime("2023-01-01"), pd.to_datetime("2023-07-24")),
    pokemon_names=(["fluffington", "pikapkia"])
)

def get_mean(data, frequency: str):
    return (
        data
        .assign(date= lambda x: x.ceil(pd.to_timedelta(frequency)))
        .groupby("date")
        [["price"]].mean()
        .assign(frequency=frequency)
    )


source_data = dag.add_node(load_data)
means = []
look_backs = ["1d", "1w", "30d"]
for look_back in look_backs:
    look_back_node = dag.const(look_back)
    mean_by_date = dag.add_node(get_mean).map(source_data, look_back_node)
    means.append(mean_by_date)

dag.add_node(lambda *x: pd.concat(*x)).map(*means)

Node

def

@dataclasses.dataclass(frozen=True)
class DataLoader:
    table_name: str

    def __call__(self, dates):
        """do your thing..."""


color_loader = DataLoader("color-features")
power_features = DataLoader("power-features")

driver = Driver({}, custom_functions={"color_features": DataLoader("color-features") })

def source_data(url: str) -> pd.Series:
    """You can also use load_from... and define a custom data loader for extra readability/metadata. Or parameterize."""
    return _load_data(url)[COLUMN] # this could also return a dataframe

LOOKBACKS = ["1d", "1w", "30d"]

@parameterize(
    {f"mean_{lookback}" : {"lookback" : value(lookback)} for lookback in LOOKBACKS}
)
def mean_lookback_n_days(data: pd.Series, lookback: str) -> pd.Series:
    return ...

@inject(lookbacks=group(*[f"mean_{lookback}" for lookback in LOOKBACKS]))
def data_with_mean_lookbacks(lookbacks: List[pd.Series]) -> pd.DataFrame:
    return _join_logic(lookbacks)

lookbacks

inject

@parameterize

lookback

def source_data(url: str) -> pd.Series:
    """You can also use load_from... and define a custom data loader for extra readability/metadata. Or parameterize."""
    return _load_data(url)[COLUMN] # this could also return a dataframe


@resolve(
  when=ResolveAt.CONFIG_AVAILABLE,
  # config now expects a lookbacks parameter that's a list of ints
  decorate_with=lambda lookbacks: parameterize(
    {f"mean_{lookback}" : {"lookback" : value(lookback)} for lookback in lookbacks}
    )
)
def mean_lookback_n_days(data: pd.Series, lookback: str) -> pd.Series:
    return ...

@resolve(
    when=ResolveAt.CONFIG_AVAILABLE,
    decorate_with=lambda lookbacks: inject(lookbacks=group(*[f"mean_{lookback}" for lookback in LOOKBACKS]))
)
def data_with_mean_lookbacks(lookbacks: List[pd.Series]) -> pd.DataFrame:
    return _join_logic(lookbacks)

@repeat_for_lookbacks(pattern="mean_{lookback}")
@inject_lookbacks(pattern="mean_{lookback}")

You could pass around dataframes and have

lookbacks

be a runtime parameter — just run the operation over the entire dataframe. No

inject

or

@parameterize

— you can even take in default arguments/send in overrides.

import sys

import pandas as pd
from hamilton.driver import Driver


def source_data() -> pd.Series:
    return pd.Series(
        pd.to_datetime(
            [
                "2021-02-21T20:26:00Z",
                "2021-02-22T20:26:00Z",
                "2021-02-23T20:26:00Z",
                "2021-02-24T20:26:00Z",
            ]
        )
    )


def data_with_lookbacks(source_data: pd.Series, lookbacks: list[str]) -> pd.DataFrame:
    return pd.concat(
        [
            (source_data - pd.to_timedelta(lookback)).rename(lookback)
            for lookback in lookbacks
        ],
        axis=1,
    )


driver = Driver({}, sys.modules[__name__])
data = driver.execute(
    ["data_with_lookbacks"], inputs={"lookbacks": ["1d", "1w", "30d"]}
)
print(data.to_markdown(index=False))