Spark DataFrame: Computing row-wise mean (or any aggregate operation)

by

All you need here is a standard SQL like this:

SELECT (US + UK + CAN) / 3 AS mean FROM df

which can be used directly with SqlContext.sql or expressed using DSL

df.select(((col("UK") + col("US") + col("CAN")) / lit(3)).alias("mean"))

If you have a larger number of columns you can generate expression as follows:

from functools import reduce
from operator import add
from pyspark.sql.functions import col, lit

n = lit(len(df.columns) - 1.0)
rowMean  = (reduce(add, (col(x) for x in df.columns[1:])) / n).alias("mean")

df.select(rowMean)

or

rowMean  = (sum(col(x) for x in df.columns[1:]) / n).alias("mean")
df.select(rowMean)

Finally its equivalent in Scala:

df.select(df.columns
  .drop(1)
  .map(col)
  .reduce(_ + _)
  .divide(df.columns.size - 1)
  .alias("mean"))

In a more complex scenario you can combine columns using array function and use an UDF to compute statistics:

import numpy as np
from pyspark.sql.functions import array, udf
from pyspark.sql.types import FloatType

combined = array(*(col(x) for x in df.columns[1:]))
median_udf = udf(lambda xs: float(np.median(xs)), FloatType())

df.select(median_udf(combined).alias("median"))

The same operation expressed using Scala API:

val combined = array(df.columns.drop(1).map(col).map(_.cast(DoubleType)): _*)
val median_udf = udf((xs: Seq[Double]) => 
    breeze.stats.DescriptiveStats.percentile(xs, 0.5))

df.select(median_udf(combined).alias("median"))

Since Spark 2.4 an alternative approach is to combine values into an array and apply aggregate expression. See for example Spark Scala row-wise average by handling null.

More Related Contents:

  1. Spark Equivalent of IF Then ELSE
  2. How to add a constant column in a Spark DataFrame?
  3. Spark Dataframe distinguish columns with duplicated name
  4. How do I add a new column to a Spark DataFrame (using PySpark)?
  5. How to perform union on two DataFrames with different amounts of columns in spark?
  6. Does spark predicate pushdown work with JDBC?
  7. How to change dataframe column names in pyspark?
  8. ‘PipelinedRDD’ object has no attribute ‘toDF’ in PySpark
  9. Explode in PySpark
  10. How can we JOIN two Spark SQL dataframes using a SQL-esque “LIKE” criterion?
  11. Dividing complex rows of dataframe to simple rows in Pyspark
  12. Pyspark ‘NoneType’ object has no attribute ‘_jvm’ error
  13. Why is Apache-Spark – Python so slow locally as compared to pandas?
  14. How to explode multiple columns of a dataframe in pyspark
  15. Spark add new column to dataframe with value from previous row
  16. PySpark: multiple conditions in when clause
  17. PySpark create new column with mapping from a dict
  18. Pyspark: Replacing value in a column by searching a dictionary
  19. Create Spark DataFrame. Can not infer schema for type
  20. Pivot String column on Pyspark Dataframe
  21. Using UDF ignores condition in when
  22. Create single row dataframe from list of list PySpark
  23. Total size of serialized results of 16 tasks (1048.5 MB) is bigger than spark.driver.maxResultSize (1024.0 MB)
  24. Concatenate two PySpark dataframes
  25. How do I convert an array (i.e. list) column to Vector
  26. Filtering DataFrame using the length of a column
  27. Add column sum as new column in PySpark dataframe
  28. How to drop all columns with null values in a PySpark DataFrame?
  29. PySpark: compute row maximum of the subset of columns and add to an exisiting dataframe
  30. how to add Row id in pySpark dataframes [duplicate]

Leave a Comment