Adding multiple columns in pyspark dataframe using a loop

Question

I need to add a number of columns (4000) into the data frame in pyspark. I am using the withColumn function, but getting assertion error.

df3 = df2.withColumn("['ftr' + str(i) for i in range(0, 4000)]", [expr('ftr[' + str(x) + ']') for x in range(0, 4000)])

Not sure what is wrong.

Chiel · Accepted Answer · 2021-06-02 08:51:59Z

10

We can use .select() instead of .withColumn() to use a list as input to create a similar result as chaining multiple .withColumn()'s. The ["*"] is used to select also every existing column in the dataframe.

import pyspark.sql.functions as F

df2:

+---+
|age|
+---+
| 10|
| 11|
| 13|
+---+

df3 = df2.select(["*"] + [F.lit(f"{x}").alias(f"ftr{x}") for x in range(0,10)])

Results in:

+---+----+----+----+----+----+----+----+----+----+----+
|age|ftr0|ftr1|ftr2|ftr3|ftr4|ftr5|ftr6|ftr7|ftr8|ftr9|
+---+----+----+----+----+----+----+----+----+----+----+
| 10|   0|   1|   2|   3|   4|   5|   6|   7|   8|   9|
| 11|   0|   1|   2|   3|   4|   5|   6|   7|   8|   9|
| 13|   0|   1|   2|   3|   4|   5|   6|   7|   8|   9|
+---+----+----+----+----+----+----+----+----+----+----+

answered Jun 2, 2021 at 8:51

Chiel

2,2601 gold badge14 silver badges29 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Kellen Donohue Over a year ago

This is a much more efficient way to do it compared to calling withColumn in a loop!

BICube · Accepted Answer · 2020-10-31 23:10:56Z

2

Try to do something like this:

df2 = df3
for i in range(0, 4000):
  df2 = df2.withColumn(f"ftr{i}", lit(f"frt{i}"))

answered Oct 31, 2020 at 23:10

BICube

4,7211 gold badge26 silver badges47 bronze badges

5 Comments

renjith Over a year ago

I dont think. it will. it will just add one field-i.e. last one -- ftr3999: string (nullable = false)

BICube Over a year ago

@renjith has you actually tried to run it?. The solutions will add all columns. Note that inside the loop I am using df2 = df2.witthColumn and not df3 = df2.withColumn

renjith Over a year ago

renjith Over a year ago

it works. not sure. why it did not work when i tried first

Sugyan sahu Over a year ago

@renjith How did this looping worked for you. It's not working for me as well.

Collectives™ on Stack Overflow

Adding multiple columns in pyspark dataframe using a loop

2 Answers 2

1 Comment

5 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

5 Comments

Your Answer

Sign up or log in

Post as a guest

Related