Pyspark: dynamically generate condition for when() clause during runtime

问题

I have read a csv file into pyspark dataframe. Now if I apply conditions in when() clause, it works fine when the conditions are given before runtime.

import pandas as pd
from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql import functions
from pyspark.sql.functions import col

sc = SparkContext('local', 'example')
sql_sc = SQLContext(sc)

pandas_df = pd.read_csv('file.csv')  # assuming the file contains a header
# Sample content of csv file
# col1,value
# 1,aa
# 2,bbb

s_df = sql_sc.createDataFrame(pandas_df)

new_df = s_df.withColumn('value', functions.when((col("col1") == 2) | (col("value") == "aa"), s_df.value).otherwise(
    2))

new_df.show(truncate=False)

But I need to dynamically form the conditions inside when clause from a list.

[{'column': 'col1', 'operator': '==', 'value': 2}, {'column': 'value', 'operator': '==', 'value': "aa"}]

Is there any way to achieve this?

Thanks in advance.

回答1:

You can:

1. dynamically generate the SQL string, Python 3.6+' f-strings are really convenient for this.
2. pass it to the pyspark.sql.functions.expr to make a pyspark.sql.column.Column out of it.

---

For your example, something like this should work:

Given s_df 's schema:

root
 |-- col1: long (nullable = false)
 |-- value: string (nullable = false)

Importing functions and instantiate your conditions collection:

[...]
from pyspark.sql.functions import col, expr, when
conditions = [
    {'column': 'col1', 'operator': '==', 'value':  3}, 
    {'column': 'value', 'operator': '==', 'value': "'aa'"}
]

• With generation of the entire if statement:

new_df = s_df.withColumn('value', expr(
    f"if({conditions[0]['column']}{conditions[0]['operator']}{conditions[0]['value']}"
    f" OR {conditions[1]['column']}{conditions[1]['operator']}{conditions[1]['value']},"
    "value, 2)")).show()

• Or with only the generation of the condition, passed to the when function.

new_df = s_df.withColumn('value',when(
    expr(
        f"{conditions[0]['column']}{conditions[0]['operator']}{conditions[0]['value']}"
        f" OR {conditions[1]['column']}{conditions[1]['operator']}{conditions[1]['value']}"
    ),
    col("value")).otherwise(2)).show()

来源：https://stackoverflow.com/questions/58762655/pyspark-dynamically-generate-condition-for-when-clause-during-runtime