Fix extract_columns fill_with for columns passed as (name, doc) tuples

hamilton/function_modifiers/expanders.py

@@ -646,14 +646,16 @@ 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):
 
             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

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"""