ENH:column-wise DataFrame.fillna and duplicated DataFrame.fillna with Series and Dict (#30922)

Author: proost

doc/source/whatsnew/v1.1.0.rst

@@ -43,7 +43,7 @@ Other enhancements
 
 - :class:`Styler` may now render CSS more efficiently where multiple cells have the same styling (:issue:`30876`)
 - When writing directly to a sqlite connection :func:`to_sql` now supports the ``multi`` method (:issue:`29921`)
--
+- :meth:`DataFrame.fillna` can fill NA values column-wise with a dictionary or :class:`Series` (:issue:`4514`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/core/generic.py

@@ -6001,19 +6001,25 @@ 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()
-                for k, v in value.items():
-                    if k not in result:
+                for i in range(result.columns.size):
+                    label = result.columns[i]
+
+                    if label not in value.keys():
                         continue
-                    obj = result[k]
-                    obj.fillna(v, limit=limit, inplace=True, downcast=downcast)
+
+                    result.iloc[:, i] = result.iloc[:, i].fillna(
+                        value[label], 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

@@ -602,9 +602,15 @@ def test_fillna_dict_series(self):
         expected = df.fillna(df.max().to_dict())
         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
@@ -703,3 +709,86 @@ def test_fill_value_when_combine_const(self):
         exp = df.fillna(0).add(2)
         res = df.add(2, fill_value=0)
         tm.assert_frame_equal(res, exp)
+
+    @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(
+        "fill_value",
+        [Series([100, 200, 300], index=[0, 1, 2]), {0: 100, 1: 200, 2: 300}],
+    )
+    def test_fillna_column_wise_duplicated_with_series_dict(self, fill_value):
+        # GH 4514
+        df = DataFrame(
+            [[np.nan, np.nan, 3], [np.nan, 5, np.nan], [7, np.nan, np.nan]],
+            columns=list("ABB"),
+            index=[0, 0, 1],
+        )
+        expected = DataFrame(
+            [[100, 100, 3], [100, 5, 100], [7, 200, 200]],
+            columns=list("ABB"),
+            index=[0, 0, 1],
+            dtype="float64",
+        )
+
+        result = df.fillna(fill_value, axis=1)
+        tm.assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "fill_value",
+        [
+            Series([100, 200, 300], index=["A", "B", "C"]),
+            {"A": 100, "B": 200, "C": 300},
+        ],
+    )
+    def test_fillna_duplicated_with_series_dict(self, fill_value):
+        # GH 4514
+        df = DataFrame(
+            [[np.nan, np.nan, 3], [np.nan, 5, np.nan], [7, np.nan, np.nan]],
+            columns=list("ABB"),
+            index=[0, 0, 1],
+        )
+        expected = DataFrame(
+            [[100, 200, 3], [100, 5, 200], [7, 200, 200]],
+            columns=list("ABB"),
+            index=[0, 0, 1],
+            dtype="float64",
+        )
+
+        result = df.fillna(fill_value)
+        tm.assert_frame_equal(result, expected)