使用 scala 基于 Spark DataFrame 中现有列的聚合添加新列

Question

我有一个如下所示的数据框。我需要根据现有列创建一个新列。

col1 col2
a      1
a      2
b      1
c      1
d      1
d      2

输出数据帧看起来像这样

col1  col2 col3 col4
a      1   1      2
a      2   1      2
b      1   0      1
c      1   0      1
d      1   1      2
d      2   1      2

我用来查找 col3 的逻辑是如果 col1 > 1且col4的计数是 col2 的最大值。

我熟悉如何在 sql 中执行此操作。但很难找到 Dataframe DSL 的解决方案。任何帮助，将不胜感激。谢谢

Answer 1

groupBy col1 和聚合以获得count和max。然后您可以将其与原始数据框连接起来以获得您想要的结果

val df2 = df1.groupBy("col1").agg(count() as col3, max("col2") as col4) 

val df3 = df1.join(df2, "col1")