如何将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

您可以在udf中创建structFields以便以后访问。

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 2

元组的索引可以像列表一样，因此您可以将第一列的值添加为get_data()[0]，并将第二列的值添加为get_data()[1]

你也可以做v1, v2 = get_data()，这样就可以将返回的元组值赋给变量v1和v2。

请在此处查看this问题以获得进一步说明。

Answer 3

例如，您有一个列的示例数据框，如下所示

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