Question

我使用的是PySpark的新pandas_udf装饰器，我试图让它接受多列作为输入并返回一系列作为输入，但是，我得到了TypeError: Invalid argument

示例代码

@pandas_udf(df.schema, PandasUDFType.SCALAR)
def fun_function(df_in):
    df_in.loc[df_in['a'] < 0] = 0.0
    return (df_in['a'] - df_in['b']) / df_in['c']

Answer 1

A SCALAR udf希望将熊猫系列作为输入而不是数据框。对于您的情况，无需使用udf。裁剪后，可以从列a，b，c直接计算：

import pyspark.sql.functions as f

df = spark.createDataFrame([[1,2,4],[-1,2,2]], ['a', 'b', 'c'])

clip = lambda x: f.when(df.a < 0, 0).otherwise(x)
df.withColumn('d', (clip(df.a) - clip(df.b)) / clip(df.c)).show()

#+---+---+---+-----+
#|  a|  b|  c|    d|
#+---+---+---+-----+
#|  1|  2|  4|-0.25|
#| -1|  2|  2| null|
#+---+---+---+-----+

如果必须使用pandas_udf，则返回类型必须为double，而不是df.schema，因为您仅返回熊猫系列而不是熊猫数据框< / em>;而且您还需要将列作为Series传递给函数，而不是整个数据帧：

@pandas_udf('double', PandasUDFType.SCALAR) def fun_function(a, b, c): clip = lambda x: x.where(a >= 0, 0) return (clip(a) - clip(b)) / clip(c) df.withColumn('d', fun_function(df.a, df.b, df.c)).show() #+---+---+---+-----+ #| a| b| c| d| #+---+---+---+-----+ #| 1| 2| 4|-0.25| #| -1| 2| 2| null| #+---+---+---+-----+

PySpark。将数据框传递给pandas_udf并返回一系列

1 个答案: