解析pyspark中数组的每个元素并应用子字符串

Posted 2023-04-13

【中文标题】解析pyspark中数组的每个元素并应用子字符串

【英文标题】：Parse through each element of an array in pyspark and apply substring

【发布时间】：2020-09-05 01:03:45

【问题描述】：

您好，我有一个 pyspark 数据框，其数组 col 如下所示。

我想遍历每个元素并只获取连字符之前的字符串并创建另一列。

+------------------------------+
|array_col                     |
+------------------------------+
|[hello-123, abc-111]          |
|[hello-234, def-22, xyz-33]   |
|[hiiii-111, def2-333, lmn-222]|
+------------------------------+

期望的输出；

+------------------------------+--------------------+
|col1                          |new_column          |
+------------------------------+--------------------+
|[hello-123, abc-111]          |[hello, abc]        |
|[hello-234, def-22, xyz-33]   |[hello, def, xyz]   |
|[hiiii-111, def2-333, lmn-222]|[hiiii, def2, lmn]  |
+------------------------------+--------------------+

我正在尝试类似下面的内容，但我无法在 udf 中应用正则表达式/子字符串。

cust_udf = udf(lambda arr: [x for x in arr],ArrayType(StringType()))
df1.withColumn('new_column', cust_udf(col("col1")))

任何人都可以帮助解决这个问题。谢谢

【参考方案1】：

从 Spark-2.4 使用 transform 高阶函数。

Example:

df.show(10,False)
#+---------------------------+
#|array_col                  |
#+---------------------------+
#|[hello-123, abc-111]       |
#|[hello-234, def-22, xyz-33]|
#+---------------------------+

df.printSchema()
#root
# |-- array_col: array (nullable = true)
# |    |-- element: string (containsNull = true)

from pyspark.sql.functions import *


df.withColumn("new_column",expr('transform(array_col,x -> split(x,"-")[0])')).\
show()
#+--------------------+-----------------+
#|           array_col|       new_column|
#+--------------------+-----------------+
#|[hello-123, abc-111]|     [hello, abc]|
#|[hello-234, def-2...|[hello, def, xyz]|
#+--------------------+-----------------+