diff --git a/hamilton/function_modifiers/expanders.py b/hamilton/function_modifiers/expanders.py index 0cfe1ea62..f705a504a 100644 --- a/hamilton/function_modifiers/expanders.py +++ b/hamilton/function_modifiers/expanders.py @@ -1168,8 +1168,8 @@ def wrapper_fn(*args, _output_columns=parameterization.outputs, **kwargs): df_out.columns = _output_columns return df_out - new_node = node_.copy_with(callabl=wrapper_fn) fn_to_call = wrapper_fn if self.reassign_columns else fn + new_node = node_.copy_with(callabl=fn_to_call) # We have to rename the underlying function so that we do not # get naming collisions. Using __ is cleaner than using a uuid # as it is easier to read/manage and naturally maeks sense. diff --git a/tests/function_modifiers/test_expanders.py b/tests/function_modifiers/test_expanders.py index fa11b5c05..022d7d478 100644 --- a/tests/function_modifiers/test_expanders.py +++ b/tests/function_modifiers/test_expanders.py @@ -918,6 +918,25 @@ def fn(input1: pd.Series, input2: pd.Series, input3: float) -> pd.DataFrame: assert nodes_by_name["outseries2b"](fn__1=pd.DataFrame({"outseries2b": [40]}))[0] == 40 +def test_parameterized_extract_columns_no_reassign_columns(): + """With reassign_columns=False the columns keep their names, so they are extracted by name.""" + annotation = function_modifiers.parameterize_extract_columns( + function_modifiers.ParameterizedExtract(("y", "x"), {"k": value(1)}), + reassign_columns=False, + ) + + def fn(k: int) -> pd.DataFrame: + return pd.DataFrame({"x": [k], "y": [k * 100]}) + + nodes_by_name = { + node_.name: node_ for node_ in annotation.expand_node(node.Node.from_fn(fn), {}, fn) + } + df = nodes_by_name["fn__0"]() + pd.testing.assert_frame_equal(df, pd.DataFrame({"x": [1], "y": [100]})) + assert nodes_by_name["x"](fn__0=df)[0] == 1 + assert nodes_by_name["y"](fn__0=df)[0] == 100 + + def test_parametrized_full_replace_groups_with_literal(): def add_n(grouped_parameter: list[int]) -> int: return sum(grouped_parameter)