more efficient way to get monthly counts in Python/Pyspark

user_6247864 0
Viewed 131

I have a table DF look like below

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

and I want to get a count as below:

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

my code is really simple and straighforward, I just applied filter for each month, and get counts, then I copied and paste into Excel, such as below:

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() ....

which is really not efficient.

I wonder if there is any easier way to do it? and create a table just like that.

Thanks at advanced!

pyspark apache-spark-sql apache-spark sql python
2 Answers

You need to first divide the number of days by 30 and then group the values:

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))

Output

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

You can make an extra month column to capture the month of each ID depends on the days, the aggregate by that month to get the count for each month

(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|
# +-----+-----------------+
user_34415100
Related