ENH:column-wise DataFrame.fillna with Series and Dict (pandas-dev#4514)

Author: proost

doc/source/whatsnew/v1.1.0.rst

@@ -18,7 +18,7 @@ Enhancements
 Other enhancements
 ^^^^^^^^^^^^^^^^^^
 
--
+- :meth:`DataFrame.fillna` can fill NA values column-wise with a dictionary or :class:`Series` (:issue:`4514`)
 -
 
 

pandas/core/generic.py

@@ -5960,19 +5960,39 @@ def fillna(
                 )
 
             elif isinstance(value, (dict, ABCSeries)):
+                result = self if inplace else self.copy()
+
                 if axis == 1:
-                    raise NotImplementedError(
-                        "Currently only can fill "
-                        "with dict/Series column "
-                        "by column"
-                    )
+                    # To access column base
+                    result = result.T
 
-                result = self if inplace else self.copy()
+                columns = result.columns
                 for k, v in value.items():
-                    if k not in result:
+                    if k not in columns.array:
                         continue
-                    obj = result[k]
-                    obj.fillna(v, limit=limit, inplace=True, downcast=downcast)
+
+                    column_locs = columns.get_loc(k)
+
+                    if is_scalar(column_locs):
+                        # if label is not duplicated
+                        result.loc[:, k] = result.loc[:, k].fillna(
+                            v, limit=limit, inplace=False, downcast=downcast
+                        )
+
+                    else:
+                        if isinstance(column_locs, slice):
+                            locs = range(column_locs.start, column_locs.stop)
+                        else:
+                            locs = np.where(column_locs)[0]
+
+                        for i in locs:
+                            result.iloc[:, i] = result.iloc[:, i].fillna(
+                                v, limit=limit, inplace=False, downcast=downcast
+                            )
+
+                if axis == 1:
+                    result = result.T
+
                 return result if not inplace else None
 
             elif not is_list_like(value):

pandas/tests/frame/test_missing.py

@@ -605,8 +605,15 @@ def test_fillna_dict_series(self):
         tm.assert_frame_equal(result, expected)
 
         # disable this for now
-        with pytest.raises(NotImplementedError, match="column by column"):
-            df.fillna(df.max(1), axis=1)
+        expected = DataFrame(
+            {
+                "a": [1.0, 1.0, 2.0, 3.0, 4.0],
+                "b": [1.0, 2.0, 3.0, 3.0, 4.0],
+                "c": [1.0, 1.0, 2.0, 3.0, 4.0],
+            }
+        )
+        result = df.fillna(df.max(1), axis=1)
+        tm.assert_frame_equal(expected, result)
 
     def test_fillna_dataframe(self):
         # GH 8377
@@ -983,3 +990,125 @@ def test_interp_time_inplace_axis(self, axis):
         result = expected.interpolate(axis=0, method="time")
         expected.interpolate(axis=0, method="time", inplace=True)
         tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "expected,fill_value",
+        [
+            (
+                DataFrame(
+                    [[100, 100], [200, 4], [5, 6]], columns=list("AB"), dtype="float64"
+                ),
+                Series([100, 200, 300]),
+            ),
+            (
+                DataFrame(
+                    [[100, 100], [np.nan, 4], [5, 6]],
+                    columns=list("AB"),
+                    dtype="float64",
+                ),
+                {0: 100, 2: 300, 3: 400},
+            ),
+        ],
+    )
+    def test_fillna_column_wise(self, expected, fill_value):
+        # GH 4514
+        df = DataFrame([[np.nan, np.nan], [np.nan, 4], [5, 6]], columns=list("AB"))
+        result = df.fillna(fill_value, axis=1)
+        tm.assert_frame_equal(expected, result)
+
+        df.fillna(fill_value, axis=1, inplace=True)
+        tm.assert_frame_equal(expected, df)
+
+    def test_fillna_column_wise_downcast(self):
+        df = DataFrame([[np.nan, 2], [3, np.nan], [np.nan, np.nan]], columns=list("AB"))
+        s = Series([100, 200, 300])
+
+        expected = DataFrame(
+            [[100, 2], [3, 200], [300, 300]], columns=list("AB"), dtype="int64"
+        )
+        result = df.fillna(s, axis=1, downcast="infer")
+        tm.assert_frame_equal(expected, result)
+
+    @pytest.mark.parametrize(
+        "target,expected",
+        [
+            (
+                DataFrame(
+                    [[np.nan, np.nan, 3], [np.nan, 5, np.nan], [7, np.nan, np.nan]],
+                    columns=list("ABB"),
+                    index=[0, 0, 1],
+                ),
+                DataFrame(
+                    [[100, 100, 3], [100, 5, 100], [7, 200, 200]],
+                    columns=list("ABB"),
+                    index=[0, 0, 1],
+                    dtype="float64",
+                ),
+            ),
+            (
+                DataFrame(
+                    [[np.nan, np.nan, 3], [np.nan, 5, np.nan], [7, np.nan, np.nan]],
+                    columns=list("ABA"),
+                    index=[0, 1, 0],
+                ),
+                DataFrame(
+                    [[100, 100, 3], [200, 5, 200], [7, 100, 100]],
+                    columns=list("ABA"),
+                    index=[0, 1, 0],
+                    dtype="float64",
+                ),
+            ),
+        ],
+    )
+    def test_fillna_column_wise_duplicated_with_series_dict(self, target, expected):
+        # GH 4514
+        s = pd.Series([100, 200, 300], index=[0, 1, 2])
+        d = {0: 100, 1: 200, 2: 300}
+
+        result = target.fillna(s, axis=1)
+        tm.assert_frame_equal(result, expected)
+
+        result = target.fillna(d, axis=1)
+        tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "target,expected",
+        [
+            (
+                DataFrame(
+                    [[np.nan, np.nan, 3], [np.nan, 5, np.nan], [7, np.nan, np.nan]],
+                    columns=list("ABB"),
+                    index=[0, 0, 1],
+                ),
+                DataFrame(
+                    [[100, 200, 3], [100, 5, 200], [7, 200, 200]],
+                    columns=list("ABB"),
+                    index=[0, 0, 1],
+                    dtype="float64",
+                ),
+            ),
+            (
+                DataFrame(
+                    [[np.nan, np.nan, 3], [np.nan, 5, np.nan], [7, np.nan, np.nan]],
+                    columns=list("ABA"),
+                    index=[0, 1, 0],
+                ),
+                DataFrame(
+                    [[100, 200, 3], [100, 5, 100], [7, 200, 100]],
+                    columns=list("ABA"),
+                    index=[0, 1, 0],
+                    dtype="float64",
+                ),
+            ),
+        ],
+    )
+    def test_fillna_duplicated_with_series_dict(self, target, expected):
+        # GH 4514
+        s = pd.Series([100, 200, 300], index=["A", "B", "C"])
+        d = {"A": 100, "B": 200, "C": 300}
+
+        result = target.fillna(s)
+        tm.assert_frame_equal(result, expected)
+
+        result = target.fillna(d)
+        tm.assert_frame_equal(result, expected)