CODEWITHSUNDEEP
pythonapache-sparkpyspark
Cowboy_Owl
Cowboy_Owl

Reputation: 109

PySpark: Calculate values in column X based on previous values in the column Y and ID

I have 3 columns in a PySpark dataframe: ID, Y and an empty column X.

data = spark.read.options(sep="|", header="true", inferschema="true").csv(csv_file)
data.createOrReplaceTempView("TABLE_NAME")
df = spark.sql("SELECT ID, Y, X FROM TABLE_NAME")

df

ID: string, Y: int, X: int

+----+-----+---+
| ID |  Y  | X |
+----+-----+---+
| V1 |   0 |   |
| V2 |   0 |   |
| V1 | 100 |   |
| V2 | 100 |   |
| V1 | 250 |   |
| V2 | 400 |   |
+----+-----+---+

I would like to calculate X as either zero if the current row is the first occurance of the ID or as a difference between the Y in the current occurance of the ID and Y in the last occurance of the ID:

+----+-----+-----+
| ID |  Y  |  X  |
+----+-----+-----+
| V1 |   0 |   0 |
| V2 |   0 |   0 |
| V1 | 100 | 100 |
| V2 | 100 | 100 |
| V1 | 250 | 150 |
| V2 | 400 | 300 |
+----+-----+-----+

Could you please help me to achieve that?

I am very sorry, I couldn't find how to read data by inputing it manually, here is my csv_file for reproducibility:

ID|Y|Date
V1|0|2018-06-22 08:33:05
V2|0|2018-06-22 08:33:05
V1|100|2018-06-22 08:34:05
V2|100|2018-06-22 08:34:05
V1|250|2018-06-22 08:35:05
V2|400|2018-06-22 08:35:05
V2|-50|2018-06-22 08:36:05
V2|400|2018-06-22 08:37:05

Spark version: 2.4.0

EDIT: Using Steven's solution with additional rows:

+---+---+----+
| id|  y|   x|
+---+---+----+
| v2|-50|   0|
| v2|  0|  50|
| v2|100| 150|
| v2|400| 350|
| v2|400| -50|
| v2|400|-450|
| v1|  0|   0|
| v1|100| 100|
| v1|250| 150|
+---+---+----+

Desirable result:

+----+-----+-----+---------------------+
| id |  y  |  x  |        Date         |
+----+-----+-----+---------------------+
| v2 |   0 |   0 | 2018-06-22 08:33:05 |
| v2 | 100 | 100 | 2018-06-22 08:34:05 |
| v2 | 400 | 300 | 2018-06-22 08:35:05 |
| v2 | -50 | -450 | 2018-06-22 08:36:05 |
| v2 | 400 |  450 | 2018-06-22 08:37:05 |
| v1 |   0 |   0 | 2018-06-22 08:33:05 |
| v1 | 100 | 100 | 2018-06-22 08:34:05 |
| v1 | 250 | 150 | 2018-06-22 08:35:05 |
+----+-----+-----+---------------------+

With orderBy("Date"):

+---+---+-------------------+----+
| id|  y|               Date|   x|
+---+---+-------------------+----+
| v2|  0|2018-06-22 08:33:05|   0|
| v2|100|2018-06-22 08:34:05| 100|
| v2|400|2018-06-22 08:35:05| 300|
| v2|-50|2018-06-22 08:36:05|-550|
| v2|400|2018-06-22 08:37:05| -50|
| v1|  0|2018-06-22 08:33:05|   0|
| v1|100|2018-06-22 08:34:05| 100|
| v1|250|2018-06-22 08:35:05| 150|
+---+---+-------------------+----+

Upvotes: 0

Views: 471

Answers (1)

Steven
Steven

Reputation: 15258

here is you dataframe :

df.show()                                                                                                    

+---+---+-------------------+                                                   
| id|  y|               date|
+---+---+-------------------+
| V1|  0|2018-06-22 08:33:05|
| V2|  0|2018-06-22 08:33:05|
| V1|100|2018-06-22 08:34:05|
| V2|100|2018-06-22 08:34:05|
| V1|250|2018-06-22 08:35:05|
| V2|400|2018-06-22 08:35:05|
| V2|-50|2018-06-22 08:36:05|
| V2|400|2018-06-22 08:37:05|
+---+---+-------------------+

You can have your result using lag: 

from pyspark.sql import Window, functions as F

df.withColumn(
    "x", 
    F.coalesce(
        F.col("y") 
        - F.lag("y").over(
            Window.partitionBy(
                "id"
            ).orderBy(
                "date"
            )
        ), 
        F.lit(0)
    )
).show()  

+---+---+-------------------+------+                                            
| id|  y|               date|     x|
+---+---+-------------------+------+
| V2|  0|2018-06-22 08:33:05|   0.0|
| V2|100|2018-06-22 08:34:05| 100.0|
| V2|400|2018-06-22 08:35:05| 300.0|
| V2|-50|2018-06-22 08:36:05|-450.0|
| V2|400|2018-06-22 08:37:05| 450.0|
| V1|  0|2018-06-22 08:33:05|   0.0|
| V1|100|2018-06-22 08:34:05| 100.0|
| V1|250|2018-06-22 08:35:05| 150.0|
+---+---+-------------------+------+

Upvotes: 1

Related Questions