2 个列表和/或 1 个二维数组的 udf 声明

Question

我想声明一个返回 2 1D 数组或 1 2D 数组的 udf（两者的示例都很好）。我知道这适用于 1D：

@udf("array<int>")

但是，我尝试了许多变体，例如以下没有运气：

@udf("array<int>,array<int>")
@udf("array<int>","array<int>")
@udf("array<int,int>")
etc. 

Answer 1

要返回两个列表，您可以使用 struct

@udf("struct<_1: array<int>, _2: array<int>>")

或者

from pyspark.sql.types import ArrayType, StructField, StructType, IntegerType 

@udf(StructType([
    StructField("_1", ArrayType(IntegerType())),
    StructField("_2", ArrayType(IntegerType()))]))

函数应该返回的位置（PEP 484输入符号）

Tuple[List[int], List[int]]

IE

return [1, 2, 3], [4, 5, 6]

返回二维数组声明：

@udf("array<array<int>>")

或者

@udf(ArrayType(ArrayType(IntegerType())))

函数应该返回的位置

List[List[int]]

IE

return [[1, 2, 3], [4, 5, 6]]

如果您返回固定大小的元组数组

List[Tuple[int, int]]

IE

return  [(1, 2), (3, 4), (5, 6)] 

架构应该是

@udf("array<struct<_1: int, _2: int>>")

或者

@udf(ArrayType(StructType([
    StructField("_1", IntegerType()),
    StructField("_2", IntegerType())])))

虽然array<array<int>>，尽管不是规范的，但在这种情况下也应该有效。

注意：

上面使用的名称（_1和_2）的选择是任意的，可以根据您的要求进行调整。