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 typing notation）

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

即

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

要返回二维数组，请声明：

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

或

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

函数应返回

List[List[int]]

即

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

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

List[Tuple[int, int]]

即

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）是任意的，可以根据您的要求进行调整。

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

1 个答案: