Merge pull request #25 from big-o/develop

column prefix bug fix in input passthroughs

Author: big-o

skdag/dag/_dag.py

@@ -33,12 +33,12 @@
 __all__ = ["DAG", "DAGStep"]
 
 
-def _get_columns(X, dep, cols, is_root, axis=1):
+def _get_columns(X, dep, cols, is_root, dep_is_passthrough, axis=1):
     if callable(cols):
         # sklearn.compose.make_column_selector
         cols = cols(X)
 
-    if not is_root:
+    if not is_root and not dep_is_passthrough:
         # The DAG will prepend output columns with the step name, so add this in to any
         # dep columns if missing. This helps keep user-provided deps readable.
         if isinstance(cols, str):
@@ -60,7 +60,14 @@ def _stack_inputs(dag, X, node):
     deps = {node.name: None} if node.is_root else node.deps
 
     cols = [
-        _get_columns(X[dep], dep, cols, node.is_root, axis=1)
+        _get_columns(
+            X[dep],
+            dep,
+            cols,
+            node.is_root,
+            _is_passthrough(dag.graph_.nodes[dep]["step"].estimator),
+            axis=1,
+        )
         for dep, cols in deps.items()
     ]
 

skdag/dag/tests/test_dag.py

@@ -407,13 +407,20 @@ def test_pandas(X, y, steps):
     assert np.allclose(y_pred_np, y_pred_pd)
 
 
-def test_pandas_indexing():
+@pytest.mark.parametrize("input_passthrough", [False, True])
+def test_pandas_indexing(input_passthrough):
     X, y = datasets.load_diabetes(return_X_y=True, as_frame=True)
 
     passcols = ["age", "sex", "bmi", "bp"]
+
+    builder = DAGBuilder(infer_dataframe=True)
+    if input_passthrough:
+        builder.add_step("inp", "passthrough")
+
     preprocessing = (
-        DAGBuilder(infer_dataframe=True)
-        .add_step("imp", SimpleImputer())
+        builder.add_step(
+            "imp", SimpleImputer(), deps=["inp"] if input_passthrough else None
+        )
         .add_step("vitals", "passthrough", deps={"imp": passcols})
         .add_step(
             "blood",