У меня есть фрейм данных Pyspark с примерно 4 миллиардами строк, поэтому эффективность операций очень важна. То, что я хочу сделать, кажется очень простым. Я хочу вычислить среднее значение из двух столбцов, и если один из них равен Null, я хочу вернуть только ненулевое значение. В Python я мог бы легко выполнить sh это с помощью np.nanmean, но я не верю, что что-то подобное реализовано в Pyspark.
Чтобы прояснить поведение, которое я ожидаю, см. Следующие строки примеров:
user_id col_1 col_2 avg_score
1 32 12 22
2 24 None 24
Ниже представлена моя текущая реализация. Обратите внимание, что все значения в col_1 гарантированно не равны нулю. Я считаю, что это, вероятно, может быть дополнительно оптимизировано:
from pyspark.sql import functions as f_
spark_df = spark_df.na.fill(0, 'col_2')
spark_df = spark_df.withColumn('avg_score',
sum([spark_df[i] for i in ['col_1','col_2']) /
sum([f_.when(spark_df[i] > 0, 1).otherwise(0) for i in ['col_1','col_2']]))
Если у кого-то есть предложения относительно того, есть ли более эффективный способ вычисления этого, я был бы очень признателен.