pyspark计数给定列的每一行中的下划线数量

时间:2018-07-20 20:39:57

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

我正在使用pyspark版本1.5.2。我有一个带有“ id”列的pyspark 数据框,如下所示:

initializer_list

我想计算DF每行中的'_'(下划线)的数目,并执行when操作,以便如果字符串中只有1个下划线,我想将'_1'添加为后缀,否则保持原样。因此所需的结果将是:

id
------------
000001_128
000123_1_3 
006745_8
000000_9_7

我正在使用pyspark.sql.functions进行其他操作。

感谢您的帮助!

2 个答案:

答案 0 :(得分:2)

from pyspark.sql.functions import udf

@udf(returnType='string')
def fmt(s):
    return s if s.count('_')!=1 else f'{s}_1'


df.withColumn('id', fmt(df.id))

答案 1 :(得分:2)

这是一种非udf方法:

您可以使用this answer中的相同方法来计算每个_id的数目,并使用pyspark.sql.functions.when()检查计数是否等于1。如果是,请使用pyspark.sql.functions.format_string()来制作new_id,否则将列保持不变:

import pyspark.sql.functions as f

df.withColumn(
    "new_id",
    f.when(
        (f.size(f.split("id", "_"))-1) == 1,
        f.format_string("%s_1",f.col("id"))
    ).otherwise(f.col("id"))
).show()
#+----------+------------+
#|        id|      new_id|
#+----------+------------+
#|000001_128|000001_128_1|
#|000123_1_3|  000123_1_3|
#|  006745_8|  006745_8_1|
#|000000_9_7|  000000_9_7|
#+----------+------------+