Window Functions

In [1]:

# Spark Session
from pyspark.sql import SparkSession

spark = (
    SparkSession
    .builder
    .appName("Unique data & Window Functions")
    .master("local[*]")
    .getOrCreate()
)

spark

# Spark Session from pyspark.sql import SparkSession spark = ( SparkSession .builder .appName("Unique data & Window Functions") .master("local[*]") .getOrCreate() ) spark

Setting default log level to "WARN".
To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).
25/02/15 12:47:54 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable
25/02/15 12:47:55 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.

Out[1]:

SparkSession - in-memory

SparkContext

Spark UI

Version

v3.5.4

Master

local[*]

AppName

Unique data & Window Functions

In [2]:

# Emp Data & Schema

emp_data = [
    ["001","101","John Doe","30","Male","50000","2015-01-01"],
    ["002","101","Jane Smith","25","Female","45000","2016-02-15"],
    ["003","102","Bob Brown","35","Male","55000","2014-05-01"],
    ["004","102","Alice Lee","28","Female","48000","2017-09-30"],
    ["005","103","Jack Chan","40","Male","60000","2013-04-01"],
    ["006","103","Jill Wong","32","Female","52000","2018-07-01"],
    ["007","101","James Johnson","42","Male","70000","2012-03-15"],
    ["008","102","Kate Kim","29","Female","51000","2019-10-01"],
    ["009","103","Tom Tan","33","Male","58000","2016-06-01"],
    ["010","104","Lisa Lee","27","Female","47000","2018-08-01"],
    ["011","104","David Park","38","Male","65000","2015-11-01"],
    ["012","105","Susan Chen","31","Female","54000","2017-02-15"],
    ["013","106","Brian Kim","45","Male","75000","2011-07-01"],
    ["014","107","Emily Lee","26","Female","46000","2019-01-01"],
    ["015","106","Michael Lee","37","Male","63000","2014-09-30"],
    ["016","107","Kelly Zhang","30","Female","49000","2018-04-01"],
    ["017","105","George Wang","34","Male","57000","2016-03-15"],
    ["018","104","Nancy Liu","29","","50000","2017-06-01"],
    ["019","103","Steven Chen","36","Male","62000","2015-08-01"],
    ["020","102","Grace Kim","32","Female","53000","2018-11-01"]
]

emp_schema = "employee_id string, department_id string, name string, age string, gender string, salary string, hire_date string"

# Emp Data & Schema emp_data = [ ["001","101","John Doe","30","Male","50000","2015-01-01"], ["002","101","Jane Smith","25","Female","45000","2016-02-15"], ["003","102","Bob Brown","35","Male","55000","2014-05-01"], ["004","102","Alice Lee","28","Female","48000","2017-09-30"], ["005","103","Jack Chan","40","Male","60000","2013-04-01"], ["006","103","Jill Wong","32","Female","52000","2018-07-01"], ["007","101","James Johnson","42","Male","70000","2012-03-15"], ["008","102","Kate Kim","29","Female","51000","2019-10-01"], ["009","103","Tom Tan","33","Male","58000","2016-06-01"], ["010","104","Lisa Lee","27","Female","47000","2018-08-01"], ["011","104","David Park","38","Male","65000","2015-11-01"], ["012","105","Susan Chen","31","Female","54000","2017-02-15"], ["013","106","Brian Kim","45","Male","75000","2011-07-01"], ["014","107","Emily Lee","26","Female","46000","2019-01-01"], ["015","106","Michael Lee","37","Male","63000","2014-09-30"], ["016","107","Kelly Zhang","30","Female","49000","2018-04-01"], ["017","105","George Wang","34","Male","57000","2016-03-15"], ["018","104","Nancy Liu","29","","50000","2017-06-01"], ["019","103","Steven Chen","36","Male","62000","2015-08-01"], ["020","102","Grace Kim","32","Female","53000","2018-11-01"] ] emp_schema = "employee_id string, department_id string, name string, age string, gender string, salary string, hire_date string"

In [3]:

# Create emp DataFrame

emp = spark.createDataFrame(data=emp_data, schema=emp_schema)

# Create emp DataFrame emp = spark.createDataFrame(data=emp_data, schema=emp_schema)

In [4]:

# Unique of department_ids
# select distinct department_id from emp
emp_dept_id = emp.select("department_id").distinct()

# Unique of department_ids # select distinct department_id from emp emp_dept_id = emp.select("department_id").distinct()

In [5]:

# Window Functions
# select *, max(salary) over(partition by department_id order by salary desc) as max_salary from emp_unique
from pyspark.sql.window import Window
from pyspark.sql.functions import max, col, desc

window_spec = Window.partitionBy(col("department_id")).orderBy(col("salary").desc())
max_func = max(col("salary")).over(window_spec)

emp_1 = emp.withColumn("max_salary", max_func)

# Window Functions # select *, max(salary) over(partition by department_id order by salary desc) as max_salary from emp_unique from pyspark.sql.window import Window from pyspark.sql.functions import max, col, desc window_spec = Window.partitionBy(col("department_id")).orderBy(col("salary").desc()) max_func = max(col("salary")).over(window_spec) emp_1 = emp.withColumn("max_salary", max_func)

In [7]:

emp_1.show()

emp_1.show()

+-----------+-------------+-------------+---+------+------+----------+----------+
|employee_id|department_id|         name|age|gender|salary| hire_date|max_salary|
+-----------+-------------+-------------+---+------+------+----------+----------+
|        007|          101|James Johnson| 42|  Male| 70000|2012-03-15|     70000|
|        001|          101|     John Doe| 30|  Male| 50000|2015-01-01|     70000|
|        002|          101|   Jane Smith| 25|Female| 45000|2016-02-15|     70000|
|        003|          102|    Bob Brown| 35|  Male| 55000|2014-05-01|     55000|
|        020|          102|    Grace Kim| 32|Female| 53000|2018-11-01|     55000|
|        008|          102|     Kate Kim| 29|Female| 51000|2019-10-01|     55000|
|        004|          102|    Alice Lee| 28|Female| 48000|2017-09-30|     55000|
|        019|          103|  Steven Chen| 36|  Male| 62000|2015-08-01|     62000|
|        005|          103|    Jack Chan| 40|  Male| 60000|2013-04-01|     62000|
|        009|          103|      Tom Tan| 33|  Male| 58000|2016-06-01|     62000|
|        006|          103|    Jill Wong| 32|Female| 52000|2018-07-01|     62000|
|        011|          104|   David Park| 38|  Male| 65000|2015-11-01|     65000|
|        018|          104|    Nancy Liu| 29|      | 50000|2017-06-01|     65000|
|        010|          104|     Lisa Lee| 27|Female| 47000|2018-08-01|     65000|
|        017|          105|  George Wang| 34|  Male| 57000|2016-03-15|     57000|
|        012|          105|   Susan Chen| 31|Female| 54000|2017-02-15|     57000|
|        013|          106|    Brian Kim| 45|  Male| 75000|2011-07-01|     75000|
|        015|          106|  Michael Lee| 37|  Male| 63000|2014-09-30|     75000|
|        016|          107|  Kelly Zhang| 30|Female| 49000|2018-04-01|     49000|
|        014|          107|    Emily Lee| 26|Female| 46000|2019-01-01|     49000|
+-----------+-------------+-------------+---+------+------+----------+----------+

In [8]:

emp_1.show()

# Window Functions - 2nd highest salary of each department
# select *, row_number() over(partition by department_id order by salary desc) as rn from emp_unique where rn = 2
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, desc, col

window_spec = Window.partitionBy(col("department_id")).orderBy(col("salary").desc())
rn = row_number().over(window_spec)

emp_2 = emp.withColumn("rn", rn).where("rn = 2")

emp_1.show() # Window Functions - 2nd highest salary of each department # select *, row_number() over(partition by department_id order by salary desc) as rn from emp_unique where rn = 2 from pyspark.sql.window import Window from pyspark.sql.functions import row_number, desc, col window_spec = Window.partitionBy(col("department_id")).orderBy(col("salary").desc()) rn = row_number().over(window_spec) emp_2 = emp.withColumn("rn", rn).where("rn = 2")

+-----------+-------------+-------------+---+------+------+----------+----------+
|employee_id|department_id|         name|age|gender|salary| hire_date|max_salary|
+-----------+-------------+-------------+---+------+------+----------+----------+
|        007|          101|James Johnson| 42|  Male| 70000|2012-03-15|     70000|
|        001|          101|     John Doe| 30|  Male| 50000|2015-01-01|     70000|
|        002|          101|   Jane Smith| 25|Female| 45000|2016-02-15|     70000|
|        003|          102|    Bob Brown| 35|  Male| 55000|2014-05-01|     55000|
|        020|          102|    Grace Kim| 32|Female| 53000|2018-11-01|     55000|
|        008|          102|     Kate Kim| 29|Female| 51000|2019-10-01|     55000|
|        004|          102|    Alice Lee| 28|Female| 48000|2017-09-30|     55000|
|        019|          103|  Steven Chen| 36|  Male| 62000|2015-08-01|     62000|
|        005|          103|    Jack Chan| 40|  Male| 60000|2013-04-01|     62000|
|        009|          103|      Tom Tan| 33|  Male| 58000|2016-06-01|     62000|
|        006|          103|    Jill Wong| 32|Female| 52000|2018-07-01|     62000|
|        011|          104|   David Park| 38|  Male| 65000|2015-11-01|     65000|
|        018|          104|    Nancy Liu| 29|      | 50000|2017-06-01|     65000|
|        010|          104|     Lisa Lee| 27|Female| 47000|2018-08-01|     65000|
|        017|          105|  George Wang| 34|  Male| 57000|2016-03-15|     57000|
|        012|          105|   Susan Chen| 31|Female| 54000|2017-02-15|     57000|
|        013|          106|    Brian Kim| 45|  Male| 75000|2011-07-01|     75000|
|        015|          106|  Michael Lee| 37|  Male| 63000|2014-09-30|     75000|
|        016|          107|  Kelly Zhang| 30|Female| 49000|2018-04-01|     49000|
|        014|          107|    Emily Lee| 26|Female| 46000|2019-01-01|     49000|
+-----------+-------------+-------------+---+------+------+----------+----------+

In [9]:

emp_2.show()

emp_2.show()

+-----------+-------------+-----------+---+------+------+----------+---+
|employee_id|department_id|       name|age|gender|salary| hire_date| rn|
+-----------+-------------+-----------+---+------+------+----------+---+
|        001|          101|   John Doe| 30|  Male| 50000|2015-01-01|  2|
|        020|          102|  Grace Kim| 32|Female| 53000|2018-11-01|  2|
|        005|          103|  Jack Chan| 40|  Male| 60000|2013-04-01|  2|
|        018|          104|  Nancy Liu| 29|      | 50000|2017-06-01|  2|
|        012|          105| Susan Chen| 31|Female| 54000|2017-02-15|  2|
|        015|          106|Michael Lee| 37|  Male| 63000|2014-09-30|  2|
|        014|          107|  Emily Lee| 26|Female| 46000|2019-01-01|  2|
+-----------+-------------+-----------+---+------+------+----------+---+

In [ ]: