Spark / Scala повторные вызовы withColumn () с использованием одной и той же функции для нескольких столбцов

В настоящее время у меня есть код, в котором я неоднократно применяю одну и ту же процедуру к нескольким столбцам фрейма данных через несколько цепочек .withColumn, и я хочу создать функцию для оптимизации процедуры. В моем случае я нахожу кумулятивные суммы по столбцам, агрегированным по ключам:

val newDF = oldDF
  .withColumn("cumA", sum("A").over(Window.partitionBy("ID").orderBy("time")))
  .withColumn("cumB", sum("B").over(Window.partitionBy("ID").orderBy("time")))
  .withColumn("cumC", sum("C").over(Window.partitionBy("ID").orderBy("time")))
  //.withColumn(...)

то, что я хотел бы, это либо что-то вроде:

def createCumulativeColums(cols: Array[String], df: DataFrame): DataFrame = {
  // Implement the above cumulative sums, partitioning, and ordering
}

или еще лучше:

def withColumns(cols: Array[String], df: DataFrame, f: function): DataFrame = {
  // Implement a udf/arbitrary function on all the specified columns
}
8
apache-spark apache-spark-sql dataframe scala user-defined-functions

2 ответов

можно использовать select С varargs в том числе *:

import spark.implicits._

df.select($"*" +: Seq("A", "B", "C").map(c => 
  sum(c).over(Window.partitionBy("ID").orderBy("time")).alias(s"cum$c")
): _*)

это:

  • сопоставляет имена столбцов с выражениями окна с Seq("A", ...).map(...)
  • добавляет все ранее существовавшие столбцы с $"*" +: ....
  • распаковывается в сочетании с последовательности ... : _*.

и может быть обобщена в виде:

import org.apache.spark.sql.{Column, DataFrame}

/**
 * @param cols a sequence of columns to transform
 * @param df an input DataFrame
 * @param f a function to be applied on each col in cols
 */
def withColumns(cols: Seq[String], df: DataFrame, f: String => Column) =
  df.select($"*" +: cols.map(c => f(c)): _*)

если вы найдете withColumn синтаксис более читаемый вы можете использовать foldLeft:

Seq("A", "B", "C").foldLeft(df)((df, c) =>
  df.withColumn(s"cum$c",  sum(c).over(Window.partitionBy("ID").orderBy("time")))
)

, который может обобщить, например, на:

/**
 * @param cols a sequence of columns to transform
 * @param df an input DataFrame
 * @param f a function to be applied on each col in cols
 * @param name a function mapping from input to output name.
 */
def withColumns(cols: Seq[String], df: DataFrame, 
    f: String =>  Column, name: String => String = identity) =
  cols.foldLeft(df)((df, c) => df.withColumn(name(c), f(c)))
20

вопрос немного старый, но я подумал, что было бы полезно (возможно, для других) отметить, что сворачивание списка столбцов с помощью DataFrame как аккумулятор и картографии по DataFrame имеют существенно разные результаты производительности, когда количество столбцов не является тривиальным (см. здесь для полного объяснения). Длинная короткая история... для нескольких столбцов foldLeft отлично, в противном случае map - это лучше.

2