Finding Ratio of Every Column Pyspark

user_13831989 0
Viewed 133

Let's say I have the following DF.

df = sqlContext.createDataFrame(sc.parallelize(
    [('a',0, 1,1), ('b',1, 3,4), ('c',2, 5,7)]), ('acc_no','col_1', 'col_2','sum'))

+------+-----+-----+---+
|acc_no|col_1|col_2|sum|
+------+-----+-----+---+
|     a|    0|    1|  1|
|     b|    1|    3|  4|
|     c|    2|    5|  7|
+------+-----+-----+---+

I want to find the create a new column for percentage of col_1 and col_2 over the sum column as below.

+------+-----+-----+---+---------+---------+
|acc_no|col_1|col_2|sum|perc_col1|perc_col2|
+------+-----+-----+---+---------+---------+
|     a|    0|    1|  1|        0|      100|
|     b|    1|    3|  4|       25|       75|
|     c|    2|    5|  7|       29|       71|
+------+-----+-----+---+---------+---------+

I know I can use the following code to add the percentage column for each col_1 and col_2

df.withColumn('perc_col1',col('col_1')/col('sum')*100).withColumn('perc_col2',col('col_2')/col('sum')*100).show()

But how can I do it if I have hundreds of columns and I'd like to avoid writing hundreds lines of code?

pyspark
1 Answers

Instead of withcolumn, you can use a for loop with df.select in this case and rename using .alias and f-strings

col_list = ['col_1','col_2'] #list of multiple columns
out = df.select("*",*[F.round((F.col(i)/F.col("sum"))*100,2).alias(f"perc_{i}")
                      for i in col_list])

out.show()
+------+-----+-----+---+----------+----------+
|acc_no|col_1|col_2|sum|perc_col_1|perc_col_2|
+------+-----+-----+---+----------+----------+
|     a|    0|    1|  1|       0.0|     100.0|
|     b|    1|    3|  4|      25.0|      75.0|
|     c|    2|    5|  7|     28.57|     71.43|
+------+-----+-----+---+----------+----------+

Or loop with withcolumn:

col_list = ['col_1','col_2']
for i in col_list:
    df = df.withColumn(f"perc_{i}",F.round((F.col(i)/F.col("sum"))*100,2))
df.show()
user_98406370
Related