web-dev-qa-db-ja.com

Spark / Scalaは、複数の列で同じ関数を使用してwithColumn（）を繰り返し呼び出しました

現在、.withColumnの複数のチェーンを介して複数のDataFrame列に同じプロシージャを繰り返し適用するコードがあり、プロシージャを合理化する関数を作成したいと考えています。私の場合、キーで集計された列の累積合計を見つけています：

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
}
scalaapache-sparkdataframeapache-spark-sqluser-defined-functions
15
Damian Satterthwaite-Phillips

_*_を含む可変引数でselectを使用できます。

_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)))
_
26
user6910411

質問は少し古いですが、アキュムレータとしてDataFrameを使用して列のリストを折り畳み、DataFrameをマッピングすると、列数が自明でない場合の異なるパフォーマンス結果（詳細な説明については here を参照）。短い話...いくつかの列ではfoldLeftで問題ありませんが、そうでない場合はmapの方が優れています。

5
Lorenzo