diff --git a/hamilton/function_modifiers/expanders.py b/hamilton/function_modifiers/expanders.py index 0cfe1ea62..2cd3309eb 100644 --- a/hamilton/function_modifiers/expanders.py +++ b/hamilton/function_modifiers/expanders.py @@ -646,6 +646,8 @@ def transform_node( """ fn = node_.callable base_doc = node_.documentation + # columns can be passed as (name, doc) tuples -- only the name is a column in the dataframe + column_names = [col[0] if isinstance(col, tuple) else col for col in self.columns] # if fn is an async function if inspect.iscoroutinefunction(fn): @@ -653,7 +655,7 @@ def transform_node( async def df_generator(*args, **kwargs) -> Any: df_generated = await fn(*args, **kwargs) if self.fill_with is not None: - for col in self.columns: + for col in column_names: if col not in df_generated: registry.fill_with_scalar(df_generated, col, self.fill_with) assert col in df_generated @@ -664,7 +666,7 @@ async def df_generator(*args, **kwargs) -> Any: def df_generator(*args, **kwargs) -> Any: df_generated = fn(*args, **kwargs) if self.fill_with is not None: - for col in self.columns: + for col in column_names: if col not in df_generated: registry.fill_with_scalar(df_generated, col, self.fill_with) assert col in df_generated diff --git a/tests/function_modifiers/test_expanders.py b/tests/function_modifiers/test_expanders.py index fa11b5c05..393838419 100644 --- a/tests/function_modifiers/test_expanders.py +++ b/tests/function_modifiers/test_expanders.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import asyncio import sys from typing import Any, TypedDict @@ -294,6 +295,37 @@ def dummy_df() -> pd.DataFrame: ) # it has to be in there now +def test_column_extractor_fill_with_documented_columns(): + def dummy_df() -> pd.DataFrame: + """dummy doc""" + return pd.DataFrame({"col_1": [1, 2, 3, 4]}) + + annotation = function_modifiers.extract_columns( + ("col_1", "col_1 doc"), ("col_3", "col_3 doc"), fill_with=0 + ) + original_node, _, col_3_node = annotation.transform_node( + node.Node.from_fn(dummy_df), {}, dummy_df + ) + original_df = original_node.callable() + assert list(original_df.columns) == ["col_1", "col_3"] + pd.testing.assert_series_equal( + col_3_node.callable(dummy_df=original_df), pd.Series([0, 0, 0, 0]), check_names=False + ) + + +def test_column_extractor_fill_with_documented_columns_async(): + async def dummy_df() -> pd.DataFrame: + """dummy doc""" + return pd.DataFrame({"col_1": [1, 2, 3, 4]}) + + annotation = function_modifiers.extract_columns( + ("col_1", "col_1 doc"), ("col_3", "col_3 doc"), fill_with=0 + ) + original_node, _, _ = annotation.transform_node(node.Node.from_fn(dummy_df), {}, dummy_df) + original_df = asyncio.run(original_node.callable()) + assert list(original_df.columns) == ["col_1", "col_3"] + + def test_column_extractor_no_fill_with(): def dummy_df_generator() -> pd.DataFrame: """dummy doc"""