PySpark: вычисление агрегации по строкам со значениями None

Question

Как вы вычисляете агрегации по строкам, когда есть значения None?

Например, следующий код дает мне результаты «Нет», когда я пытаюсь суммировать:

columns = ['US', 'UK', 'Can']
countries.withColumn('sum', sum(countries[col] for col in columns))

+--+--+-----+---+
|id|US| UK  |Can|
+--+--+-----+---+
| 1|50|None |  0|
| 1| 0|100  |  0|
| 1| 0|  0  |125|
| 2|75|  0  |  0|
+--+--+-----+---+

Answer 1

Вот рабочий ответ, который я нашел, но был бы заинтересован чем-нибудь еще более элегантным. Спасибо.

columns = ['US', 'UK', 'Can']
countries.withColumn('sum', sum([func.coalesce(col(x), func.lit(0)) for x in columns]))

Answer 2

Можно преобразовать в pandas DataFrame, вычислить сумму, а затем обратно в Spark DataFrame (но я не рекомендую это). Обратите внимание на следующие работы, потому что pandas.Dataframe.sum по умолчанию skipna=True

# Generate data
data = [(1, 50, None, 0), (1, 0, 100, 0), (1, 0, 0, 125), (2, 75, 0, 0)]
countries = spark.createDataFrame(data, ["id", "US", "UK", "Can"])
columns = ['US', 'UK', 'Can']
# Solution
pd_countries = countries.toPandas()
pd_countries['sum'] = pd_countries[columns].sum(axis=1) # default is ignore nan
result = spark.createDataFrame(pd_countries)
result.show()
+---+---+-----+---+-----+                                                       
| id| US|   UK|Can|  sum|
+---+---+-----+---+-----+
|  1| 50|  NaN|  0| 50.0|
|  1|  0|100.0|  0|100.0|
|  1|  0|  0.0|125|125.0|
|  2| 75|  0.0|  0| 75.0|
+---+---+-----+---+-----+

Также обратите внимание, что ваш оригинальный код выполняет следующую операцию

sum(countries[c] for c in columns)
# Column<b'(((US + 0) + UK) + Can)'>
# equivalent to (col(US) + 0 + col(UK) + col(Can))

На мой взгляд, самый простой способ должен использовать na.fill(0) как предложено комментатором и ответом, а затем вычислить сумму.

Answer 3

Создайте копию своего фрейма данных, используя copy.deepcopy в pyspark, а затем замените свои нули на ноль и выполните суммирование и усреднение.

import copy

countries = spark.createDataFrame([[1,50,None,0], [1,0,100,0],[1,0,0,125],[2,75,0,0]], ['id', 'US','UK','Can'])
_schema = copy.deepcopy(countries.schema)
_countries = countries.rdd.zipWithIndex().toDF(_schema)

_countries=_countries.na.fill(0)

columns = ['US', 'UK', 'Can']
_countries.withColumn('sum', sum(_countries[col] for col in columns))

Это должно дать желаемое значение без необходимости изменения исходного фрейма данных.