spark-计算2列或更多列中的平均值，并在每行中放入新列

Question

假设我有一个包含以下内容的数据集/数据框：-

name, marks1, marks2
Alice, 10, 20
Bob, 20, 30

我想添加一个新列，该列应具有B列和C列的平均值。

预期结果：-

name, marks1, marks2, Result(Avg)
Alice, 10, 20, 15
Bob, 20, 30, 25

对于求和或任何其他算术运算，我使用df.withColumn("xyz", $"marks1"+$"marks2")。我找不到平均值的类似方法。请帮忙。

此外：-列数不是固定的。就像有时它可能是2列的平均值，有时是3列甚至更多列。所以我想要一个通用的代码，它应该可以工作。

Answer 1

一种最简单且优化的方法是创建标记列的列列表，并将其与withColumn一起用作

pyspark

from pyspark.sql.functions import col

marksColumns = [col('marks1'), col('marks2')]

averageFunc = sum(x for x in marksColumns)/len(marksColumns)

df.withColumn('Result(Avg)', averageFunc).show(truncate=False)

您应该获得

+-----+------+------+-----------+
|name |marks1|marks2|Result(Avg)|
+-----+------+------+-----------+
|Alice|10    |20    |15.0       |
|Bob  |20    |30    |25.0       |
+-----+------+------+-----------+

scala-spark

在 scala 中的过程几乎与上述 python 中的过程相同

import org.apache.spark.sql.functions.{col, lit}

val marksColumns = Array(col("marks1"), col("marks2"))

val averageFunc = marksColumns.foldLeft(lit(0)){(x, y) => x+y}/marksColumns.length

df.withColumn("Result(Avg)", averageFunc).show(false)

应该给您与 pyspark

中相同的输出

我希望答案会有所帮助

Answer 2

就像使用用户定义函数一样简单。通过创建一个特定的UDF以处理平均许多列，您将能够根据需要多次重复使用它。

Python

在此代码段中，我正在创建一个UDF，该UDF将使用一组列，并计算其平均值。

from pyspark.sql.functions import udf, array
from pyspark.sql.types import DoubleType

avg_cols = udf(lambda array: sum(array)/len(array), DoubleType())

df.withColumn("average", avg_cols(array("marks1", "marks2"))).show()

输出：

+-----+------+------+--------+
| name|marks1|marks2| average|
+-----+------+------+--------+
|Alice|    10|    20|    15.0|
|  Bob|    20|    30|    25.0|
+-----+------+------+--------+

Scala

使用Scala API，您必须将选定的列作为行处理。您只需要使用Spark struct函数选择列即可。

import org.apache.spark.sql.functions._
import spark.implicits._
import scala.util.Try

def average = udf((row: Row) => {
  val values = row.toSeq.map(x => Try(x.toString.toDouble).toOption).filter(_.isDefined).map(_.get)
  if(values.nonEmpty) values.sum / values.length else 0.0
})

df.withColumn("average", average(struct($"marks1", $"marks2"))).show()

如您所见，我使用Try将所有值都强制转换为Double，这样，如果无法强制转换该值，则不会抛出任何异常，仅对那些定义。

仅此而已：）