Как обрабатывать возвращаемые значения PySpark UDF в разных типах?

Question

У меня есть фрейм данных с одним столбцом. В каждой строке этого фрейма данных есть список. Например:

df = spark.createDataFrame(
    [
        [[13,23]],
        [[55,65]],
    ],
    ['col',]
)

Затем я определил UDF, который в основном добавляет 1 к первому номеру в списке и добавляет 1,5 ко второму номеру списка.

def calculate(mylist) :
  x = mylist[0] + 1
  y = mylist[1] + 1.5
  return x,y

Проблема в том, что что, когда я применяю эту функцию к моему фрейму данных, он возвращает значение X, но не возвращает значение Y. Я думаю, это потому, что значение Y не является целым числом. Это способ, которым я делаю это.

import pyspark.sql.functions as F
from pyspark.sql.types import IntegerType, ArrayType
func = F.udf(lambda x: calculate(x), ArrayType(IntegerType()))
df.withColumn('vals', func('col')).show()

Что я могу сделать, чтобы получить значение Y, а также значение X? Я упростил UDF и образец фрейма данных, чтобы его было легко понять и решить.

Answer 1

вычисление udf возвращает integer, а также float тип с заданным входом. Если ваше значение варианта использования является целым числом, а второе - float, вы можете вернуть StructType

Если оба должны быть одного типа, вы можете использовать один и тот же код и изменить вычисление udf, которое возвращает оба целых числа

func = F.udf(lambda x: calculate(x), T.StructType(
        [T.StructField("val1", T.IntegerType(), True),
         T.StructField("val2", T.FloatType(), True)]))