суммирующий вложенный список pySpark

Question

У меня есть датафрейм, похожий на этот

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 в тип списка.

Answer 1

Полагаю, если вам нужен только вектор, который возвращает сумму строк, это можно сделать немного проще.

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]|
+---------+