PySpark如何仅将key为total的行转为新列并在所有行保留其值

实现方案 #

方法1：先提取total值再追加列（小数据量场景） #

先提取total行的val值，再过滤原数据后追加常量列即可：

from pyspark.sql.functions import lit

# 提取total对应的数值
total_val = df.filter(df.key == "total").first()["val"]
# 过滤total行并新增全量填充的total列
res = df.filter(df.key != "total").withColumn("total", lit(total_val))

res.show()

输出结果：

+----+---+-----+
| key|val|total|
+----+---+-----+
|num1|  1|   10|
|num2|  5|   10|
+----+---+-----+

方法2：窗口函数实现（大数据量无driver拉取场景） #

如果数据量较大不想把值拉取到driver端，可以用全局窗口函数提取total值填充所有行，再过滤即可：

from pyspark.sql.functions import first, when
from pyspark.sql.window import Window

# 全局窗口取total对应的val值填充所有行
res = df.withColumn("total", first(when(df.key == "total", df.val), ignorenulls=True).over(Window.partitionBy())) \
        .filter(df.key != "total")

res.show()

输出结果和方法1一致。

内容的提问来源于stack exchange，提问作者md2614