pyspark: for loop calculations over the columns

Question

Anyone know how can I do theses calculations in pyspark?

data = {
    'Name': ['Tom', 'nick', 'krish', 'jack'],
    'Age': [20, 21, 19, 18],
    'CSP': [2, 6, 8, 7],
    'coef': [2, 2, 3, 3]
}
  
# Create DataFrame
df = pd.DataFrame(data)
colsToRecalculate = ['Age','CSP']

for i in range(len(colsToRecalculate)):
    df[colsToRecalculate[i]] =df[colsToRecalculate[i]]/df["coef"]

there are some good answers related to this [here](https://stackoverflow.com/q/33681487/8279585) — samkart, Aug 22 '22 at 15:08

score 1 · Answer 1 · answered Aug 22 '22 at 15:19

You can use select() on spark dataframe and include multiple columns (with different calculations) as parameters. In your case:

df2 = spark.createDataFrame(pd.DataFrame(data))
df2.select(*[(F.col(c) / F.col('coef')).alias(c) for c in colsToRecalculate], 'coef').show()

score 0 · Answer 2 · answered Aug 23 '22 at 07:44

Slight variation to bzu's answer which selects non-listed columns manually within the select. We can use dataframe.columns and check the columns against the colsToRecalculate list - If column is in the list, do the calculation, else leave column as is.

data_sdf. \
    select(*[(func.col(k) / func.col('coef')).alias(k) if k in colsToRecalculate else k for k in data_sdf.columns])

pyspark: for loop calculations over the columns

2 Answers2