在我的Scala序列中添加一个计数列 [英] Adding a count column to my sequence in Scala

问题描述

给出下面的代码，我将如何添加一个计数列? (例如.count("*").as("count"))

Given the code below, how would I go about adding a count column? (e.g. .count("*").as("count"))

最终输出看起来像这样:

Final output to look like something like this:

+---+------+------+-----------------------------+------
| id|sum(d)|max(b)|concat_ws(,, collect_list(s))|count|
+---+------+------+-----------------------------+------
|  1|   1.0|  true|                          a. | 1 |
|  2|   4.0|  true|                          b,b| 2 |
|  3|   3.0|  true|                          c. | 1 |

当前代码如下:

val df =Seq(
  (1, 1.0, true, "a"),
  (2, 2.0, false, "b")
  (3, 3.0, false, "b")
  (2, 2.0, false, "c")
).toDF("id","d","b","s")

val dataTypes: Map[String, DataType] = df.schema.map(sf => (sf.name,sf.dataType)).toMap

def genericAgg(c:String) = {
  dataTypes(c) match {
    case DoubleType => sum(col(c))
    case StringType => concat_ws(",",collect_list(col(c))) // "append"
    case BooleanType => max(col(c))
  }
}

val aggExprs: Seq[Column] = df.columns.filterNot(_=="id")
.map(c => genericAgg(c))

df
.groupBy("id")
.agg(aggExprs.head,aggExprs.tail:_*)
  .show()

推荐答案

您只需在agg中的aggExprs.tail后面附加count("*").as("count")，如下所示:

You can simply append count("*").as("count") to aggExprs.tail in your agg, as shown below:

df.
  groupBy("id").agg(aggExprs.head, aggExprs.tail :+ count("*").as("count"): _*).
  show
// +---+------+------+-----------------------------+-----+
// | id|sum(d)|max(b)|concat_ws(,, collect_list(s))|count|
// +---+------+------+-----------------------------+-----+
// |  1|   1.0|  true|                            a|    1|
// |  3|   3.0| false|                            b|    1|
// |  2|   4.0| false|                          b,c|    2|
// +---+------+------+-----------------------------+-----+

这篇关于在我的Scala序列中添加一个计数列的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！