将Spark数据框的列求和并创建另一个数据框

Question

我有一个如下数据框-

我正在尝试从中创建另一个具有2列的数据框-列名和像这样的每列中的值之和-

到目前为止，我已经尝试过（在Spark 2.2.0中），但是抛出了堆栈跟踪-

val get_count: (String => Long) = (c: String) => {
    df.groupBy("id")
      .agg(sum(c) as "s")
      .select("s")
      .collect()(0)
      .getLong(0)
}
val sqlfunc = udf(get_count)

summary = summary.withColumn("sum_of_column", sqlfunc(col("c")))

完成此任务还有其他选择吗？

Answer 1

我认为最有效的方法是进行聚合，然后构建新的数据框。这样一来，您就避免了昂贵的explode。

首先，让我们创建数据框。顺便说一句，当您提出问题时，提供执行此操作的代码总是很不错的。这样我们可以在几秒钟内重现您的问题。

val df = Seq((1, 1, 0, 0, 1), (1, 1, 5, 0, 0),
             (0, 1, 0, 6, 0), (0, 1, 0, 4, 3))
    .toDF("output_label", "ID", "C1", "C2", "C3")

然后，我们构建我们感兴趣的列的列表，聚合并计算结果。

val cols = (1 to 3).map(i => s"C$i")
val aggs = cols.map(name => sum(col(name)).as(name))
val agg_df = df.agg(aggs.head, aggs.tail :_*) // See the note below
agg_df.show
+---+---+---+
| C1| C2| C3|
+---+---+---+
|  5| 10|  4|
+---+---+---+

我们几乎有需要的东西，我们只需要收集数据并建立一个新的数据框：

val agg_row = agg_df.first
cols.map(name => name -> agg_row.getAs[Long](name))
    .toDF("column", "sum")
    .show
+------+---+
|column|sum|
+------+---+
|    C1|  5|
|    C2| 10|
|    C3|  4|
+------+---+

编辑：

NB：df.agg(aggs.head, aggs.tail :_*)似乎很奇怪。这个想法只是计算aggs中计算出的所有聚合。人们会期望更简单的东西，例如df.agg(aggs : _*)。 agg方法的签名如下：

def agg(expr: org.apache.spark.sql.Column,exprs: org.apache.spark.sql.Column*)

也许要确保至少使用一列，这就是为什么需要在aggs和aggs.head中拆分aggs.tail的原因。

Answer 2

我要做的是定义一种从所需值创建结构的方法：

  def kv (columnsToTranspose: Array[String]) = explode(array(columnsToTranspose.map {
    c => struct(lit(c).alias("k"), col(c).alias("v"))
  }: _*))

此函数接收要转置的列的列表（在您的情况下，您的最后3列）并将其转换为以列名作为键且列值作为值的结构

然后使用该方法创建一个结构并根据需要对其进行处理

df.withColumn("kv", kv(df.columns.tail.tail))
.select( $"kv.k".as("column"), $"kv.v".alias("values"))
.groupBy("column")
.agg(sum("values").as("sum"))

首先应用先前定义的函数以具有所需的列作为所述结构，然后解构该结构以在每行中具有列键和列值。 然后，您可以按列名称进行汇总并求和

输入

+------------+---+---+---+---+
|output_label| id| c1| c2| c3|
+------------+---+---+---+---+
|           1|  1|  0|  0|  1|
|           1|  1|  5|  0|  0|
|           0|  1|  0|  6|  0|
|           0|  1|  0|  4|  3|
+------------+---+---+---+---+

输出

+------+---+
|column|sum|
+------+---+
|    c1|  5|
|    c3|  4|
|    c2| 10|
+------+---+