def input_lf() -> pl.LazyFrame:
    return pl.scan_csv("example_customer_data.csv")


@resolve_from_config(
    decorate_with=lambda column_config: parameterize_values(
        parameter="column_names",
        assigned_output={
            (key + "_required_columns", "test"): [
                column_config[key]["variable"],
                *column_config[key].get("segment", []),
            ]
            for key in column_config
        },
    )
)
def required_columns(input_lf: pl.LazyFrame, column_names: list[str]) -> pl.LazyFrame:
    return input_lf.select(column_names)


@resolve_from_config(
    decorate_with=lambda column_config: parameterize(
        **{
            col: {
                "input_lf": source(col + "_required_columns"),
                "group_by_col": value(column_config[col].get("segment", None)),
            }
            for col in column_config
        }
    )
)
def group_by(input_lf: pl.LazyFrame, group_by_col: str) -> LazyGroupBy:
    return input_lf.group_by(group_by_col)