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

pyspark count number of underscores in each row of a given column

我使用的是 pyspark 版本 1.5.2。我有一个带有列 "id" 的 pyspark dataframe,如下所示:

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

我想统计DF的每一行中'_'(下划线)的个数,并执行一个when操作,如果字符串中只有1个下划线,我想添加'_1'作为后缀, 否则保留原值。所以期望的结果是:

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

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

感谢任何帮助!

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))

这是一个非 udf 方法:

您可以使用 to count the number of _ in each id, and use pyspark.sql.functions.when() to check if the count is equal to 1. If yes, use 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|
#+----------+------------+