Repository: NVIDIA-Merlin/NVTabular

Issue:
Fill Missing Op option to create a new binary column indicating the value was replaced.
Fill Missing Op should have the option to create a new binary column indicating whether the column was filled or not for continuous variables.

This is a common feature used when dealing with missing values of categoricals.


Discussion:
+1

Resolving patch:
diff --git a/nvtabular/ops/fill.py b/nvtabular/ops/fill.py
--- a/nvtabular/ops/fill.py
+++ b/nvtabular/ops/fill.py
@@ -36,18 +36,34 @@ class FillMissing(Operator):
     -----------
     fill_val : float, default 0
         The constant value to replace missing values with.
+    add_binary_cols : boolean, default False
+        When True, adds binary columns that indicate whether cells in each column were filled
     """
 
-    def __init__(self, fill_val=0):
+    def __init__(self, fill_val=0, add_binary_cols=False):
         super().__init__()
         self.fill_val = fill_val
+        self.add_binary_cols = add_binary_cols
 
     @annotate("FillMissing_op", color="darkgreen", domain="nvt_python")
     def transform(self, columns, gdf: cudf.DataFrame) -> cudf.DataFrame:
-        return gdf[columns].fillna(self.fill_val)
+        if self.add_binary_cols:
+            for col in columns:
+                gdf[f"{col}_filled"] = gdf[col].isna()
+                gdf[col] = gdf[col].fillna(self.fill_val)
+        else:
+            gdf[columns] = gdf[columns].fillna(self.fill_val)
+
+        return gdf
 
     transform.__doc__ = Operator.transform.__doc__
 
+    def output_column_names(self, columns: ColumnNames) -> ColumnNames:
+        output_cols = columns[:]
+        if self.add_binary_cols:
+            output_cols.extend([f"{col}_filled" for col in columns])
+        return output_cols
+
 
 class FillMedian(StatOperator):
     """
@@ -64,10 +80,16 @@ class FillMedian(StatOperator):
 
         # Add FillMedian to the workflow for continuous columns
         proc.add_cont_feature(nvt.ops.FillMedian())
+
+    Parameters
+    -----------
+    add_binary_cols : boolean, default False
+        When True, adds binary columns that indicate whether cells in each column were filled
     """
 
-    def __init__(self):
+    def __init__(self, add_binary_cols=False):
         super().__init__()
+        self.add_binary_cols = add_binary_cols
         self.medians = {}
 
     @annotate("FillMedian_transform", color="darkgreen", domain="nvt_python")
@@ -76,6 +98,8 @@ def transform(self, columns: ColumnNames, gdf: cudf.DataFrame) -> cudf.DataFrame
             raise RuntimeError("need to call 'fit' before running transform")
 
         for col in columns:
+            if self.add_binary_cols:
+                gdf[f"{col}_filled"] = gdf[col].isna()
             gdf[col] = gdf[col].fillna(self.medians[col])
         return gdf
 
@@ -96,3 +120,9 @@ def fit_finalize(self, dask_stats):
 
     def clear(self):
         self.medians = {}
+
+    def output_column_names(self, columns: ColumnNames) -> ColumnNames:
+        output_cols = columns[:]
+        if self.add_binary_cols:
+            output_cols.extend([f"{col}_filled" for col in columns])
+        return output_cols


Test patch:
diff --git a/tests/unit/test_ops.py b/tests/unit/test_ops.py
--- a/tests/unit/test_ops.py
+++ b/tests/unit/test_ops.py
@@ -134,8 +134,9 @@ def test_target_encode_multi(tmpdir, npartitions):
 @pytest.mark.parametrize("gpu_memory_frac", [0.01, 0.1])
 @pytest.mark.parametrize("engine", ["parquet", "csv", "csv-no-header"])
 @pytest.mark.parametrize("op_columns", [["x"], ["x", "y"]])
-def test_fill_median(tmpdir, df, dataset, gpu_memory_frac, engine, op_columns):
-    cont_features = op_columns >> nvt.ops.FillMedian()
+@pytest.mark.parametrize("add_binary_cols", [True, False])
+def test_fill_median(tmpdir, df, dataset, gpu_memory_frac, engine, op_columns, add_binary_cols):
+    cont_features = op_columns >> nvt.ops.FillMedian(add_binary_cols=add_binary_cols)
     processor = nvt.Workflow(cont_features)
     processor.fit(dataset)
     new_gdf = processor.transform(dataset).to_ddf().compute()
@@ -144,6 +145,9 @@ def test_fill_median(tmpdir, df, dataset, gpu_memory_frac, engine, op_columns):
         col_median = df[col].dropna().quantile(0.5, interpolation="linear")
         assert math.isclose(col_median, processor.column_group.op.medians[col], rel_tol=1e1)
         assert np.all((df[col].fillna(col_median) - new_gdf[col]).abs().values <= 1e-2)
+        assert (f"{col}_filled" in new_gdf.keys()) == add_binary_cols
+        if add_binary_cols:
+            assert df[col].isna().sum() == new_gdf[f"{col}_filled"].sum()
 
 
 @pytest.mark.parametrize("gpu_memory_frac", [0.01, 0.1])
@@ -207,9 +211,10 @@ def test_hash_bucket_lists(tmpdir):
 
 
 @pytest.mark.parametrize("engine", ["parquet"])
-def test_fill_missing(tmpdir, df, dataset, engine):
+@pytest.mark.parametrize("add_binary_cols", [True, False])
+def test_fill_missing(tmpdir, df, dataset, engine, add_binary_cols):
     cont_names = ["x", "y"]
-    cont_features = cont_names >> nvt.ops.FillMissing(fill_val=42)
+    cont_features = cont_names >> nvt.ops.FillMissing(fill_val=42, add_binary_cols=add_binary_cols)
 
     for col in cont_names:
         idx = np.random.choice(df.shape[0] - 1, int(df.shape[0] * 0.2))
@@ -220,9 +225,13 @@ def test_fill_missing(tmpdir, df, dataset, engine):
     processor = nvt.Workflow(cont_features)
     processor.fit(dataset)
     new_gdf = processor.transform(dataset).to_ddf().compute()
+
     for col in cont_names:
         assert np.all((df[col].fillna(42) - new_gdf[col]).abs().values <= 1e-2)
         assert new_gdf[col].isna().sum() == 0
+        assert (f"{col}_filled" in new_gdf.keys()) == add_binary_cols
+        if add_binary_cols:
+            assert df[col].isna().sum() == new_gdf[f"{col}_filled"].sum()
 
 
 @pytest.mark.parametrize("engine", ["parquet"])