Pyspark UDF - снижение производительности при сложном типе возврата

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()]))

Это дает мне гораздо лучшую производительность.