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 CTEs but hiveql doesn't support it.

Could you give me some suggestions?

I think that the natural way of partition the data is withthe `city` column. I tried with window functions, but it failed (or maybe I failed) But I couldn't figure out how to use in the next row a column that doesn't exist yet — nanounanue, Nov 12 '15 at 05:24
https://stackoverflow.com/questions/29109916/updating-a-dataframe-column-in-spark — Ibraim Ganiev, Nov 12 '15 at 06:49

zero323 · Accepted Answer · 2015-11-13T00:50:57.013

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|
## +-----+----------+----+----+------------------+

Awesome answer! And It works! Just one question, Could you point to the documentation of the method `coalesce`? Or some explanation? — nanounanue, Nov 12 '15 at 20:20
It is a standard SQL coalesce - it chooses first not null column. Since `wr` will be empty on the first element in a group, log sum evaluates to `null` and we move to population. [Docs](https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.functions.coalesce) — zero323, Nov 12 '15 at 23:11

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

1 Answers1