Fix DataFrame.replace with NaN/None values #1907

databricks/koalas/series.py

@@ -4264,7 +4264,7 @@ def replace(self, to_replace=None, value=None, regex=False) -> "Series":
         dtype: float64
         """
         if to_replace is None:
-            return self
+            return self.fillna(method="ffill")
         if not isinstance(to_replace, (str, list, dict, int, float)):
             raise ValueError("'to_replace' should be one of str, list, dict, int, float")
         if regex:
@@ -4283,14 +4283,23 @@ def replace(self, to_replace=None, value=None, regex=False) -> "Series":
                 current = self.spark.column
             else:
                 for to_replace_, value in to_replace.items():
+                    cond = (
+                        (F.isnan(self.spark.column) | self.spark.column.isNull())
+                        if pd.isna(to_replace_)
+                        else (self.spark.column == F.lit(to_replace_))
+                    )
                     if is_start:
-                        current = F.when(self.spark.column == F.lit(to_replace_), value)
+                        current = F.when(cond, value)
                         is_start = False
                     else:
-                        current = current.when(self.spark.column == F.lit(to_replace_), value)
+                        current = current.when(cond, value)
                 current = current.otherwise(self.spark.column)
         else:
-            current = F.when(self.spark.column.isin(to_replace), value).otherwise(self.spark.column)
+            cond = self.spark.column.isin(to_replace)
+            # to_replace may be a scalar
+            if np.array(pd.isna(to_replace)).any():
+                cond = cond | F.isnan(self.spark.column) | self.spark.column.isNull()
+            current = F.when(cond, value).otherwise(self.spark.column)
 
         return self._with_new_scol(current)
 

databricks/koalas/tests/test_dataframe.py

@@ -2385,7 +2385,7 @@ def test_replace(self):
         self.assert_eq(kser, pser)
 
         pdf = pd.DataFrame(
-            {"A": [0, 1, 2, 3, 4], "B": [5, 6, 7, 8, 9], "C": ["a", "b", "c", "d", "e"]},
+            {"A": [0, 1, 2, 3, np.nan], "B": [5, 6, 7, 8, np.nan], "C": ["a", "b", "c", "d", None]},
             index=np.random.rand(5),
         )
         kdf = ks.from_pandas(pdf)
@@ -2399,10 +2399,21 @@ def test_replace(self):
 
         self.assert_eq(kdf.replace({0: 10, 1: 100, 7: 200}), pdf.replace({0: 10, 1: 100, 7: 200}))
 
-        self.assert_eq(kdf.replace({"A": 0, "B": 5}, 100), pdf.replace({"A": 0, "B": 5}, 100))
+        self.assert_eq(
+            kdf.replace({"A": [0, np.nan], "B": [5, np.nan]}, 100),
+            pdf.replace({"A": [0, np.nan], "B": [5, np.nan]}, 100),
+        )
+
+        self.assert_eq(
+            kdf.replace({"A": {0: 100, 4: 400, np.nan: 700}}),
+            pdf.replace({"A": {0: 100, 4: 400, np.nan: 700}}),
+        )
+        self.assert_eq(
+            kdf.replace({"X": {0: 100, 4: 400, np.nan: 700}}),
+            pdf.replace({"X": {0: 100, 4: 400, np.nan: 700}}),
+        )
 
-        self.assert_eq(kdf.replace({"A": {0: 100, 4: 400}}), pdf.replace({"A": {0: 100, 4: 400}}))
-        self.assert_eq(kdf.replace({"X": {0: 100, 4: 400}}), pdf.replace({"X": {0: 100, 4: 400}}))
+        self.assert_eq(kdf.replace({"C": ["a", None]}, "e"), pdf.replace({"C": ["a", None]}, "e"))
 
         # multi-index columns
         columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")])
@@ -2419,15 +2430,21 @@ def test_replace(self):
         self.assert_eq(kdf.replace({0: 10, 1: 100, 7: 200}), pdf.replace({0: 10, 1: 100, 7: 200}))
 
         self.assert_eq(
-            kdf.replace({("X", "A"): 0, ("X", "B"): 5}, 100),
-            pdf.replace({("X", "A"): 0, ("X", "B"): 5}, 100),
+            kdf.replace({("X", "A"): [0, np.nan], ("X", "B"): 5}, 100),
+            pdf.replace({("X", "A"): [0, np.nan], ("X", "B"): 5}, 100),
         )
 
         self.assert_eq(
-            kdf.replace({("X", "A"): {0: 100, 4: 400}}), pdf.replace({("X", "A"): {0: 100, 4: 400}})
+            kdf.replace({("X", "A"): {0: 100, 4: 400, np.nan: 700}}),
+            pdf.replace({("X", "A"): {0: 100, 4: 400, np.nan: 700}}),
         )
         self.assert_eq(
-            kdf.replace({("X", "B"): {0: 100, 4: 400}}), pdf.replace({("X", "B"): {0: 100, 4: 400}})
+            kdf.replace({("X", "B"): {0: 100, 4: 400, np.nan: 700}}),
+            pdf.replace({("X", "B"): {0: 100, 4: 400, np.nan: 700}}),
+        )
+
+        self.assert_eq(
+            kdf.replace({("Y", "C"): ["a", None]}, "e"), pdf.replace({("Y", "C"): ["a", None]}, "e")
         )
 
     def test_update(self):

databricks/koalas/tests/test_series.py

@@ -1422,12 +1422,14 @@ def test_pop(self):
             kser.pop(("lama", "speed", "x"))
 
     def test_replace(self):
-        pser = pd.Series([10, 20, 15, 30, 45], name="x")
+        pser = pd.Series([10, 20, 15, 30, np.nan], name="x")
         kser = ks.Series(pser)
 
         self.assert_eq(kser.replace(), pser.replace())
         self.assert_eq(kser.replace({}), pser.replace({}))
 
+        self.assert_eq(kser.replace(np.nan, 45), pser.replace(np.nan, 45))
+
         msg = "'to_replace' should be one of str, list, dict, int, float"
         with self.assertRaisesRegex(ValueError, msg):
             kser.replace(ks.range(5))