在 Python/Pyspark 中获取每月计数的更有效方法

Question

我有一个表格 DF，如下所示

ID   Days
1    30
2    55
3    32
4    12
5    100
.....

我想得到如下计数：

month                           count
30 days and greater             20,000
60 days and greater             15,323
90 days and greater             11,232
.....
3600 days and greater           55

我的代码非常简单明了，我只是为每个月应用过滤器，并获取计数，然后复制并粘贴到 Excel 中，如下所示：

month1 = df.filter("Days >= 30").agg(countDistinct('ID')).show() 
month2 = df.filter("Days>= 60").agg(countDistinct('ID')).show() 
month3 = df.filter("Days >= 90").agg(countDistinct('ID')).show() ....

确实效率不高。

我想知道是否有更简单的方法可以做到这一点？并像这样创建一个表。

感谢高级！

Answer 1

您可以额外制作一个month 列来捕获每个ID 的月份取决于天数，按该月聚合以获得每个月的计数

(df
    .withColumn('month', F
        .when(F.col('days') >= 90, 3)
        .when(F.col('days') >= 60, 2)
        .when(F.col('days') >= 30, 1)
        .otherwise(0)
    )
    .groupBy('month')
    .agg(F.countDistinct('ID').alias('count_distinct_id'))
    .show()
)

# +-----+-----------------+
# |month|count_distinct_id|
# +-----+-----------------+
# |    1|                3|
# |    3|                1|
# |    0|                1|
# +-----+-----------------+

Answer 2

您需要先将天数除以 30，然后对值进行分组：

import pyspark.sql.functions as f
from pyspark.shell import spark
from pyspark.sql import Row
from pyspark.sql.window import Window

df = spark.createDataFrame([
    Row(ID=1, Days=30),
    Row(ID=2, Days=55),
    Row(ID=3, Days=32),
    Row(ID=4, Days=12),
    Row(ID=5, Days=100),
    Row(ID=6, Days=3600)
])

# Calculating quantity of months
df = df.withColumn('total_months', f.floor(f.col('Days') / f.lit(30)))

# Group and count distinct
df = df.groupBy('total_months').agg(f.countDistinct('ID').alias('count'))

# Adding description
df = df.withColumn('month', f.concat(f.col('total_months') * f.lit(30), f.lit(' days and greater')))

# Cumulative sum
window = Window.orderBy(f.col('total_months').desc()).rangeBetween(Window.unboundedPreceding, Window.currentRow)
df = df.withColumn('count', f.sum('count').over(window))

# Selecting only required columns and sorting asc
(df
 .select('month', 'count')
 .sort('total_months')
 .show(truncate=False))

输出

+---------------------+-----+
|month                |count|
+---------------------+-----+
|0 days and greater   |6    |
|30 days and greater  |5    |
|90 days and greater  |2    |
|3600 days and greater|1    |
+---------------------+-----+