совокупность вектора в искре

Question

Я новичок в спарк и мне нужна поддержка в решении проблемы ниже. У меня есть данные, как показано ниже:

Country value
India   [1,2,3,4,5]
US  [8,9,10,11,12]
US  [7,6,5,4,3]
India   [8,7,6,5,4]

, и требуемый вывод - это совокупность элементов вектора той же страны, что и ниже в искре.

Output:
Country value
India   [9,9,9,9,9]
US  [15,15,15,15,15]    

Answer 1

AFAIK, spark не обеспечивает функции агрегации для массивов. Поэтому, если размер массивов фиксирован, вы можете создать один столбец для каждого элемента массива, агрегировать, а затем заново создать массив.

В общем случае это может выглядеть следующим образом:

from pyspark.sql.functions import col, sum

# first, let's get the size of the array
size = len(df.first()['value'])

# Then, summing each element separately:
aggregation = df.groupBy("country")\
    .agg(*[sum(df.value.getItem(i)).alias("v"+str(i)) for i in range(size)])
aggregation.show()
+-------+---+---+---+---+---+                                                   
|country| v0| v1| v2| v3| v4|
+-------+---+---+---+---+---+
|  India|  9|  9|  9|  9|  9|
|     US| 15| 15| 15| 15| 15|
+-------+---+---+---+---+---+


# Finally, we recreate the array
result = aggregation.select(df.country,\
    functions.array(*[col("v"+str(i)) for i in range(size)]).alias("value"))
result.show()
+-------+--------------------+
|country|               value|
+-------+--------------------+
|  India|     [9, 9, 9, 9, 9]|
|     US|[15, 15, 15, 15, 15]|
+-------+--------------------+