如何对PySpark特征向量中所有元素求平方？

输入DataFrame #

schema = ["id", "vector"]
data = [
(0,[1.0,1.0,1.0]),
(1,[2.0,2.0,1.0]),
]
df = spark.createDataFrame(data, schema)

方法一：使用内置transform函数（推荐，Spark 3.1+） #

利用PySpark内置的transform函数遍历数组元素，结合pow函数计算平方，性能优于自定义UDF：

from pyspark.sql.functions import transform, pow

# 对vector列的每个元素求平方
df_result = df.withColumn(
    "vector",
    transform("vector", lambda x: pow(x, 2))
)

# 查看结果
df_result.show(truncate=False)

方法二：自定义UDF（兼容旧版Spark） #

如果使用的Spark版本低于3.1，可以通过自定义UDF实现：

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

# 定义UDF：接收数组，返回每个元素平方后的新数组
square_vector_udf = udf(lambda vec: [x**2 for x in vec], ArrayType(DoubleType()))

df_result = df.withColumn("vector", square_vector_udf("vector"))

# 查看结果
df_result.show(truncate=False)

输出结果 #

+---+---------------+
|id |vector         |
+---+---------------+
|0  |[1.0, 1.0, 1.0]|
|1  |[4.0, 4.0, 1.0]|
+---+---------------+

内容的提问来源于stack exchange，提问作者NHUV