[SPARK-14932][SQL] Allow DataFrame.replace() to replace values with None

python/pyspark/sql/dataframe.py

@@ -1277,10 +1277,10 @@ def replace(self, to_replace, value, subset=None):
             If the value is a dict, then `value` is ignored and `to_replace` must be a
             mapping from column name (string) to replacement value. The value to be
             replaced must be an int, long, float, or string.
-        :param value: int, long, float, string, or list.
+        :param value: int, long, float, string, list or None.
             Value to use to replace holes.
-            The replacement value must be an int, long, float, or string. If `value` is a
-            list or tuple, `value` should be of the same length with `to_replace`.
+            The replacement value must be an int, long, float, string or None. If `value`
+            is a list or tuple, `value` should be of the same length with `to_replace`.
         :param subset: optional list of column names to consider.
             Columns specified in subset that do not have matching data type are ignored.
             For example, if `value` is a string, and subset contains a non-string column,
@@ -1296,6 +1296,16 @@ def replace(self, to_replace, value, subset=None):
         |null|  null| null|
         +----+------+-----+
 
+        >>> df4.na.replace('Alice', None).show()
+        +----+------+----+
+        | age|height|name|
+        +----+------+----+
+        |  10|    80|null|
+        |   5|  null| Bob|
+        |null|  null| Tom|
+        |null|  null|null|
+        +----+------+----+
+
         >>> df4.na.replace(['Alice', 'Bob'], ['A', 'B'], 'name').show()
         +----+------+----+
         | age|height|name|
@@ -1310,8 +1320,8 @@ def replace(self, to_replace, value, subset=None):
             raise ValueError(
                 "to_replace should be a float, int, long, string, list, tuple, or dict")
 
-        if not isinstance(value, (float, int, long, basestring, list, tuple)):
-            raise ValueError("value should be a float, int, long, string, list, or tuple")
+        if value is not None and not isinstance(value, (float, int, long, basestring, list, tuple)):
+            raise ValueError("value should be a float, int, long, string, list, tuple or None")
 
         rep_dict = dict()
 
@@ -1328,7 +1338,7 @@ def replace(self, to_replace, value, subset=None):
             if len(to_replace) != len(value):
                 raise ValueError("to_replace and value lists should be of the same length")
             rep_dict = dict(zip(to_replace, value))
-        elif isinstance(to_replace, list) and isinstance(value, (float, int, long, basestring)):
+        elif isinstance(to_replace, list) and (value is None or isinstance(value, (float, int, long, basestring))):
             rep_dict = dict([(tr, value) for tr in to_replace])
         elif isinstance(to_replace, dict):
             rep_dict = to_replace

sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala

@@ -342,11 +342,14 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) {
     }
 
     // replacementMap is either Map[String, String] or Map[Double, Double] or Map[Boolean,Boolean]
-    val replacementMap: Map[_, _] = replacement.head._2 match {
-      case v: String => replacement
-      case v: Boolean => replacement
-      case _ => replacement.map { case (k, v) => (convertToDouble(k), convertToDouble(v)) }
-    }
+    val replacementMap: Map[_, _] =
+      if (replacement.head._2 == null)
+        replacement
+      else replacement.head._2 match {
+        case v: String => replacement
+        case v: Boolean => replacement
+        case _ => replacement.map { case (k, v) => (convertToDouble(k), convertToDouble(v)) }
+      }
 
     // targetColumnType is either DoubleType or StringType or BooleanType
     val targetColumnType = replacement.head._1 match {

sql/core/src/test/scala/org/apache/spark/sql/DataFrameNaFunctionsSuite.scala

@@ -208,16 +208,16 @@ class DataFrameNaFunctionsSuite extends QueryTest with SharedSQLContext {
     assert(out(4) === Row("Amy", null, null))
     assert(out(5) === Row(null, null, null))
 
-    // Replace only the age column
-    val out1 = input.na.replace("age", Map(
-      16 -> 61,
-      60 -> 6,
-      164.3 -> 461.3  // Alice is really tall
+    // Replace only the name column
+    val out1 = input.na.replace("name", Map(
+      "Bob" -> "Bravo",
+      "Alice" -> "Jessie",
+      "David" -> null
     )).collect()
 
-    assert(out1(0) === Row("Bob", 61, 176.5))
-    assert(out1(1) === Row("Alice", null, 164.3))
-    assert(out1(2) === Row("David", 6, null))
+    assert(out1(0) === Row("Bravo", 16, 176.5))
+    assert(out1(1) === Row("Jessie", null, 164.3))
+    assert(out1(2) === Row(null, 60, null))
     assert(out1(3).get(2).asInstanceOf[Double].isNaN)
     assert(out1(4) === Row("Amy", null, null))
     assert(out1(5) === Row(null, null, null))