Python pySpark求和嵌套列表

我有一个像这样的数据框

import pandas as pd

df = pd.DataFrame({'hist': [[1, 2, 3], [0, 1, 0], [1, 0, 0]]})
df = spark.createDataFrame(df)

我想聚合数据帧并返回

hist

的元素和。除此之外的结果是

[2,3,3]

我试着这么做

from pyspark.sql.types import ArrayType, IntegerType
from pyspark.sql import functions as F

df.agg(F.udf(lambda x: list(np.sum(x, axis=0).astype('int')), ArrayType(IntegerType()))(F.collect_list(df.hist))).show(10, False)

但我犯了这个错误

expected zero arguments for construction of ClassDict (for numpy.dtype)

---

我不希望出现这种情况，因为我显式地将UDF的输出转换为列表类型。

我想如果您只需要一个返回行和的向量，那么可以更容易地完成

import pandas as pd
import pyspark.sql.functions as F

df = pd.DataFrame({'hist': [[1, 2, 3], [0, 1, 0], [1, 0, 0]]})
df = spark.createDataFrame(df)
import pyspark.sql.functions as F
sumOfRows = df.agg(F.array(*[F.sum(F.col("hist")[i]) for i in range(n+1)]).alias("sum"))
sumOfRows.show()

输出：

+---------+
|      sum|
+---------+
|[2, 3, 3]|
+---------+

---

我想如果你只需要一个向量，它返回行和，这可以做得简单一点

import pandas as pd
import pyspark.sql.functions as F

df = pd.DataFrame({'hist': [[1, 2, 3], [0, 1, 0], [1, 0, 0]]})
df = spark.createDataFrame(df)
import pyspark.sql.functions as F
sumOfRows = df.agg(F.array(*[F.sum(F.col("hist")[i]) for i in range(n+1)]).alias("sum"))
sumOfRows.show()

输出：

+---------+
|      sum|
+---------+
|[2, 3, 3]|
+---------+

---

非常感谢。这是在耍花招。我唯一的一点意见是，

n

需要硬编码，这不是很好，所以使用

F.size（'hist'）

可能是一个更好的选择谢谢！这是在耍花招。我唯一的一点意见是，

n

需要硬编码，这并不好，因此使用

F.size（'hist'）

可能是更好的选择