Pyspark window lead function with offset value from value of different column

user_96180 0
Viewed 678

I am trying to derive a new column using window lead function, but my offset value to lead function varies depending on the value of the column, here is my example data

inputdata = (("James", "Sales", 3000), \
    ("James", "Sales", 4600),  \
    ("James", "Sales", 4100),   \
    ("James", "Finance", 3000),  \
    ("James", "Sales", 3000),    \
    ("James", "Finance", 3300),  \
    ("James", "Finance", 3900),    \
    ("Kumar", "Marketing", 3000), \
    ("Kumar", "Marketing", 2000),\
    ("James", "Sales", 4100) \
  )
  
columns= ["employee_name", "department", "salary"]
df = spark.createDataFrame(data = inputdata, schema = columns)

input

+-------------+----------+------+
|employee_name|department|salary|
+-------------+----------+------+
|James        |Sales     |3000  |
|James        |Sales     |3000  |
|James        |Sales     |4100  |
|James        |Sales     |4100  |
|James        |Sales     |4600  |
|James        |Finance   |3000  |
|James        |Finance   |3300  |
|James        |Finance   |3900  |
|Kumar        |Marketing |2000  |
|Kumar        |Marketing |3000  |
+-------------+----------+------+

expected output:

Here I have a column called expected value

it needs to be derived from salary column, and the value needs to be next immediate big value

if the next value of the column is same then it should consider next value until it finds the different value,

in below example expected value for first two columns is 4100 which is next big value here 9999 is the default value

enter image description here

I have used window lead function but the problem is offset needs to be constant I mean lead can go forward only till fixed number of records, is there any way to solve this ?

this is what I have tried:

created an offset column that, lead function needs to go far

windowSpec  = Window.partitionBy("employee_name","department","salary").orderBy('employee_name','department','salary')
dfk = df.withColumn("row_number",row_number().over(windowSpec))
dfk = dfk.withColumn("max_row_number", max(col("row_number")+1).over(windowSpec))
dfk = dfk.orderBy("employee_name","department","salary").withColumn('offset', dfk['max_row_number']-dfk['row_number'])

+-------------+----------+------+----------+--------------+------+
|employee_name|department|salary|row_number|max_row_number|offset|
+-------------+----------+------+----------+--------------+------+
|        James|   Finance|  3000|         1|             2|     1|
|        James|   Finance|  3300|         2|             3|     1|
|        James|   Finance|  3900|         3|             4|     1|
|        James|     Sales|  3000|         1|             3|     2|
|        James|     Sales|  3000|         2|             3|     1|
|        James|     Sales|  4100|         3|             5|     2|
|        James|     Sales|  4100|         4|             5|     1|
|        James|     Sales|  4600|         5|             6|     1|
|        Kumar| Marketing|  2000|         1|             2|     1|
|        Kumar| Marketing|  3000|         2|             3|     1|
+-------------+----------+------+----------+--------------+------+

but I am not able to pass offset column as offset to lead function or something like this

windowSpec_2  = Window.partitionBy("employee_name","department").orderBy('employee_name','department','salary','row_number')
dfk.withColumn("*** expected value ****",lead('salary', **dfk['offset']**, 9999).over(windowSpec_2)).show()

I have tried using express/eval but I still see the

iter raise TypeError("Column is not iterable") TypeError: Column is not iterable

is there any way to achieve the above expected results ?

thanks in advance

rolling-computation pyspark apache-spark-sql apache-spark
1 Answers

Your use case is a bit tricky, so I think lead may not be appropriate. You may be better off using a rangeBetweeen window, and get the min value in that window:

import pyspark.sql.functions as F
from pyspark.sql.window import Window

df2 = df.withColumn(
    'expected_value',
    F.coalesce(
        F.min('salary').over(
            Window.partitionBy(
                "employee_name","department"
            ).orderBy(
                'salary'
            ).rangeBetween(
                1, Window.unboundedFollowing
            )
        ),
        F.lit(9999)
    )
)

df2.show()
+-------------+----------+------+--------------+
|employee_name|department|salary|expected_value|
+-------------+----------+------+--------------+
|        James|     Sales|  3000|          4100|
|        James|     Sales|  3000|          4100|
|        James|     Sales|  4100|          4600|
|        James|     Sales|  4100|          4600|
|        James|     Sales|  4600|          9999|
|        Kumar| Marketing|  2000|          3000|
|        Kumar| Marketing|  3000|          9999|
|        James|   Finance|  3000|          3300|
|        James|   Finance|  3300|          3900|
|        James|   Finance|  3900|          9999|
+-------------+----------+------+--------------+
user_141657300
Related