根据 pyspark 中的列值过滤分组的 dataframe

Question

我有低于df。 我想按公司和日期对 dataframe 进行分组，并且对于此类分组的子集，如果可用，则根据优先级 QF 的类别过滤行，如果不是 SAF，如果不是，则为 AF。 我正在尝试使用 window function 分配等级，但也许有更简单的方法。

    company     date     value  category
    ------------------------------------
      xyz    31-12-2020    12      AF
      xyz    31-12-2020    10      SAF
      xyz    31-12-2020    11      QF
      xyz    30-06-2020    14      AF
      xyz    30-06-2020    16      SAF
      xyz    30-09-2020    13      SAF
      xyz    31-03-2019    20      AF

预期 output：

   company     date      value  category
    ------------------------------------
      xyz    31-12-2020    11      QF
      xyz    30-06-2020    16      SAF
      xyz    30-09-2020    13      SAF
      xyz    31-03-2019    20      AF

Answer 1

我们可以使用when().otherwise()为类别分配排名，并保留组中排名最低的记录。

data_sdf. \
    withColumn('cat_rank',
               func.when(func.col('cat') == 'QF', func.lit(1)).
               when(func.col('cat') == 'SAF', func.lit(2)).
               when(func.col('cat') == 'AF', func.lit(3))
               ). \
    withColumn('min_cat_rank', 
               func.min('cat_rank').over(wd.partitionBy('company', 'dt'))
               ). \
    filter(func.col('min_cat_rank').isNotNull()). \
    filter(func.col('min_cat_rank') == func.col('cat_rank')). \
    drop('cat_rank', 'min_cat_rank'). \
    show()

# +-------+----------+---+---+
# |company|        dt|val|cat|
# +-------+----------+---+---+
# |    xyz|30-09-2020| 13|SAF|
# |    xyz|30-06-2020| 16|SAF|
# |    xyz|31-03-2019| 20| AF|
# |    xyz|31-12-2020| 11| QF|
# +-------+----------+---+---+

Answer 2

假设只有有限数量的类别，并且每个类别没有重复的条目，我建议 map 将类别转换为整数，您可以对其进行排序。 之后，您可以简单地分区、排序并选择每个分区的第一个条目。

df = df.withColumn('mapping',
            f.when(f.col('category') == 'QF', f.lit('1')).otherwise(
            f.when(f.col('category') == 'SAF', f.lit('2')).otherwise(
            f.when(f.col('category') == 'AF', f.lit('3')).otherwise(f.lit(None)))))

w = Window.partitionBy('date').orderBy(f.col('mapping'))
df.withColumn('row', f.row_number().over(w))\
   .filter(f.col('row') == 1)\
   .drop('row', 'mapping')\
   .show()

Answer 3

假设在company和date的组合中同一类别可以有多个值，并且我们希望保留首选类别的value ，这是一个具有两个 window 函数的解决方案：

import pyspark.sql.functions as F
from pyspark.sql.window import Window

w_company_date = Window.partitionBy('company', 'date')
w_company_date_category = Window.partitionBy('company', 'date', 'category')

df = (df
  .withColumn('priority', F.when(F.col('category') == 'QF', 1)
                           .when(F.col('category') == 'SAF', 2)
                           .when(F.col('category') == 'AF', 3)
                           .otherwise(None))
  .withColumn('top_choice', F.when((F.col('priority') == F.min('priority').over(w_company_date))
                                   & (F.col('value') == F.max('value').over(w_company_date_category)), 1)
                             .otherwise(0))
  .filter(F.col('top_choice') == 1)
  .drop('priority', 'top_choice')
)

df.show()

+-------+----------+-----+--------+
|company|      date|value|category|
+-------+----------+-----+--------+
|    xyz|2020-03-31|   20|      AF|
|    xyz|2020-06-30|   16|     SAF|
|    xyz|2020-09-30|   13|     SAF|
|    xyz|2020-12-31|   11|      QF|
+-------+----------+-----+--------+