PySpark DataFrame - Add Row Number via row_number() Function

Code description

In Spark SQL, row_number can be used to generate a series of sequential number starting from 1 for each record in the specified window. Examples can be found in this page:  Spark SQL - ROW_NUMBER Window Functions.

This code snippet provides the same approach to implement row_number directly using PySpark DataFrame APIs instead of Spark SQL. It created a window that partitions the data by ACCT attribute and sorts the records in each partition via TXN_DT column in descending order. The frame boundary of the window is defined as unbounded preceding and current row. 

Output:

    +----+------+-------------------+
    |ACCT|   AMT|             TXN_DT|
    +----+------+-------------------+
    | 101| 10.01|2021-01-01 00:00:00|
    | 101|102.01|2021-01-01 00:00:00|
    | 102|  93.0|2021-01-01 00:00:00|
    | 103| 913.1|2021-01-02 00:00:00|
    | 101|900.56|2021-01-03 00:00:00|
    +----+------+-------------------+
    
    +----+------+-------------------+------+
    |ACCT|   AMT|             TXN_DT|rownum|
    +----+------+-------------------+------+
    | 101|900.56|2021-01-03 00:00:00|     1|
    | 101| 10.01|2021-01-01 00:00:00|     2|
    | 101|102.01|2021-01-01 00:00:00|     3|
    | 102|  93.0|2021-01-01 00:00:00|     1|
    | 103| 913.1|2021-01-02 00:00:00|     1|
    +----+------+-------------------+------+  
    

Code snippet

    from pyspark.sql import SparkSession, Window
    from datetime import datetime
    from pyspark.sql.functions import row_number, desc
    
    app_name = "PySpark row_number Window Function"
    master = "local"
    
    spark = SparkSession.builder         .appName(app_name)         .master(master)         .getOrCreate()
    
    spark.sparkContext.setLogLevel("WARN")
    
    data = [
        [101, 10.01, datetime.strptime('2021-01-01', '%Y-%m-%d')],
        [101, 102.01, datetime.strptime('2021-01-01', '%Y-%m-%d')],
        [102, 93.0, datetime.strptime('2021-01-01', '%Y-%m-%d')],
        [103, 913.1, datetime.strptime('2021-01-02', '%Y-%m-%d')],
        [101, 900.56, datetime.strptime('2021-01-03', '%Y-%m-%d')]
    ]
    
    df = spark.createDataFrame(data, ['ACCT', 'AMT', 'TXN_DT'])
    df.show()
    
    window = Window.partitionBy('ACCT').orderBy(desc("TXN_DT")).rowsBetween(
        Window.unboundedPreceding, Window.currentRow)
    df = df.withColumn('rownum', row_number().over(window))
    
    df.show()