Dataframe 上的 Pyspark UDF 列

Question

我正在尝试根据某些列的值在数据框上创建一个新列。它在所有情况下都返回 null。任何人都知道这个简单的例子出了什么问题？

df = pd.DataFrame([[0,1,0],[1,0,0],[1,1,1]],columns = ['Foo','Bar','Baz'])

spark_df = spark.createDataFrame(df)

def get_profile():

    if 'Foo'==1:

        return 'Foo'

    elif 'Bar' == 1:

        return 'Bar'

    elif 'Baz' ==1 :

        return 'Baz'

spark_df = spark_df.withColumn('get_profile', lit(get_profile()))

spark_df.show()

   Foo  Bar  Baz get_profile

    0    1    0        None

    1    0    0        None

    1    1    1        None

我希望所有行都填写 get_profile 列。

我也试过这个：

spark_udf = udf(get_profile,StringType())

spark_df = spark_df.withColumn('get_profile', spark_udf())

print(spark_df.toPandas())

达到同样的效果。