объявление udf для 2 списков и / или 1 2D-массива - PullRequest
Новая
       37

объявление udf для 2 списков и / или 1 2D-массива

0 голосов
/ 29 января 2019

Я хотел бы объявить udf, который возвращает 2 1D массива или 1 2D массив (пример обоих был бы хорош).Я знаю, что это работает для 1D: 

@udf("array<int>")

Однако я безуспешно пробовал много вариантов этого, таких как следующие: 

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

1 Ответ

0 голосов
/ 29 января 2019

Для возврата двух списков вы можете использовать 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]]

т.е. 

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) может быть произвольным, его можно настроить в соответствии с вашими требованиями.

...