Question

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

--prune-empty

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

我也尝试过：

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

达到相同的效果。

Answer 1

udf不知道列名是什么。因此，它将检查if / elif块中的每个条件，所有条件的评估结果均为False。 function will return None。

您必须重写udf才能包含要检查的列：

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

def get_profile(foo, bar, baz):
    if foo == 1:
        return 'Foo'
    elif bar == 1:
        return 'Bar'
    elif baz == 1 :
        return 'Baz'

spark_udf = udf(get_profile, StringType())
spark_df = spark_df.withColumn('get_profile',spark_udf('Foo', 'Bar', 'Baz'))
spark_df.show()
#+---+---+---+-----------+
#|Foo|Bar|Baz|get_profile|
#+---+---+---+-----------+
#|  0|  1|  0|        Bar|
#|  1|  0|  0|        Foo|
#|  1|  1|  1|        Foo|
#+---+---+---+-----------+

如果您有很多列，并且想要全部传递（按顺序）：

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

更一般而言，您可以解压缩任何有序的列列表：

cols_to_pass_to_udf = ['Foo', 'Bar', 'Baz']
spark_df = spark_df.withColumn('get_profile', spark_udf(*cols_to_pass_to_udf ))

但是此特定操作不需要udf。我会这样：

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

spark_df.withColumn(
    "get_profile",
    coalesce(*[when(col(c)==1, lit(c)) for c in spark_df.columns])
).show()
#+---+---+---+-----------+
#|Foo|Bar|Baz|get_profile|
#+---+---+---+-----------+
#|  0|  1|  0|        Bar|
#|  1|  0|  0|        Foo|
#|  1|  1|  1|        Foo|
#+---+---+---+-----------+

之所以可行，是因为如果条件的值为pyspark.sql.functions.when()并且未指定null，则False将默认返回otherwise。然后，pyspark.sql.functions.coalesce的列表理解将返回第一个非空列。

请注意，仅当列的顺序与在udf函数中求值的顺序相同时，这才等效于get_profile。为了更加明确，您应该执行以下操作：

spark_df.withColumn(
    "get_profile",
    coalesce(*[when(col(c)==1, lit(c)) for c in ['Foo', 'Bar', 'Baz'])
).show()

数据框上的Pyspark UDF列

1 个答案: