根据另一列更新pyspark数据框列

时间:2018-05-16 23:56:44

标签: apache-spark pyspark apache-spark-sql

以下是pyspark中的数据框。我想根据val列中的值更新data frame中的列tests

df.show()
+---------+----+---+
|    tests| val|asd|
+---------+----+---+
|    test1|   Y|  1|
|    test2|   N|  2|
|    test2|   Y|  1|
|    test1|   N|  2|
|    test1|   N|  3|
|    test3|   N|  4|
|    test4|   Y|  5|
+---------+----+---+

我希望在任何给定的test val Y时更新该值,然后该特定测试的所有val's都应更新为Y。如果不是那么他们有什么价值。

基本上我希望data frame如下所示。

result_df.show()

+---------+----+---+
|    tests| val|asd|
+---------+----+---+
|    test1|   Y|  1|
|    test2|   Y|  2|
|    test2|   Y|  1|
|    test1|   Y|  2|
|    test1|   Y|  3|
|    test3|   N|  4|
|    test4|   Y|  5|
+---------+----+---+

我应该怎么做才能实现这一目标。

2 个答案:

答案 0 :(得分:3)

这是一个解决方案。 首先,我们发现每个测试是否有val Y.

import pyspark.sql.functions as sf
by_test = df.groupBy('tests').agg(sf.sum((sf.col('val') == 'Y').cast('int')).alias('HasY'))
by_test.show()
+-----+----+
|tests|HasY|
+-----+----+
|test4|   1|
|test3|   0|
|test1|   1|
|test2|   1|
+-----+----+

加入回原始数据框

df = df.join(by_test, on='tests')
df.show()
+-----+---+---+----+
|tests|val|asd|HasY|
+-----+---+---+----+
|test4|  Y|  5|   1|
|test3|  N|  4|   0|
|test1|  Y|  1|   1|
|test1|  N|  2|   1|
|test1|  N|  3|   1|
|test2|  N|  2|   1|
|test2|  Y|  1|   1|
+-----+---+---+----+

使用when / otherwise

创建一个具有相同名称的新列
df = df.withColumn('val', sf.when(sf.col('HasY') > 0, 'Y').otherwise(sf.col('val')))
df = df.drop('HasY')
df.show()
+-----+---+---+
|tests|val|asd|
+-----+---+---+
|test4|  Y|  5|
|test3|  N|  4|
|test1|  Y|  1|
|test1|  Y|  2|
|test1|  Y|  3|
|test2|  Y|  2|
|test2|  Y|  1|
+-----+---+---+

答案 1 :(得分:3)

使用max窗口函数和selectExpr

df.selectExpr(
    'tests', 'max(val) over (partition by tests) as val', 'asd'
).show()

+-----+---+---+
|tests|val|asd|
+-----+---+---+
|test4|  Y|  5|
|test3|  N|  4|
|test1|  Y|  1|
|test1|  Y|  2|
|test1|  Y|  3|
|test2|  Y|  2|
|test2|  Y|  1|
+-----+---+---+