Pyspark keep state within tasks

user_12024969 0
Viewed 58

This is related to this question: Pyspark dataframe column value dependent on value from another row but this one gets even more complicated. I have a dataframe:

columns = ['id','seq','manufacturer']
data = [("1",1,"Factory"), ("1",2,"Sub-Factory-1"), ("1",3,"Order"),("1",4,"Sub-Factory-1"),("2",1,"Factory"), ("2",2,"Sub-Factory-1"), ("2",5,"Sub-Factory-1"),("3",1, "Sub-Factory-1"),("3",2,"Order"), ("3",4, "Sub-Factory-1"), ("4", 1,"Factory"), ("4",3, "Sub-Factory-1"),("4",4, "Sub-Factory-1"),("5",1,"Sub-Factory-1"), ("5",2, "Sub-Factory-1"), ("5", 6,"Order"), ("6",2,"Factory"), ("6",3, "Order"), ("6",4,"Sub-Factory-1"), ("6", 6,"Sub-Factory-1"), ("6",7,"Order"), ("7",1,"Sub-Factory-1"), ("7",2,"Factory" ), ("7", 3,"Order"), ("7", 4,"Sub-Factory-1"),("7",5,"Factory"), ("7",8, "Sub-Factory-1"),("7",10,"Sub-Factory-1")]
rdd = spark.sparkContext.parallelize(data)
df = rdd.toDF(columns)
df.orderBy('id','seq').show(40)
       
+---+---+-------------+
| id|seq| manufacturer|
+---+---+-------------+
|  1|  1|      Factory|
|  1|  2|Sub-Factory-1|
|  1|  3|        Order|
|  1|  4|Sub-Factory-1|
|  2|  1|      Factory|
|  2|  2|Sub-Factory-1|
|  2|  5|Sub-Factory-1|
|  3|  1|Sub-Factory-1|
|  3|  2|        Order|
|  3|  4|Sub-Factory-1|
|  4|  1|      Factory|
|  4|  3|Sub-Factory-1|
|  4|  4|Sub-Factory-1|
|  5|  1|Sub-Factory-1|
|  5|  2|Sub-Factory-1|
|  5|  6|        Order|
|  6|  2|      Factory|
|  6|  3|        Order|
|  6|  4|Sub-Factory-1|
|  6|  6|Sub-Factory-1|
|  6|  7|        Order|
|  7|  1|Sub-Factory-1|
|  7|  2|      Factory|
|  7|  3|        Order|
|  7|  4|Sub-Factory-1|
|  7|  5|      Factory|
|  7|  8|Sub-Factory-1|
|  7| 10|Sub-Factory-1|
+---+---+-------------+

What I want to do is to assign hierarchical values to another column(not saying its the best idea) that I can use with the logic from Pyspark dataframe column value dependent on value from another row. So within id group and seq order I want only the first Sub-Factory to attribute to Factory, if there is a Factory within same id and seq order above the Sub-Factory. So end result should look like:

+---+---+-------------+-------+
| id|seq| manufacturer|checker|
+---+---+-------------+-------+
|  1|  1|      Factory|      1|
|  1|  2|Sub-Factory-1|      1|
|  1|  3|        Order|      0|
|  1|  4|Sub-Factory-1|      0|
|  2|  1|      Factory|      1|
|  2|  2|Sub-Factory-1|      1|
|  2|  5|Sub-Factory-1|      0|
|  3|  1|Sub-Factory-1|      0|
|  3|  2|        Order|      0|
|  3|  4|Sub-Factory-1|      0|
|  4|  1|      Factory|      1|
|  4|  3|Sub-Factory-1|      1|
|  4|  4|Sub-Factory-1|      0|
|  5|  1|Sub-Factory-1|      0|
|  5|  2|Sub-Factory-1|      0|
|  5|  6|        Order|      0|
|  6|  2|      Factory|      1|
|  6|  3|        Order|      0|
|  6|  4|Sub-Factory-1|      1|
|  6|  6|Sub-Factory-1|      0|
|  6|  7|        Order|      0|
|  7|  1|Sub-Factory-1|      0|
|  7|  2|      Factory|      1|
|  7|  3|        Order|      0|
|  7|  4|Sub-Factory-1|      1|
|  7|  5|      Factory|      1|
|  7|  8|Sub-Factory-1|      1|
|  7| 10|Sub-Factory-1|      0|
+---+---+-------------+-------+

The dataset is large so I can't use something like df.collect() and then loop over data because it crashes memory. My first idea was to use an accumulator like:

acc = sc.accumulator(0)
def myFunc(manufacturer):
    if  manufacturer == 'Factory':
        acc.value = 1
        return 1
    elif manufacturer == 'Sub-Factory-1' and acc.value == 1:
        acc.value = 0
        return 1
    else:
        return 0
  
myFuncUDF = F.udf(myFunc, IntegerType())
df = df.withColumn('test', myFuncUDF(col('manufacturer')))

But it's a bad idea since accumulator cannot be accessed within tasks. Also Window function solves it if I want to attribute all Sub-Factories from above Factory within same id but now only the first Sub-Factory should get attributed. Any ideas?

pyspark
1 Answers
from pyspark.sql.window import Window
from pyspark.sql.functions import *

df_mod = df.filter(df.manufacturer == 'Sub-Factory-1')

W  = Window.partitionBy("id").orderBy("seq")

df_mod = df_mod.withColumn("rank",rank().over(W))
df_mod = df_mod.filter(col('rank') == 1)

df_mod2 = df.filter(col('manufacturer') == 'Factory')\
            .select('id', 'seq', col('manufacturer').alias('Factory_chk_2'))

df_f = df\
          .join(df_mod, ['id', 'seq'], 'left')\
          .select('id', 'seq', df.manufacturer, 'rank')\
          .join(df_mod2, 'id', 'left')\
          .select('id', df.seq, df.manufacturer, 'rank', 'Factory_chk_2')\
          .withColumn('Factory_chk', when(df.manufacturer=='Factory', 1))\
          .withColumn('Factory_chk_2', when(col('Factory_chk_2')=='Factory', 1))\
          .withColumn('checker',when(col('Factory_chk_2')=='1', coalesce(col('rank'),col('Factory_chk'))).otherwise(lit(0)))\
          .select('id', 'seq', 'manufacturer', 'checker')\
          .na.fill(value=0)\
          .orderBy('id', 'seq')

df_f.show()

+---+---+-------------+-------+
| id|seq| manufacturer|checker|
+---+---+-------------+-------+
|  1|  1|      Factory|      1|
|  1|  2|Sub-Factory-1|      1|
|  1|  3|        Order|      0|
|  1|  4|Sub-Factory-1|      0|
|  2|  1|      Factory|      1|
|  2|  2|Sub-Factory-1|      1|
|  2|  5|Sub-Factory-1|      0|
|  3|  1|Sub-Factory-1|      0|
|  3|  2|        Order|      0|
|  3|  4|Sub-Factory-1|      0|
|  4|  1|      Factory|      1|
|  4|  3|Sub-Factory-1|      1|
|  4|  4|Sub-Factory-1|      0|
|  5|  1|Sub-Factory-1|      0|
|  5|  2|Sub-Factory-1|      0|
|  5|  6|        Order|      0|
|  6|  2|      Factory|      1|
|  6|  3|        Order|      0|
|  6|  4|Sub-Factory-1|      1|
|  6|  6|Sub-Factory-1|      0|
+---+---+-------------+-------+
only showing top 20 rows
user_116985620
Related