Apache Spark/PySpark - How to Caculate a Column Values Incrementally?

Question

Data frame below is already populated with some data.

(A)
--------------------------------
| id | date       | some_count |
--------------------------------
| 3  | 2020-03-31 |    5       | 
| 2  | 2020-03-24 |    6       |
| 1  | 2020-03-17 |    3       |
--------------------------------

I want to create another data frame based on the above, but with an additional column, that contains change from the previous week, of some_count field for every week. (change of the first record is taken as 0, because it has no previous week record to compare with)

(B)
-----------------------------------------------
| id | date       | some_count | count_change |
-----------------------------------------------
| 3  | 2020-03-31 |    5       |    -1        | 
| 2  | 2020-03-24 |    6       |     3        |
| 1  | 2020-03-17 |    3       |     0        |
-----------------------------------------------

How to do this calculation with Apache Spark (SQL/PySpark)?

Does the order of id correspond to the order of weekly date? — ernest_k
yes, just included it for the clarity, but column to be considered is the date actually. — Keerthi Bandara

Kideok Kim Kideok Kim · Accepted Answer · 2020-03-31T23:29:18

Because order is important to calculate 'count_change', I think we can load it into driver memory, calculate, and re-create another dataframe.
The sample code is implemented by Java, but I believe there is exactly same way in python also.

@Test
public void test() {
    StructType schema = createStructType(Arrays.asList(
            createStructField("id", IntegerType, true),
            createStructField("date", StringType, true),
            createStructField("some_count", IntegerType, true)));

    // assume source data is already sorted by desc.
    Dataset<Row> data = ss.createDataFrame(Arrays.asList(
            RowFactory.create(3, "2020-03-31", 5),
            RowFactory.create(2, "2020-03-24", 6),
            RowFactory.create(1, "2020-03-17", 3)), schema);

    // add column and set 0 as default value.
    Dataset<Row> dataWithColumn = data.withColumn("count_change", lit(0));

    // load driver memory to calculate 'count_change' based on order.
    List<Row> dataWithColumnList = dataWithColumn.collectAsList();
    List<Row> newList = new ArrayList<>();

    // add first row which has count_change 0.
    newList.add(dataWithColumnList.get(dataWithColumnList.size() - 1));

    for (int i = dataWithColumnList.size() - 2; i >= 0; i--) {
        Row currWeek = dataWithColumnList.get(i);
        Row prevWeek = dataWithColumnList.get(i+1);
        int currCount = currWeek.getInt(currWeek.fieldIndex("some_count"));
        int prevCount = prevWeek.getInt(prevWeek.fieldIndex("some_count"));
        int countChange = currCount - prevCount;
        newList.add(RowFactory.create(currWeek.get(0), currWeek.get(1), currWeek.get(2), countChange));
    }

    Dataset<Row> result = ss.createDataFrame(newList, dataWithColumn.schema()).sort(col("date").desc());
    result.show();
}

This is a result of show():

+---+----------+----------+------------+
| id|      date|some_count|count_change|
+---+----------+----------+------------+
|  3|2020-03-31|         5|          -1|
|  2|2020-03-24|         6|           3|
|  1|2020-03-17|         3|           0|
+---+----------+----------+------------+

Apache Spark/PySpark - How to Caculate a Column Values Incrementally?

1 Answers