当密钥出现在地图中时,有条件地替换数据框中的值

时间:2017-12-15 02:14:37

标签: scala dictionary apache-spark

我在scala spark环境中有两个对象。第一个是具有许多String类型的列的数据帧。第二个是带有column_names作为键的Map,每个键的值是另一个定义替换值的Map。

我想迭代地图Map,将column_name键与dataframe键匹配,然后使用替换值Map替换值。如果替换值的Map中没有该值,那么我想返回列的原始值。

实现这一目标的最佳方法是什么:

val replacements = Map("Col_1" -> Map("Male" -> "0", "Female" -> "1"),
                       "Col_2" -> Map("2" -> "1", "3" -> "1", "4" -> "1")

val df = Seq(("Mary", "Female", "0"), ("John", "Male", "2"), ("Luke", "Male", "1")).toDF(List("name", "Col_1", "Col_2").toSeq:_*)

df.show
+----+------+-----+
|name| Col_1|Col_2|
+----+------+-----+
|Mary|Female|    0|
|John|  Male|    2|
|Luke|  Male|    1|
+----+------+-----+

//I hoped to do something like this:

val fixed_df = replacements.foldLeft(df) { 
case (tmpDF, (colName, foo: Map[key, v])) => tmpDF.withColumn(colName, when($"$colName" isin(foo.keys.to[List]:_*), foo($"$colName")).otherwise($"$colName"))
    }

//to get output like this:
fixed_df.take(3)
+----+------+-----+
|name| Col_1|Col_2|
+----+------+-----+
|Mary|     1|    0|
|John|     0|    1|
|Luke|     0|    1|
+----+------+-----+

请注意,感谢@tzachz寻求类似的不同解决方案的帮助。

1 个答案:

答案 0 :(得分:1)

您需要的只是na.replace

replacements.keys.foldLeft(df)(
  (acc, key) => acc.na.replace(key, replacements(key))
).show

// +----+-----+-----+
// |name|Col_1|Col_2|
// +----+-----+-----+
// |Mary|    1|    0|
// |John|    0|    1|
// |Luke|    0|    1|
// +----+-----+-----+