How to find total and average of columns in PySpark Azure Databricks?

Are you looking to find out how to find the total of PySpark DataFrame column in Azure Databricks cloud using sum() function or maybe you are looking for a solution, to find the average of column in PySpark Databricks using the avg() function? If you are looking for any of these problem solutions, you have landed on the correct page. I will also help you how to use PySpark sum() & avg() function with multiple examples in Azure Databricks. I will explain it by taking a practical example. So please don’t waste time let’s start with a step-by-step guide to understand how to use the sum() & avg() function in PySpark.

In this blog, I will teach you the following with practical examples:

  • Syntax of sum() & avg() functions
  • Aggregation on a selected column
  • Aggregation on a grouped column

The PySpark sum() function is used to find the total value of numeric column in PySpark DataFrame.

Syntax:

sum()

What is the syntax of the sum() function in PySpark Azure Databricks?

The syntax is as follows:

sum(column)
Parameter NameRequiredDescription
column (str, Column)YesIt represents the column to be considered aggregation.
Table 1: sum() Method in PySpark Databricks Parameter list with Details

Apache Spark Official Documentation Link: sum()

The PySpark avg() function is used to find the average value of numeric column in PySpark DataFrame.

Syntax:

avg()

What is the syntax of the avg() function in PySpark Azure Databricks?

The syntax is as follows:

avg(column)
Parameter NameRequiredDescription
column (str, Column)YesIt represents the column to be considered aggregation.
Table 1: avg() Method in PySpark Databricks Parameter list with Details

Apache Spark Official Documentation Link: avg()

Create a simple DataFrame

Let’s understand the use of the sum() & avg() functions with various examples. Let’s start by creating a DataFrame.

Gentle reminder:

In Databricks,

  • sparkSession made available as spark
  • sparkContext made available as sc

In case, you want to create it manually, use the below code.

from pyspark.sql.session import SparkSession

spark = SparkSession.builder 
    .master("local[*]") 
    .appName("azurelib.com") 
    .getOrCreate()

sc = spark.sparkContext

a) Create manual PySpark DataFrame

data = [
    ("chevrolet vega 2300","USA",None,90,28.0,"1970-01-01"),
    ("chevrolet vega 2300","USA",15.5,90,28.0,"1970-01-01"),
    ("toyota corona","Japan",14.0,95,25.0,"1970-01-01"),
    ("ford pinto","USA",19.0,75,25.0,"1971-01-01"),
    ("amc gremlin","USA",13.0,100,19.0,"1971-01-01"),
    ("plymouth satellite custom","USA",15.5,105,16.0,"1971-01-01"),
    ("datsun 510 (sw)","Japan",17.0,92,28.0,"1972-01-01"),
    ("toyouta corona mark ii (sw)","Japan",14.5,97,23.0,"1972-01-01"),
    ("dodge colt (sw)","USA",15.0,80,28.0,"1972-01-01"),
    ("toyota corolla 1600 (sw)","Japan",16.5,88,27.0,"1972-01-01")
]

columns = ["name","origin","acceleration","horse_power","miles_per_gallon","year"]
df = spark.createDataFrame(data, schema=columns)
df.printSchema()
df.show(5, truncate=False)

"""
root
 |-- name: string (nullable = true)
 |-- origin: string (nullable = true)
 |-- acceleration: double (nullable = true)
 |-- horse_power: long (nullable = true)
 |-- miles_per_gallon: double (nullable = true)
 |-- year: string (nullable = true)

+-------------------+------+------------+-----------+----------------+----------+
|name               |origin|acceleration|horse_power|miles_per_gallon|year      |
+-------------------+------+------------+-----------+----------------+----------+
|chevrolet vega 2300|USA   |null        |90         |28.0            |1970-01-01|
|chevrolet vega 2300|USA   |15.5        |90         |28.0            |1970-01-01|
|toyota corona      |Japan |14.0        |95         |25.0            |1970-01-01|
|ford pinto         |USA   |19.0        |75         |25.0            |1971-01-01|
|amc gremlin        |USA   |13.0        |100        |19.0            |1971-01-01|
+-------------------+------+------------+-----------+----------------+----------+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

# replace the file_path with the source file location which you have downloaded.

df_2 = spark.read.format("csv").option("inferSchema", True).option("header", True).load(file_path)
df_2.printSchema()

"""
root
 |-- name: string (nullable = true)
 |-- origin: string (nullable = true)
 |-- acceleration: double (nullable = true)
 |-- horse_power: long (nullable = true)
 |-- miles_per_gallon: double (nullable = true)
 |-- year: string (nullable = true)
"""

Note: Here, I will be using the manually created DataFrame.

How to use sum() & avg() on selected columns of PySpark DataFrame using Azure Databricks?

Let’s see how to use sum() & avg() on selected columns of PySpark DataFrame in Azure Databricks using various methods.

Example:

from pyspark.sql.functions import sum, avg

df.select(
    sum("acceleration").alias("total_acc"),
    avg("miles_per_gallon").alias("avg_mileage")
).show()

"""
Output:

+---------+-----------+
|total_acc|avg_mileage|
+---------+-----------+
|    140.0|       24.7|
+---------+-----------+

"""

How to use sum() & avg() on grouped columns of PySpark DataFrame using Azure Databricks?

Let’s see how to use sum() & avg() on grouped columns of PySpark DataFrame in Azure Databricks using various methods.

Example:

from pyspark.sql.functions import sum, avg

df.groupBy("origin") \
.agg(
    sum("acceleration").alias("total_acc"),
    avg("miles_per_gallon").alias("avg_mileage")
).show()

"""
Output:

+------+---------+-----------+
|origin|total_acc|avg_mileage|
+------+---------+-----------+
|   USA|     78.0|       24.0|
| Japan|     62.0|      25.75|
+------+---------+-----------+

"""

Note: By default, the sum() and avg() functions omit the null/None value of a column.

I have attached the complete code used in this blog in a notebook format to this GitHub link. You can download and import this notebook in databricks, jupyter notebook, etc.

When should you use the PySpark sum() & avg() in Azure Databricks?

These could be the possible reasons:

  1. To find the total summation of numeric columns
  2. To find the average value of numeric columns

Real World Use Case Scenarios for PySpark DataFrame sum() & avg() in Azure Databricks?

Assume you were given a student dataset, such as student ID, name, subject, and mark. It would help if you found the total and average marks scored by the student. The PySpark sum() and avg() aggregation function helps in summing and averaging the numeric column and returns a single record of it.

What are the alternatives to the sum() and avg() functions in PySpark Azure Databricks?

There are multiple alternatives to the sum() and avg() functions, which are as follows:

  • mean(): also used for finding out the average value of numeric columns
  • sum_distinct(): used for summing the unique records

Final Thoughts

In this article, we have learned about the PySpark sum() & avg() method of DataFrame in Azure Databricks along with the examples explained clearly. I have also covered different scenarios with practical examples that could be possible. I hope the information that was provided helped in gaining knowledge.

Please share your comments and suggestions in the comment section below and I will try to answer all your queries as time permits.

Arud Seka Berne S

As a big data engineer, I design and build scalable data processing systems and integrate them with various data sources and databases. I have a strong background in Python and am proficient in big data technologies such as Hadoop, Hive, Spark, Databricks, and Azure. My interest lies in working with large datasets and deriving actionable insights to support informed business decisions.