Fill a column according to a weight list in another Spark DataFrame

Cri 0
Viewed 37

I have:

Dataframe A

ID Name
1 Null
2 Null
3 Null
4 Null

Dataframe B

Name weight
Max 0.5
Mike 0.25
John 0.25

I need to fill (in PySpark) the column "Name" (Dataset A) according to the weighted Name of Dataset B:

ID Name
1 Max
2 Max
3 Mike
4 John

Max 50%; Mike 25%; John 25%

join dataframe pyspark apache-spark apache-spark-sql
1 Answers

This may not be the best approach, but it would work for weights >= 0.01. I wanted to create an approach where equality would be used in the join as opposed to > < >= and <=.

Input:

from pyspark.sql import functions as F, Window as W
dfA = spark.range(1, 5).toDF("ID")
dfB = spark.createDataFrame(
    [('Max', 0.5),
     ('Mike', 0.25),
     ('John', 0.25)],
    ['Name', 'weight'])

Script:

cs = F.round(F.sum('weight').over(W.rowsBetween(W.unboundedPreceding, 0)), 2)
dfB = dfB.withColumn('w', F.sequence(((cs - F.col('weight'))*100+1).cast('int'), (cs*100).cast('int')))
dfB = dfB.withColumn('w', F.explode('w'))

cntA = dfA.count()
dfA = dfA.withColumn('w', F.ceil(F.count('ID').over(W.rowsBetween(W.unboundedPreceding, 0))/cntA*100))

dfA = dfA.join(dfB, 'w', 'left').drop('w', 'weight')

dfA.show()
# +---+----+
# | ID|Name|
# +---+----+
# |  1| Max|
# |  2| Max|
# |  3|Mike|
# |  4|John|
# +---+----+
ZygD0
Related