CODEWITHSUNDEEP
apache-sparkpysparkapache-spark-sql
PRAKASH NAYAK
PRAKASH NAYAK

Reputation: 33

cumulative product in pySpark data frame

I have the following spark DataFrame:

+---+---+
|  a|  b|
+---+---+     
|  1|  1|  
|  1|  2|  
|  1|  3|
|  1|  4|
+---+---+

I want to make another column named "c" which contains the cumulative product of "b" over "a". The resulting DataFrame should look like:

+---+---+---+
|  a|  b|  c|
+---+---+---+     
|  1|  1|  1|
|  1|  2|  2|
|  1|  3|  6|
|  1|  4| 24|
+---+---+---+

How can this be done?

Upvotes: 2

Views: 2374

Answers (4)

Yannick
Yannick

Reputation: 51

from pyspark.sql import functions as F, Window
df = spark.createDataFrame([(1, 1), (1, 2), (1, 3), (1, 4), (1, 5)], ['a', 'b'])

window = Window.orderBy('b').rowsBetween(Window.unboundedPreceding,Window.currentRow)

df = df.withColumn('c', F.product(F.col('b')).over(window))

df.show()

+---+---+-----+
|  a|  b|    c|
+---+---+-----+
|  1|  1|  1.0|
|  1|  2|  2.0|
|  1|  3|  6.0|
|  1|  4| 24.0|
|  1|  5|120.0|
+---+---+-----+

Upvotes: 0

Michael
Michael

Reputation: 2566

Here is an alternative approach not using a user-defined function

df = spark.createDataFrame([(1, 1), (1, 2), (1, 3), (1, 4), (1, 5)], ['a', 'b'])
wind = Window.partitionBy("a").rangeBetween(Window.unboundedPreceding, Window.currentRow).orderBy("b")
df2 = df.withColumn("foo", collect_list("b").over(wind))
df2.withColumn("foo2", expr("aggregate(foo, cast(1 as bigint), (acc, x) -> acc * x)")).show()

+---+---+---------------+----+
|  a|  b|            foo|foo2|
+---+---+---------------+----+
|  1|  1|            [1]|   1|
|  1|  2|         [1, 2]|   2|
|  1|  3|      [1, 2, 3]|   6|
|  1|  4|   [1, 2, 3, 4]|  24|
|  1|  5|[1, 2, 3, 4, 5]| 120|
+---+---+---------------+----+

and if you don't really care about the precision you can build a shorter version

import pyspark.sql.functions as psf

df.withColumn("foo", psf.exp(psf.sum(psf.log("b")).over(wind))).show()
+---+---+------------------+
|  a|  b|               foo|
+---+---+------------------+
|  1|  1|               1.0|
|  1|  2|               2.0|
|  1|  3|               6.0|
|  1|  4|23.999999999999993|
|  1|  5|119.99999999999997|
+---+---+------------------

Upvotes: 5

Henrique Florencio
Henrique Florencio

Reputation: 3751

You have to set an order column. In your case I used column 'b'

from pyspark.sql import functions as F, Window, types
from functools import reduce
from operator import mul

df = spark.createDataFrame([(1, 1), (1, 2), (1, 3), (1, 4), (1, 5)], ['a', 'b'])

order_column = 'b'

window = Window.orderBy(order_column)

expr = F.col('a') * F.col('b')

mul_udf = F.udf(lambda x: reduce(mul, x), types.IntegerType())

df = df.withColumn('c', mul_udf(F.collect_list(expr).over(window)))

df.show()

+---+---+---+
|  a|  b|  c|
+---+---+---+
|  1|  1|  1|
|  1|  2|  2|
|  1|  3|  6|
|  1|  4| 24|
|  1|  5|120|
+---+---+---+

Upvotes: 2

Prathik Kini
Prathik Kini

Reputation: 1708

You answer is something similar to this.

import pandas as pd
df = pd.DataFrame({'v':[1,2,3,4,5,6]})
df['prod'] = df.v.cumprod()
   v   prod
0  1     1
1  2     2
2  3     6
3  4    24
4  5   120
5  6   720

Upvotes: -1

Related Questions