Add decorator to value mapper (#160)

* add decorator to value mapper * simplify
nubank · Feb 9, 2021 · d5a1723 · d5a1723
1 parent cf0cd0f
commit d5a1723
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/src/fklearn/training/transformation.py b/src/fklearn/training/transformation.py
@@ -361,6 +361,7 @@ def apply_replacements(df: pd.DataFrame,
     return df.assign(**categ_columns)
 
 
+@column_duplicatable('value_maps')
 @curry
 @log_learner_time(learner_name="value_mapper")
 def value_mapper(df: pd.DataFrame,

diff --git a/tests/training/test_transformation.py b/tests/training/test_transformation.py
@@ -173,6 +173,12 @@ def test_value_mapper():
     }
 
     pred_fn, data_ignore, log = value_mapper(input_df, value_maps)
+    pred_fn2, data_ignore2, log2 = value_mapper(input_df, value_maps, suffix="_suffix")
+    pred_fn3, data_ignore3, log3 = value_mapper(input_df, value_maps, prefix="prefix_")
+    pred_fn4, data_ignore4, log4 = value_mapper(input_df, value_maps,
+                                                columns_mapping={"feat1": "feat1_raw",
+                                                                 "feat2": "feat2_raw",
+                                                                 "feat3": "feat3_raw"})
     pred_fn, data_not_ignore, log = value_mapper(
         input_df, value_maps, ignore_unseen=False
     )
@@ -196,6 +202,18 @@ def test_value_mapper():
     assert expected_ignore.equals(data_ignore)
     assert expected_not_ignore.equals(data_not_ignore)
 
+    assert pd.concat(
+        [expected_ignore, input_df.copy().add_suffix("_suffix")], axis=1
+    ).equals(data_ignore2)
+
+    assert pd.concat(
+        [expected_ignore, input_df.copy().add_prefix("prefix_")], axis=1
+    ).equals(data_ignore3)
+
+    assert pd.concat(
+        [expected_ignore, input_df.copy().add_suffix("_raw")], axis=1
+    ).equals(data_ignore4)
+
 
 def test_truncate_categorical():
     input_df_train = pd.DataFrame(