如何与 pyspark 中的 ArrayType 列的每个元素进行交互？

Question

如果我在 pyspark 中有一个 ArrayType 列

from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame(((1,[]),(2,[1,2,3]),(3,[-2])),schema=StructType([StructField("a",IntegerType()),StructField("b",ArrayType(IntegerType()))]))
df.show()
output:
+---+---------+
|  a|        b|
+---+---------+
|  1|       []|
|  2|[1, 2, 3]|
|  3|     [-2]|

现在，我希望能够与列b的每个元素进行交互，例如，

1. 将每个元素除以 5 output：

+---+---------------+
|  a|              b|
+---+---------------+
|  1|             []|
|  2|[0.2, 0.4, 0.6]|
|  3|         [-0.4]|
+---+---------------+

2. 添加到每个元素等。

我如何 go 关于这样的转换，其中某些运算符或 function 应用于数组类型列的每个元素？

Answer 1

您正在寻找tranform function。 Transform 能够对数组的每个元素应用计算。

from pyspark.sql import functions as F

# Spark < 3.1.0
df.withColumn("b", F.expr("transform(b, x ->  x / 5)")).show()

"""
+---+---------------+
|  a|              b|
+---+---------------+
|  1|             []|
|  2|[0.2, 0.4, 0.6]|
|  3|         [-0.4]|
+---+---------------+
"""

# Spark >= 3.1.0

df.withColumn("b", F.transform("b", lambda x: x / 5)).show()
"""
+---+---------------+
|  a|              b|
+---+---------------+
|  1|             []|
|  2|[0.2, 0.4, 0.6]|
|  3|         [-0.4]|
+---+---------------+
"""