UDF？与列？在 pyspark 中更新列哪个更好？

Question

如果我们只是实现一个简单的函数来更新 pyspark 中的列（就地），我们可以使用：

1. when

   语法，例如

  df.withColumn("col_name", when(col("reference")==1, False).otherwise(col("col_name"))

2. udf

   功能。例如

  def update_col(reference, col_name):
     if reference == 1:
        return False
     else:
        return col_name

  update_udf = udf(update_col, BooleanType())
  df.withColumn("col_name", update_udf(col("reference"), col("col_name")))

假设

df

非常大，比如十亿行。

我们将使用哪一个？有没有人体验过这两种方式并比较性能，如速度、内存成本？谢谢！

Answer 1

除了 samkart 评论中尽可能选择 Spark SQL 之外，UDF 无论语言如何对于 Spark 来说都是一个黑匣子，无法优化其使用。

此外，不要在循环或映射中使用 withColumn 来修改列，它的成本非常高。 优先选择您需要的确切列，稍后 Spark api 提供 withColumns，它也应该优于 withColumn。