Python to Pyspark函数UDF如何输出列表列表

Question

我在Python语言中有一个函数(有许多不同的函数，但情况相同)，我正在将它转换为PySpark，然而，这个函数有一个不同整数类型的列表作为输入，有一个输出是一个列表，其中包含n个整数类型的列表，举个例子：

#I know some libraries are not necessary righ now
import pyspark
from pyspark import SQLContext
from pyspark.sql.types import StructType, StructField, IntegerType, FloatType, StringType, ArrayType
from pyspark.sql.functions import udf
from pyspark.sql import Row
from pyspark.sql import functions as F

my_function_input([4,5,7,8,10, 11]) 
my_function_output[[4, 5], [7, 8], [10, 11]]

这就是我正在尝试的，但在尝试使用它时出现错误

pyspark_my_function = udf(my_function, ArrayType(IntegerType()))

TypeError:无效参数，不是字符串或列: 4、5、6、8、9类型。对于列文字，使用'lit‘、'array’、'struct‘或'create_map’函数。

我还有一些其他的函数，它们有2或3个输出，每个输出中也有列表列表。我该如何改变它们呢？我已经试过了

schema = StructType([StructField("output1", ArrayType(IntegerType()), nullable=False), 
                     StructField("output2", ArrayType(IntegerType()), nullable=False)])

pyspark_function = udf(my_function, schema)

谢谢你们！

Answer 1

看起来UDF的输入和输出都有问题。看一下下面的示例代码。

from pyspark.sql import functions as F
from pyspark.sql import types as T

df = (spark
    .sparkContext
    .parallelize([
        ([4,5,7,8,9],),
    ])
    .toDF(['A'])
 )

def myfunc(a):
    return [[4, 5], [7, 8], [10, 11]]

(df
    .withColumn('test', F.udf(myfunc, T.ArrayType(T.ArrayType(T.IntegerType())))('A'))
    .show()
)

# Output
# +---------------+--------------------+
# |              A|                test|
# +---------------+--------------------+
# |[4, 5, 7, 8, 9]|[[4, 5], [7, 8], ...|
# +---------------+--------------------+