Question

我有一个PySpark UDF，它返回String的元组，我将其编码为结构。这是一个玩具示例，

def my_func(x):
  return "1", x, "3"

spark.udf.register("my_func", lambda x: my_func(x), StructType([StructField("one", StringType(),
                                                                StructField("two", StringType(), 
                                                                StructField("three", StringType()])

我叫

spark.sql("select col1, my_func(col1) from sdf").show()

与返回元组的一个元素相比，返回整个元组的性能下降了10到20倍。

spark.udf.register("my_func", lambda x: my_func(x)[1], StringType())

这是一个已知问题，有没有办法避免转换速度变慢？

Answer 1

这是我如何使其工作的方法-如果有更有效的方法，请使用lmk。为了解决性能问题，

1) Transform the DataFrame to an RDD[Row]

2) Apply the function to transform into a Row of the final output

3) Convert back to a DataFrame

代码：

def map_to_new_row(row):
  NewRow = Row("one", "two", "three")
  return NewRow("1", row.col1, "3")

rdd1 = df1.rdd.map(map_to_new_row)
df2 = spark.createDataFrame(rdd1, StructType([StructField("one", StringType(), 
                                              StructField("two", StringType(), 
                                              StructField("three", StringType()]))

这给了我更好的表现。

Pyspark UDF-复杂返回类型的性能下降

1 个答案: