将逗号分隔的字符串转换为 PySpark 数据帧中的数组

有时候数据以逗号分隔的形式存在字符串中，我们需要将这些字符串转换为 PySpark 数据帧中的数组以进行后续处理。下面我们介绍一种实现方法。

首先，我们定义一个包含逗号分隔字符串的 PySpark 数据帧：

from pyspark.sql.functions import split
from pyspark.sql.types import ArrayType, StringType
from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("split_string_to_array").getOrCreate()

data = [("1,2,3",), ("4,5,6",), ("7,8,9",)]
df = spark.createDataFrame(data, ["str"])
df.show()

输出结果为：

+-----+
|  str|
+-----+
|1,2,3|
|4,5,6|
|7,8,9|
+-----+

我们使用 PySpark 内置的 split 函数，将逗号分隔的字符串转换为数组。此时，数组元素的数据类型为字符串类型。我们可以使用 UDF 函数将字符串类型的数组转为整型数组，也可以使用内置函数 cast 将字符串类型的数组转为整型数组。

下面是使用 UDF 函数的示例代码：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType

# 定义 UDF 函数，将字符串类型的数组转为整型数组
@udf(returnType=ArrayType(IntegerType()))
def str_to_int_array(str_value):
    str_array = str_value.split(",")
    int_array = [int(x) for x in str_array]
    return int_array

# 使用 split 函数将逗号分隔的字符串转换为数组，并应用 UDF 函数将整型数组转换为 PySpark 中的整型数组
df = df.withColumn("int_array", str_to_int_array("str"))
df.show()

输出结果为：

+-----+---------+
|  str|int_array|
+-----+---------+
|1,2,3|[1, 2, 3]|
|4,5,6|[4, 5, 6]|
|7,8,9|[7, 8, 9]|
+-----+---------+

下面是使用 cast 函数的示例代码：

from pyspark.sql.functions import split

# 使用 split 函数将逗号分隔的字符串转换为数组，并将每个元素转为整型
df = df.withColumn("int_array", split("str", ",").cast(ArrayType(IntegerType())))
df.show()

输出结果与上述方法相同。

以上就是将逗号分隔的字符串转换为 PySpark 数据帧中的数组的实现方法。实际使用时，我们可以根据需求选择上述两种方法中的任意一种。