更改dataframe pyspark中的列值

时间:2016-09-13 10:41:02

标签: apache-spark machine-learning pyspark

我在Spark中有2个数据帧,它们是训练和测试。我在两者中都有一个分类列,比如Product_ID,我想要做的是,我想为这些类别设置-1值,这些类别在测试中但在火车中不存在。 因此,我首先在p_not_in_test中找到该列的不同类别。但我无法继续前进。该怎么做.....

p_not_in_test = test.select('Product_ID').subtract(train.select('Product_ID'))

p_not_in_test  = p_not_in_test.distinct()

此致

1 个答案:

答案 0 :(得分:1)

这是一个可重复的例子,首先我们创建虚拟数据:

test = sc.parallelize([("ID1", 1,5),("ID2", 2,4),
                       ("ID3", 5,8),("ID4", 9,0),
                       ("ID5", 0,3)]).toDF(["PRODUCT_ID", "val1", "val2"])

train = sc.parallelize([("ID1", 4,7),("ID3", 1,4),
                        ("ID5", 9,2)]).toDF(["PRODUCT_ID", "val1", "val2"])

现在我们需要扩展p_not_in_test的定义,以便我们得到一个列表作为输出:

p_not_in_test = (test.select('PRODUCT_ID')
                 .subtract(train.select('PRODUCT_ID'))
                 .rdd.map(lambda x: x[0]).collect())

最后,我们可以创建udf,在"-1"中不存在的每个ID前面添加train

from pyspark.sql.types import StringType
from pyspark.sql.functions import udf

addString = udf(lambda x: '-1 ' + x if x in p_not_in_test else x, StringType())

test.withColumn("NEW_ID",addString(test["PRODUCT_ID"])).show()
+----------+----+----+------+
|PRODUCT_ID|val1|val2|NEW_ID|
+----------+----+----+------+
|       ID1|   1|   5|   ID1|
|       ID2|   2|   4|-1 ID2|
|       ID3|   5|   8|   ID3|
|       ID4|   9|   0|-1 ID4|
|       ID5|   0|   3|   ID5|
+----------+----+----+------+