代码之家 › 专栏 › 技术社区 › Markus

如何将自定义项函数的返回值保存为两列?

apache-spark-sql pyspark apache-spark python-3.x python

Markus · 技术社区 · 6 年前

我的职能 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 ,然后把它分成两列。或者有没有更短的路?

3 回复 | 直到 6 年前

hamza tuna 6 年前

您可以在自定义项中创建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'))

Cut7er 6 年前

元组可以像列表一样编制索引,因此可以将列1的值添加为 get_data()[0] 对于第二列中的第二个值 get_data()[1]

v1, v2 = get_data() 通过这种方式将返回的元组值赋给变量 v1 v2

看一看 this 问题在此进一步澄清。

Chandan Ray 6 年前

例如,您有一个列的示例dataframe,如下所示

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

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

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