[英]Adding new Columns based on aggregation on existing column in Spark DataFrame using scala
我有一个如下所示的 DataFrame。 我需要根据现有列创建一个新列。
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 中执行此操作。 但是很难找到数据框 DSL 的解决方案。 任何帮助,将不胜感激。 谢谢
spark df 具有名为withColumn的属性,您可以根据需要添加任意数量的派生列。 但是该列不会添加到现有 DF 中,而是创建一个带有添加列的新 DF。
例如向数据添加静态日期
val myFormattedData = myData.withColumn("batchdate",addBatchDate(myData("batchdate")))
val addBatchDate = udf { (BatchDate: String) => "20160101" }
要添加 col3,您可以使用 withcolumn + when/otherwise:
val df2 = df.withColumn("col3",when($"col2" > 1, 1).otherwise(0))
要添加 col4,已经提到的 groupBy/max + join 应该可以完成这项工作:
val df3 = df2.join(df.groupBy("col1").max("col2"), "col1")
要在没有连接的情况下实现这一点,您需要使用count
和max
作为窗口函数。 这需要使用Window
创建一个窗口,并告诉count
和max
t 在此窗口上运行。
from pyspark.sql import Window, functions as fn
df = sc.parallelize([
{'col1': 'a', 'col2': 1},
{'col1': 'a', 'col2': 2},
{'col1': 'b', 'col2': 1},
{'col1': 'c', 'col2': 1},
{'col1': 'd', 'col2': 1},
{'col1': 'd', 'col2': 2}
]).toDF()
col1_window = Window.partitionBy('col1')
df = df.withColumn('col3', fn.when(fn.count('col1').over(col1_window) > 1, 1).otherwise(0))
df = df.withColumn('col4', fn.max('col2').over(col1_window))
df.orderBy(['col1', 'col2']).show()
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.