How to modify a subset of rows based on condition in PySpark

Viewed 332

I am try to convert all MPa values to Pa. The code that I use in pandas is shown below. How would I translate this into pyspark?

file_df.loc[file_df['Unit'] == 'MPa', 'Value'] = file_df['Value'] * 1000000 #coverts Value to Pa from MPa
file_df.loc[file_df['Unit'] == 'MPa', 'Unit'] = 'Pa' # replace the MPa with Pa
1 Answers

You can replicate those in-place assignments with with when/otherwise as:

from pyspark.sql.functions import when, col, lit

m = sparkdf.Unit == 'MPa'
(sparkdf.withColumn("Value", when(m, col('Value')*1000).otherwise(col('Value')))
        .withColumn("Unit",  when(m, lit('Pa')).otherwise(col('Unit'))))

Small working example:

df = pd.DataFrame({'Unit':['MPa', 'MPb', 'MPc'],
                   'Value':[5, 4, 3]})

sparkdf = spark.createDataFrame(df)
m = sparkdf.Unit == 'MPa'

(sparkdf.withColumn("Value",  when(m, col('Value')*1000).otherwise(col('Value')))
        .withColumn("Unit",  when(m, lit('Pa')).otherwise(col('Unit')))).show()

+----+-----+
|Unit|Value|
+----+-----+
|  Pa| 5000|
| MPb|    4|
| MPc|    3|
+----+-----+
Related