Pyspark 基本问题 - If Else Equivalent

Question

您好，这是一个非常基本的问题，但我是 Pyspark 的新手。我希望我的函数根据输入参数 return 不同的列，但不知道如何执行此操作。 Python 等价物是：

if model='a': return df[[colA,colB]] if model ='b': return df[[colA,colB,colC]]

提前致谢

Answer 1

Pyspark 等效项是使用 select 获取所需的列 -

if model = 'a':
   return df.select(*[colA,colB])
elif model = 'b':
   return df.select(*[colA,colB,colC])

Basic Pyspark Question - If Else Equivalent