pyspark計數給定列的每一行中的下划線數量

Question

我正在使用pyspark版本1.5.2。 我有一個帶有列“ id”的pyspark 數據框，如下所示：

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用於其他操作。

任何幫助表示贊賞！

Answer 1

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

Answer 2

這是一種非udf方法：

您可以使用與該答案相同的方法來計算每個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|
#+----------+------------+