Weighted running total with conditional in Pyspark/Hive - pyspark

Question

I have product, brand, percentage and price columns. I want to calculate the sum of the percentage column for the rows above the current row for those with different brand than the current row and also for those with same brand as the current row. I want to weigh them by price. If the price of the products above the current row are more than the current row, I want to down-weigh it by multiplying it by 0.8. How can I do it in PySpark or using using spark.sql? The answer without using multiplying with weight is here.
df = pd.DataFrame({'a': ['a1','a2','a3','a4','a5','a6'],
'brand':['b1','b2','b1', 'b3', 'b2','b1'],
'pct': [40, 30, 10, 8,7,5],
'price':[0.6, 1, 0.5, 0.8, 1, 0.5]})
df = spark.createDataFrame(df)
What I am looking for
product brand pct pct_same_brand pct_different_brand
a1 b1 40 null null
a2 b2 30 null 40
a3 b1 10 32 30
a4 b3 8 null 80
a5 b2 7 24 58
a6 b1 5 40 45
Update:
I have added the below data points to help clarify the problem. As can be seen, one row can be multiplied by 0.8 in one row and by 1.0 in another row.
product brand pct price pct_same_brand pct_different_brand
a1 b1 30 0.6 null null
a2 b2 20 1.3 null 30
a3 b1 10 0.5 30*0.8 20
a4 b3 8 0.8 null 60
a5 b2 7 0.5 20*0.8 48
a6 b1 6 0.8 30*1 + 10*1 35
a7 b2 5 1.5 20*1 + 7*1 54
Update2: In the data that I provided above, the weight per row is the same number (0.8 or 1) but it can also be 1 and 0.8 (0.8 for some of the rows and 1 for other rows)
Example in the below data frame, the multiplier for the last row , for example, should be 0.8 for a6 and 1.0 for the rest of brand b1. :
df = pd.DataFrame({'a': ['a1','a2','a3','a4','a5','a6', 'a7', 'a8', 'a9', 'a10'],
'brand':['b1','b2','b1', 'b3', 'b2','b1','b2', 'b1', 'b1', 'b1'],
'pct': [30, 20, 10, 8, 7,6,5,4,3,2],
'price':[0.6, 1.3, 0.5, 0.8, 0.5, 0.8, 1.5, 0.5, 0.65, 0.7]
})
df = spark.createDataFrame(df)

Answer 1

You can add a weight column to facilitate calculation:
from pyspark.sql import functions as F, Window
df2 = df.withColumn(
'weight',
F.when(
F.col('price') <= F.lag('price').over(
Window.partitionBy('brand')
.orderBy(F.desc('pct'))
),
0.8
).otherwise(1.0)
).withColumn(
'pct_same_brand',
F.col('weight')*F.sum('pct').over(
Window.partitionBy('brand')
.orderBy(F.desc('pct'))
.rowsBetween(Window.unboundedPreceding, -1)
)
).withColumn(
'pct_different_brand',
F.sum('pct').over(
Window.orderBy(F.desc('pct'))
.rowsBetween(Window.unboundedPreceding, -1)
) - F.coalesce(F.col('pct_same_brand'), F.lit(0)) / F.col('weight')
)
df2.show()
+---+-----+---+-----+------+--------------+-------------------+
| a|brand|pct|price|weight|pct_same_brand|pct_different_brand|
+---+-----+---+-----+------+--------------+-------------------+
| a1| b1| 40| 0.6| 1.0| null| null|
| a2| b2| 30| 1.0| 1.0| null| 40.0|
| a3| b1| 10| 0.5| 0.8| 32.0| 30.0|
| a4| b3| 8| 0.8| 1.0| null| 80.0|
| a5| b2| 7| 1.0| 0.8| 24.0| 58.0|
| a6| b1| 5| 0.5| 0.8| 40.0| 45.0|
+---+-----+---+-----+------+--------------+-------------------+
Output for the edited question:
+---+-----+---+-----+------+--------------+-------------------+
| a|brand|pct|price|weight|pct_same_brand|pct_different_brand|
+---+-----+---+-----+------+--------------+-------------------+
| a1| b1| 30| 0.6| 1.0| null| null|
| a2| b2| 20| 1.3| 1.0| null| 30.0|
| a3| b1| 10| 0.5| 0.8| 24.0| 20.0|
| a4| b3| 8| 0.8| 1.0| null| 60.0|
| a5| b2| 7| 0.5| 0.8| 16.0| 48.0|
| a6| b1| 6| 0.8| 1.0| 40.0| 35.0|
| a7| b2| 5| 1.5| 1.0| 27.0| 54.0|
+---+-----+---+-----+------+--------------+-------------------+

Answer 2

If anyone has a similar question, this worked for me.
Basically, I used outer join of the dataframe with itself and assigned the weights. Finally, I used window functions.
df_copy = df.withColumnRenamed('a', 'asin')\
.withColumnRenamed('brand', 'brandd')\
.withColumnRenamed('pct', 'pct2')\
.withColumnRenamed('price', 'price2')
df2 = df.join(df_copy, on = [df.brand == df_copy.brandd], how = 'outer').orderBy('brand')
df3 = df2.filter(~((df2.a == df2.asin) & (df2.brand == df2.brandd))
& (df2.pct <= df2.pct2))
df3 = df3.withColumn('weight', F.when(df3.price2 > df3.price, 0.8).otherwise(1))
df4 = df3.groupBy(['a', 'brand', 'pct', 'price']).agg(F.sum(df3.pct2*df3.weight)
.alias('same_brand_pct'))
df5 = df.join(df4, on = ['a', 'brand', 'pct', 'price'], how = 'left')
df6 = df5.withColumn(
'pct_same_brand_unscaled',
F.sum('pct').over(
Window.partitionBy('brand')
.orderBy(F.desc('pct'))
.rowsBetween(Window.unboundedPreceding, -1)
)
).withColumn(
'pct_different_brand',
F.sum('pct').over(
Window.orderBy(F.desc('pct'))
.rowsBetween(Window.unboundedPreceding, -1)
) - F.coalesce(F.col('pct_same_brand_unscaled'), F.lit(0))
).drop('pct_same_brand_unscaled')
gives:
+---+-----+---+-----+--------------+-------------------+
| a|brand|pct|price|same_brand_pct|pct_different_brand|
+---+-----+---+-----+--------------+-------------------+
| a1| b1| 30| 0.6| null| null|
| a2| b2| 20| 1.3| null| 30|
| a3| b1| 10| 0.5| 24.0| 20|
| a4| b3| 8| 0.8| null| 60|
| a5| b2| 7| 0.5| 16.0| 48|
| a6| b1| 6| 0.8| 40.0| 35|
| a7| b2| 5| 1.5| 27.0| 54|
| a8| b1| 4| 0.5| 38.8| 40|
| a9| b1| 3| 0.65| 48.8| 40|
|a10| b1| 2| 0.7| 51.8| 40|```