Question

我有一个包含不同模式的spark数据帧列表。示例：

list_df = [df1, df2, df3, df4]
# df1.columns = ['a', 'b']
# df2.columns = ['a', 'b', 'c']
# df3.columns = ['a', 'b', 'c', 'd']
# df4.columns = ['a', 'b', 'c', 'd', 'e']

现在，我想编写一个能够在具有不同列数的数据帧列表上运行的单个udf。

有一篇关于如何使用scala执行此操作的帖子：Spark UDF with varargs，其中udf接收列数组。

但似乎这种方法对python不起作用。有什么建议吗？

感谢。

Answer 1

实际上这种方法在Python中运行得很好：

from pyspark.sql.functions import array, udf

df = sc.parallelize([("a", "b", "c", "d")]).toDF()

f = udf(lambda xs: "+".join(xs))

df.select(f("_1")).show()
## +------------+
## |<lambda>(_1)|
## +------------+
## |           a|
## +------------+

df.select(f(array("_1", "_2"))).show()
## +-----------------------+
## |<lambda>(array(_1, _2))|
## +-----------------------+
## |                    a+b|
## +-----------------------+

df.select(f(array("_1", "_2", "_3"))).show()
## +---------------------------+
## |<lambda>(array(_1, _2, _3))|
## +---------------------------+
## |                      a+b+c|
## +---------------------------+

由于Python UDF不是同一类型的实体，因为它们的Scala对应物不受输入参数的类型和数量的限制，所以你也使用args：

g = udf(lambda *xs: "+".join(xs))

df.select(g("_1", "_2", "_3", "_4")).show()
## +------------------------+
## |<lambda>(_1, _2, _3, _4)|
## +------------------------+
## |                 a+b+c+d|
## +------------------------+

避免使用array包装输入。

您还可以使用struct作为替代包装器来访问列名：

h = udf(lambda row: "+".join(row.asDict().keys()))

df.select(h(struct("_1", "_2", "_3"))).show()
## +----------------------------+
## |<lambda>(struct(_1, _2, _3))|
## +----------------------------+
## |                    _1+_3+_2|
## +----------------------------+

Spark UDF，包含未知数量的列

1 个答案: