How to calculate a column in a Row using two columns of the previous Row in Spark Data Frame?

Question

I'm trying to calculate a compound interest (kind-of) using spark (any flavor: pyspark , spark , spark sql , etc).

My data has the following shape:

+------------+------+------+--------+
| population | rate | year | city   |
+------------+------+------+--------+
| 100        | 0.1  | 1    | one    |
+------------+------+------+--------+
| 100        | 0.11 | 2    | one    |
+------------+------+------+--------+
| 100        | 0.12 | 3    | one    |
+------------+------+------+--------+
| 200        | 0.1  | 1    | two    |
+------------+------+------+--------+
| 1000       | 0.21 | 2    | three  |
+------------+------+------+--------+
| 1000       | 0.22 | 3    | three  |
+------------+------+------+--------+

The population column is wrong (it came from a join between two tables, not shown).

I want to update the population column with the result of the population*(1 + rate) of the previous row. I know that in sql I could use recursive CTE s but hiveql doesn't support it.

Could you give me some suggestions?

Answer 1

As far as I understand your description all you need is some basic algebra and window functions. First lets recreate example data:

import pandas as pd  # Just to make a reproducible example

pdf = pd.DataFrame({
    'city': {0: 'one', 1: 'one', 2: 'one', 3: 'two', 4: 'three', 5: 'three'},
    'population': {0: 100, 1: 100, 2: 100, 3: 200, 4: 1000, 5: 1000},
    'rate': {0: 0.10000000000000001,
     1: 0.11,
     2: 0.12,
     3: 0.10000000000000001,
     4: 0.20999999999999999,
     5: 0.22},
    'year': {0: 1, 1: 2, 2: 3, 3: 1, 4: 2, 5: 3}})

df = sqlContext.createDataFrame(pdf)

df.show()
## +-----+----------+----+----+
## | city|population|rate|year|
## +-----+----------+----+----+
## |  one|       100| 0.1|   1|
## |  one|       100|0.11|   2|
## |  one|       100|0.12|   3|
## |  two|       200| 0.1|   1|
## |three|      1000|0.21|   2|
## |three|      1000|0.22|   3|
## +-----+----------+----+----+

Next we lets define the windows:

import sys
from pyspark.sql.window import Window
from pyspark.sql.functions import exp, log, sum, first, col, coalesce

# Base window
w = Window.partitionBy("city").orderBy("year")

# ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING
wr = w.rowsBetween(-sys.maxsize, -1)

and some columns:

# Take a sum of logarithms of rates over the window
log_sum = sum(log(col("rate") + 1)).over(wr)

# Take sum of logs and exponentiate to go back to original space 
cumulative_rate = exp(log_sum).alias("cumulative_rate")

# Find base population for each group
base_population = first("population").over(w).alias("base_population")

# Prepare final column (base population * cumulative product of rates)
current_population = coalesce(
     # This is null for the first observation in a group
     cumulative_rate * base_population, 
     # so we provide population as an alternative
     col("population")  
).alias("current_population")

Finally we can use these as follows

df.select("*", current_population).show()

## +-----+----------+----+----+------------------+
## | city|population|rate|year|current_population|
## +-----+----------+----+----+------------------+
## |three|      1000|0.21|   2|            1000.0|
## |three|      1000|0.22|   3|            1210.0|
## |  two|       200| 0.1|   1|             200.0|
## |  one|       100| 0.1|   1|             100.0|
## |  one|       100|0.11|   2|110.00000000000001|
## |  one|       100|0.12|   3|122.10000000000004|
## +-----+----------+----+----+------------------+