Как сохранить возвращенные значения функции UDF в два столбца?

Question

Моя функция get_data возвращает кортеж: два целых значения.

get_data_udf = udf(lambda id: get_data(spark, id), (IntegerType(), IntegerType()))

Мне нужно разделить их на две колонки val1 и val2. Как я могу это сделать?

dfnew = df \
    .withColumn("val", get_data_udf(col("id")))

Должен ли я сохранить кортеж в столбце, например val, а затем разбить его как-нибудь на две колонки. Или есть более короткий путь?

Answer 1

Например, у вас есть примерный кадр данных из одного столбца, как показано ниже

val df = sc.parallelize(Seq(3)).toDF()
df.show()

// Ниже UDF, который вернет кортеж

def tupleFunction(): (Int,Int) = (1,2)

// мы создадим два новых столбца из указанного выше UDF

df.withColumn("newCol",typedLit(tupleFunction.toString.replace("(","").replace(")","")
.split(","))).select((0 to 1)
.map(i => col("newCol").getItem(i).alias(s"newColFromTuple$i")):_*).show

Answer 2

Вы можете создать structFields в udf для доступа к более позднему времени.

from pyspark.sql.types import *

get_data_udf = udf(lambda id: get_data(spark, id), 
      StructType([StructField('first', IntegerType()), StructField('second', IntegerType())]))
dfnew = df \
    .withColumn("val", get_data_udf(col("id"))) \
    .select('*', 'val.`first`'.alias('first'), 'val.`second`'.alias('second'))

Answer 3

кортежи можно индексировать точно так же, как списки, поэтому вы можете добавить значение для первого столбца как get_data()[0], а для второго значения во втором столбце вы можете сделать get_data()[1]

также вы можете сделать v1, v2 = get_data() и таким образом присвоить возвращенные значения кортежей переменным v1 и v2.

Взгляните на этот вопрос здесь для дальнейшего разъяснения.