根据 PySpark 中不同列的重复条目中的列值选择行

Question

我有一个 PySpark DataFrame，我将它分组在一个字段（列）上，目的是消除每个组的记录，这些记录具有另一个字段的特定值。 例如，表格看起来像

colA    colB
'a'     1
'b'     1
'a'     0
'c'     0

这里我想要删除有重复 colA 和 colB 为 0 的记录，以便获得

colA    colB
'a'     1
'b'     1
'c'     0

'c' 的行仍然存在，因为我只想删除重复（在 colA 上）行的 0。

我想不出一种方法来实现这一点，因为如果expr 不是“avg”、“max”之一，我不精通在groupBy 之后使用agg 的方法等。

Answer 1

简单的max怎么样？

from pyspark.sql.functions import max as max_

df = sc.parallelize([
  ('a', 1), ('b', 1), ('a', 0), ('c', 0)
]).toDF(('colA', 'colB'))

df.groupBy('colA').agg(max_('colB')).show()
## +----+---------+
## |colA|max(colB)|
## +----+---------+
## |   a|        1|
## |   b|        1|
## |   c|        0|
## +----+---------+

这种方法适用于任何支持排序并使用二进制标签的列，并且可以对您使用的聚合函数进行可选调整 (min / max)。

使用窗口函数可以实现更高级的规则，但成本会更高。

这里有一个例子：

from pyspark.sql.functions import col, sum as sum_, when
from pyspark.sql import Window
import sys

w = Window.partitionBy("colA").rowsBetween(-sys.maxsize, sys.maxsize)

this_non_zero = col("colB") != 0
any_non_zero = sum_(this_non_zero.cast("long")).over(w) != 0

(df
  .withColumn("this_non_zero", this_non_zero)
   .withColumn("any_non_zero", any_non_zero)
   .where(
       (col("this_non_zero") & col("any_non_zero")) | 
       ~col("any_non_zero")
   ))