From d0442a0f6496c2ecd794d9ef03e08409828f9ae2 Mon Sep 17 00:00:00 2001 From: breken-ai <312387581+breken-ai@users.noreply.github.com> Date: Fri, 25 Sep 2026 03:13:00 -0700 Subject: [PATCH] Respect reassign_columns=False in parameterize_extract_columns The node handed to @parameterize always wrapped the function in the column-renaming wrapper, so reassign_columns=False had no effect: the output columns were still renamed by position, and an extracted column could silently carry another column's data. Build the node from the function chosen by reassign_columns. Generated-by: Claude Code (Claude Opus 5.5) --- hamilton/function_modifiers/expanders.py | 2 +- tests/function_modifiers/test_expanders.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) 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)