根据条件火花替换值

Question

I have dataset I want to replace the result column based on the least value of quantity by grouping id,date

id,date,quantity,result
1,2016-01-01,245,1
1,2016-01-01,345,3
1,2016-01-01,123,2
1,2016-01-02,120,5
2,2016-01-01,567,1
2,2016-01-01,568,1
2,2016-01-02,453,1

Here the output, replace the quantity which has least value in that groupby(id,date). Here ordering of rows doesn't matter, any order it can be.

id,date,quantity,result
1,2016-01-01,245,2
1,2016-01-01,345,2
1,2016-01-01,123,2
1,2016-01-02,120,5
2,2016-01-01,567,1
2,2016-01-01,568,1
2,2016-01-02,453,1

Answer 1

Use the Window and get the maximum by max .

import pyspark.sql.functions as f
from pyspark.sql import Window

w = Window.partitionBy('id', 'date')

df.withColumn('result', f.when(f.col('quantity') == f.min('quantity').over(w), f.col('result'))) \
  .withColumn('result', f.max('result').over(w)).show(10, False)

+---+----------+--------+------+
|id |date      |quantity|result|
+---+----------+--------+------+
|1  |2016-01-02|120     |5     |
|1  |2016-01-01|245     |2     |
|1  |2016-01-01|345     |2     |
|1  |2016-01-01|123     |2     |
|2  |2016-01-02|453     |1     |
|2  |2016-01-01|567     |1     |
|2  |2016-01-01|568     |1     |
+---+----------+--------+------+