How to number records in PySpark Azure Databricks?

Are you looking to find out how to add the row number of PySpark DataFrame in Azure Databricks cloud or maybe you are looking for a solution, to add row number based on grouped records in PySpark Databricks using the row_number() 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 row_number() 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 row_number() function in PySpark.

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

  • Syntax of row_number() functions
  • Adding row number to each record
  • Adding row numbers based on column values in descending order
  • Adding row numbers based on grouped column

The PySpark function row_number() is a window function used to assign a sequential row number, starting with 1, to each window partition’s result in Azure Databricks.

Syntax:

row_number().over()

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

The syntax is as follows:

row_number().over(window_spec)
Parameter NameRequiredDescription
window_spec(WindowSpec)YesIt represents the windowing column.
Table 1: row_number() Method in PySpark Databricks Parameter list with Details

Apache Spark Official Documentation Link: row_number()

Create a simple DataFrame

Let’s understand the use of the row_number() function with a variety of 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 = [
    ("Lewis","McLaren",10.0),
    ("Nick","McLaren",2.0),
    ("Nico","McLaren",6.0),
    ("Fernando","McLaren",3.0),
    ("Heikki","McLaren",8.0),
    ("Kazuki","Ferrari",9.0),
    ("Sébastien","Ferrari",7.0),
    ("Kimi","Ferrari",6.0)
]

df = spark.createDataFrame(data, schema=["driver_name","team","points"])
df.printSchema()
df.show(truncate=False)

"""
root
 |-- driver_name: string (nullable = true)
 |-- team: string (nullable = true)
 |-- points: double (nullable = true)

+-----------+-------+------+
|driver_name|team   |points|
+-----------+-------+------+
|Lewis      |McLaren|10.0  |
|Nick       |McLaren|2.0   |
|Nico       |McLaren|6.0   |
|Fernando   |McLaren|3.0   |
|Heikki     |McLaren|8.0   |
|Kazuki     |Ferrari|9.0   |
|Sébastien  |Ferrari|7.0   |
|Kimi       |Ferrari|6.0   |
+-----------+-------+------+
"""

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
 |-- driver_name: string (nullable = true)
 |-- team: string (nullable = true)
 |-- points: double (nullable = true)
"""

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

How to assign row numbers to each record of PySpark DataFrame using Azure Databricks?

Let’s see how to assign row numbers to each record of a PySpark DataFrame in Azure Databricks using various methods.

Example:

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

window_spec = Window.orderBy("driver_name")

df \
.withColumn("id", row_number().over(window_spec)) \
.select("id", "driver_name", "team", "points").show()

"""
Output:

+---+-----------+-------+------+
| id|driver_name|   team|points|
+---+-----------+-------+------+
|  1|   Fernando|McLaren|   3.0|
|  2|     Heikki|McLaren|   8.0|
|  3|     Kazuki|Ferrari|   9.0|
|  4|       Kimi|Ferrari|   6.0|
|  5|      Lewis|McLaren|  10.0|
|  6|       Nick|McLaren|   2.0|
|  7|       Nico|McLaren|   6.0|
|  8|  Sébastien|Ferrari|   7.0|
+---+-----------+-------+------+

"""

As you can, we have numbered each row based on “driver_name” column ascendingly.

How to assign row numbers to each record based on column values descendingly of PySpark DataFrame using Azure Databricks?

Let’s see how to assign row numbers to each record based on column values descendingly of a PySpark DataFrame in Azure Databricks using various methods.

Example:

from pyspark.sql.functions import row_number, col
from pyspark.sql.window import Window

window_spec = Window.orderBy(col("points").desc())
# The window orderBy() -> acts as on which order the row has be numbered

df \
.withColumn("id", row_number().over(window_spec)) \
.select("id", "driver_name", "team", "points").show()

"""
Output:

+---+-----------+-------+------+
| id|driver_name|   team|points|
+---+-----------+-------+------+
|  1|      Lewis|McLaren|  10.0|
|  2|     Kazuki|Ferrari|   9.0|
|  3|     Heikki|McLaren|   8.0|
|  4|  Sébastien|Ferrari|   7.0|
|  5|       Nico|McLaren|   6.0|
|  6|       Kimi|Ferrari|   6.0|
|  7|   Fernando|McLaren|   3.0|
|  8|       Nick|McLaren|   2.0|
+---+-----------+-------+------+

"""

As you can, the row gets numbered from 1 to 8 based on the points descendingly.

How to assign row numbers based on a group of records of PySpark DataFrame using Azure Databricks?

Let’s see how to assign row numbers based on a group of records of a PySpark DataFrame in Azure Databricks using various methods.

Example:

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

window_spec = Window.partitionBy("team").orderBy(col("points").desc())
# The window partitionBy() -> acts as groupBy

df\
.withColumn("id", row_number().over(window_spec))\
.select("id","team", "driver_name", "points").show()

"""
Output:

+---+-------+-----------+------+
| id|   team|driver_name|points|
+---+-------+-----------+------+
|  1|Ferrari|     Kazuki|   9.0|
|  2|Ferrari|  Sébastien|   7.0|
|  3|Ferrari|       Kimi|   6.0|
|  1|McLaren|      Lewis|  10.0|
|  2|McLaren|     Heikki|   8.0|
|  3|McLaren|       Nico|   6.0|
|  4|McLaren|   Fernando|   3.0|
|  5|McLaren|       Nick|   2.0|
+---+-------+-----------+------+

"""

As you can, each team gets numbered from 1 based on the points descendingly.

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 row_number() in PySpark Azure Databricks?

These could be the possible reasons:

  1. To number each record, for eg. the Serial number
  2. To add a unique identifier for each row

Real World Use Case Scenarios of using row_number() in PySpark Azure Databricks?

Assume that you have a dataset without any identifier column, for example, serial number, or unique ID. The PySpark row_number() function helps in numbering each row consecutively. By using this function you can number your records.

What are the alternatives to the row_number() function in PySpark Azure Databricks?

There are some alternatives to the row_number() function, which are as follows:

  • monotonically_increasing_id(): used for unique increasing ID, but the order will not be consecutive.

Final Thoughts

In this article, we have learned about the PySpark srow_number() 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.