如何避免在Spark Scala中迭代使用withColumn？

Question

当前，我们有一个迭代使用withColumn的代码。 否则-条件检查＆在此基础上进行算术运算。

样例代码：

df.withColumn("col4", when(col("col1")>10, col("col2").+col("col3")).otherwise(col("col2")))

其他算术运算针对另外40列进行迭代。 总记录数-2M。

重申的代码：

  df.withColumn(colName1, when((col(amountToSubtract).<("0")) && (col(colName1).===("0")) && (col(amountToSubtract).<=(col(colName2))) && (col(colName2).!==("0")), col(colName2)).otherwise(col(colName1))).
      withColumn(amountToSubtract, when(col(amountToSubtract).!==("0"), col(amountToSubtract).-(col(colName2))).otherwise(col(amountToSubtract))).
      withColumn(colName1, when((col(amountToSubtract).>("0")) && (col(colName1).===("0")), col(colName2).+(col(amountToSubtract))).otherwise(col(colName1))).
      withColumn(amountToSubtract, when((col(amountToSubtract).>("0")) && (col(colName1).===("0")), "0").otherwise(col(amountToSubtract)))

接下来是其他7或8组计算。

此时，该作业将挂起更长的时间。 有时会引发GC开销错误。 意识到.withColumn的迭代使用不会提高性能，因此我找不到实现上述条件检查的替代方法。 请协助。

Answer 1

您可以尝试将当前使用'when / otherwise'实现的所有逻辑封装到一个UDF中，该UDF将生成新列值时需要考虑的所有列值的数组作为输入，并作为输出返回一个所有生成的列值。 UDF有时会遇到自己的性能问题，但这可能值得一试。 这是我正在考虑的技术的简单说明：

object SO extends App {

  val sparkSession = SparkSession.builder().appName("simple").master("local[*]").getOrCreate()
  sparkSession.sparkContext.setLogLevel("ERROR")

  import sparkSession.implicits._

  case class Record(col1: Int, col2: Int, amtToSubtract: Int)

  val recs = Seq(
    Record(1, 2, 3),
    Record(11, 2, 3)
  ).toDS()


  val colGenerator : Seq[Int] => Seq[Int] =
    (arr: Seq[Int]) =>  {
      val (in_c1, in_c2, in_amt_sub) = (arr(0), arr(1), arr(2))

      val newColName1_a = if (in_amt_sub < 0 && in_c1 == 0 && in_amt_sub < in_c2 &&  in_c2 != 0) {
        in_c2
      }
      else {
        in_c1
      }
      val newAmtSub_a = if (in_amt_sub != 0) {
        in_amt_sub - in_c2
      } else {
        in_amt_sub
      }

      val newColName1_b  = if (  newAmtSub_a  > 0 &&  newColName1_a  == 0 ) {
        in_c2  + newAmtSub_a
      } else {
        newColName1_a
      }

      val newAmtSub_b = if (newAmtSub_a  > 0  && newColName1_b   == 0) {
        0
      } else {
        newAmtSub_a
      }

      Seq(newColName1_b,  newAmtSub_b)
    }

  val colGeneratorUdf = udf(colGenerator)

  // Here the first column in the generated array is 'col4', the UDF could equivalently generate as many
  // other values as you want from the input array of column values.
  //
  val afterUdf = recs.withColumn("colsInStruct",   colGeneratorUdf (array($"col1", $"col2", $"amtToSubtract")))
  afterUdf.show()
  // RESULT
  //+----+----+-------------+------------+
  //|col1|col2|amtToSubtract|colsInStruct|
  //+----+----+-------------+------------+
  //|   1|   2|            3|      [1, 1]|
  //|  11|   2|            3|     [11, 1]|
  //+----+----+-------------+------------+


}

Answer 2

我不确定您要应用的逻辑。 但是，这里的foldleft的想法是：

  val colList = List("col1", "col2", "col3")
  val df: DataFrame = ???
  colList.foldLeft(df){case(df, colName1) => df
     .withColumn(colName1, when((col(amountToSubtract).<("0")) && (col(colName1).=== ("0")) && (col(amountToSubtract).<=(col(colName2))) && (col(colName2).!==("0")), col(colName2)).otherwise(col(colName1))).
     .withColumn(amountToSubtract, when(col(amountToSubtract).!==("0"), col(amountToSubtract).-(col(colName2))).otherwise(col(amountToSubtract))).
     .withColumn(colName1, when((col(amountToSubtract).>("0")) && (col(colName1).===("0")), col(colName2).+(col(amountToSubtract))).otherwise(col(colName1))).
     .withColumn(amountToSubtract, when((col(amountToSubtract).>("0")) && (col(colName1).===("0")), "0").otherwise(col(amountToSubtract)))

}