How to use the count() function in PySpark Azure Databricks?

Are you looking to find how to use the count() function of a PySpark Dataframe using Azure Databricks cloud or maybe you are looking for a solution, to count grouped records of a Dataframe in PySpark using Azure Databricks? If you are looking for any of these problem solutions, then you have landed on the correct page. I will also show you how to use both PySpark and Spark SQL to perform counts in Azure Databricks. I will explain it with a practical example. So don’t waste time let’s start with a step-by-step guide to understanding how to use the count() function in PySpark Azure Databricks.

There are various count() functions in PySpark, and you should choose the one that best suits your needs based on the use case. So, let’s learn the following things:

  • PySpark DataFrame.count()
  • PySpark count() function
  • PySpark GroupedData.count() function
  • PySpark SQL count

The PySpark count() method is used to count the number of records in PySpark DataFrame on Azure Databricks by excluding null/None values.

Syntax: dataframe_name.count()

Apache Spark Official documentation link: count()

Create a simple 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 = [
    (1,"Em","VN"),
    (2,"Emalia","DE"),
    (3,"Odette",None),
    (4,"Mandy","DE"),
    (4,"Mandy","DE")
]

df = spark.createDataFrame(data, schema=["id","name","country"])
df.printSchema()
df.show(truncate=False)

"""
root
 |-- id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- country: string (nullable = true)

+---+------+-------+
|id |name  |country|
+---+------+-------+
|1  |Em    |VN     |
|2  |Emalia|DE     |
|3  |Odette|null   |
|4  |Mandy |DE     |
|4  |Mandy |DE     |
+---+------+-------+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

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

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

"""
root
 |-- id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- country: string (nullable = true)
"""

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

How to count number of rows in PySpark DataFrame using Azure Databricks?

In this section, let’s see how to count the number of records in PySpark DataFrame using Azure Databricks with some examples.

Example 1:

In the below example, we are trying to get the number of records in our DataFrame.

rows = df.count()
print(f"Number of rows: {rows}")

"""
Output:

Number of rows: 5

"""

Example 2:

In the below example, we are trying to get the number of unique records in our DataFrame.

uniques_rows = df.distinct().count()
print(f"Number of unique rows: {uniques_rows}")

"""
Output:

Number of unique rows: 4

"""

As you can see that the record ‘4,Mandy,DE’ was repeated twice. The distinct() function returns all the unique records and that’s how we get the number of the unique row as 4.

How to find number of records in PySpark Azure Databricks using count() function

In this section, let’s see how to find the number of records in PySpark DataFrame Azure Databricks using the count() function with an example.

Example:

In the below example, we are trying to get the number of non-null records of each column of DataFrame.

from pyspark.sql.functions import col, count

df.select(count("id"), count("name"), count("country")).show()

"""
Output:

+---------+-----------+--------------+
|count(id)|count(name)|count(country)|
+---------+-----------+--------------+
|        5|          5|             4|
+---------+-----------+--------------+

"""

As you can see that the third column ‘country’ value has been reduced to 4. This is because the count() neglects the null/None value for counting.

How to count records by grouping data in PySpark DataFrame using Azure Databricks?

In this section, let’s see how to count records by grouping data in PySpark DataFrame using Azure Databricks using various methods.

Example:

In the below example, we are trying to count the records of each country in our DataFrame.

from pyspark.sql.functions import count

# Method 1:
df.select("country").groupBy("country").count().show()

# Method 2:
df.groupBy("country").agg(count("country").alias("count")).show()

# Method 3:
df.groupBy("country").agg({'country': 'count'}) \
.withColumnRenamed("count(country)", "count").show()

# The above all codes generates the following output.

"""
Output:

+-------+-----+
|country|count|
+-------+-----+
|     VN|    1|
|     DE|    3|
|   null|    1|
+-------+-----+

"""

How to count number of records in PySpark Azure Databricks using SQL expression?

In this section, let’s see how to count the number of records in PySpark Azure Databricks using SQL expression with some examples. In order to use a raw SQL expression, we have to convert our DataFrame into a SQL view.

df.createOrReplaceTempView("people")

Example 1:

# Example 1: Count
spark.sql('SELECT COUNT(*) FROM people').show()

"""
Output:

+--------+
|count(1)|
+--------+
|       5|
+--------+

"""

Example 2:

# Example 2: Distinct count
spark.sql('SELECT COUNT(DISTINCT country) FROM people').show()

"""
Output:

+-----------------------+
|count(DISTINCT country)|
+-----------------------+
|                      2|
+-----------------------+

"""

Example 3:

# Example 3: Groupby count
spark.sql('''
    SELECT country, COUNT(country)
    FROM people
    GROUP BY country
''').show()

"""
Output:

+-------+--------------+
|country|count(country)|
+-------+--------------+
|     VN|             1|
|     DE|             3|
|   null|             0|
+-------+--------------+

"""

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

When should you count records in PySpark Azure Databricks?

These could be the possible reasons:

  1. When you want to count the number of records in a DataFrame.
  2. When you want to get a number of records particular to a Grouped value.

Real World Use Case Scenarios for counting records in PySpark Azure Databricks?

  • Assume you have a large Dataset and you want to get the number of records in a DataFrame, you can use the count() function all get the number of records.
  • Assume you have an employee DataFrame with their id, name, and their state. You can find the number of employees from each state by using the group() function with the count() function.

What are the alternatives for counting records in PySpark Azure Databricks?

There are multiple alternatives for counting records in PySpark Azure Databricks, which are as follows:

  • collect(): Using this function we can fetch all the records from DataFrame and by using the Python len() function on top of the list we can have the row counts. Always keep an eye on this function, because the collect() function sends result back to driver memory and might result in out of memory issue.

Final Thoughts

In this article, we have learned about counting records in PySpark 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.

PySpark in Azure Databricks, as explained by Arud Seka Berne S on azurelib.com.

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.