How to convert a lot of columns from long type to integer type in PySpark?

Viewed 445

I have a dataframe with around 50+ columns all in the "long" format. I would like to batch process 40 of them to be converted to "integer" format.

Do I have to keep repeating below?


df = df \
    .withColumn('colA', col('colA').cast(IntegerType())) \
    .withColumn('colB', col('colB').cast(IntegerType())) \
    .withColumn('colC', col('colC').cast(IntegerType())) \
 ....

Above looks pretty manual to me. I am new to PySpark, so not sure if I can put all columns into a list, and only use cast once (like what I would have done in Python).

Greatly appreciate your help!

3 Answers

Use something like below (if you want to cast all your columns at once) -

from pyspark.sql.functions import col

df.select(*(col(c).cast("integer").alias(c) for c in df.columns))

In this case I would probably use reduce, because in python 3, it has been turned into a c wrapper and it quite fast. Warning- It may be more computationally expensive where you are selectively casting columns because it will scan through the entire dataframe before casting specified columns. Can explore this by running .explain()

from functools import reduce 
out = reduce(
    lambda df, c: df.withColumn(c, df[c].cast('integer')), 
    df.columns,
    df
)#.explain()
from pyspark.sql import functions as F

my_cols = ['col1', 'col2']

for c in my_cols:
  df = df.withColumn(c, F.col(c).cast('Integer'))
Related