calculation between two date in YYYYMM format in Pyspark or python

user_6247864 0
Viewed 566

I have a table, such in below

ID  Index_month  Month_ID
1     201701      201701
1     201701      201702
1     201701      201703
1     201701      201704
1     201701      201705
1     201701      201706
2     201501      201701
2     201501      201702
2     201501      201703
2     201501      201704
2     201501      201705
2     201501      201706

and I want to calculate the length from their Index month to the max month_ID, for example, for ID 1, the length should be 201706 - 201701 which 6 months; and for ID 2, the length 201706 - 201501 which is 30 months

the desire output is

ID  length
1     6
2     30

both Index_month and Month_ID are integer, at the beginning, I just use Month_ID - Index_month, but 201706 - 201501 will get 305.

Is there any month_between function in Pyspark that can do the trick?

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

You could write a quick and dirty function to convert your strings to date time objects, such as

def datestring_to_datetime(datestring):
    return datetime.strptime(datestring, '%Y%m')

It is then possible to calculate the difference in months like this:

datestring1 = "201706"
datestring2 = "201501"
difference = (datestring_to_datetime(datestring1).year - datestring_to_datetime(datestring2).year) * 12 + (datestring_to_datetime(datestring1).month - datestring_to_datetime(datestring2).month) + 1

Which will output 30

user_148775440

You can use to_date to convert the columns to date type, and manually calculate the month difference from the year difference x 12 + month difference:

import pyspark.sql.functions as F

df2 = df.groupBy('ID','Index_month').max('Month_ID').select(
    'ID',
    (
        12 * (
            F.year(F.to_date(F.col('max(Month_ID)').cast('string'), 'yyyyMM')) - 
            F.year(F.to_date(F.col('Index_month').cast('string'), 'yyyyMM'))
        ) + 
        F.month(F.to_date(F.col('max(Month_ID)').cast('string'), 'yyyyMM')) - 
        F.month(F.to_date(F.col('Index_month').cast('string'), 'yyyyMM')) + 1
    ).alias('length')
)

df2.show()
+---+------+
| ID|length|
+---+------+
|  1|     6|
|  2|    30|
+---+------+
user_141657300
Related