Faster (vectorized) way to do this pandas formula with dates

user_14157390 0
Viewed 410

I'm building a time series, trying to get a more efficient way to do this - ideally vectorized. The pandas apply with list comprehension step is very slow (on a big data set).

import datetime
import pandas as pd

# Dummy data:
todays_date = datetime.datetime.now().date()
xdates = pd.date_range(todays_date-datetime.timedelta(10), periods=4, freq='D')
categories = list(2*'A') + list(2*'B')
d = {'xdate': xdates, 'periods': [8]*2 + [2]*2, 'interval': [3]*2 + [12]*2}
df = pd.DataFrame(d,index=categories)

# This step is slow:
df['sdates'] = df.apply(lambda x: [x.xdate + pd.DateOffset(months=k*x.interval) for k in range(x.periods)], axis=1)
# This step is quite quick, but shown here for completeness
df = df.explode('sdates')

Maybe something like this:

df['sdates'] = [df.xdate + df.periods * [df.interval.astype('timedelta64[M]')]]

but the syntax isn't quite right. This code

df = pd.DataFrame(d,index=categories)
df['m_offsets'] = df.interval.apply(lambda x: list(range(0, 72, x)))
df = df.explode('m_offsets')
df['sdate'] = df.xdate + df.m_offsets * pd.DateOffset(months=1)

I think is similar to one of the answers, but the last step, pd.DateOffset gives a warning:

PerformanceWarning: Adding/subtracting array of DateOffsets to DatetimeArray not vectorized

I tried building something along the lines of one answer, but as mentioned the modular arithmatic needs tweaking a lot to deal with edge cases, and haven't figured that out yet (calendar monthrange wasn't playing nicely). This function doesn't run:

from calendar import monthrange
def add_months(df, date_col, n_col):
    """ Adds ncol months do date_col """
    z = df.copy()
    # calculate new year/month/day and convert to datetime
    z['year'] = (z[date_col].dt.year * 12 + (z[date_col].dt.month-1) + z[n_col]) // 12
    z['month'] = ((z[date_col].dt.month + z[n_col] - 1) % 12) + 1
    x,x = monthrange(z.year, z.month)
    z['days_in_month'] = monthrange(z.year, z.month)
    z['target_day'] = z[date_col].dt.day
    # z['day'] = min(z.target_day, z.days_in_month)
    z['day'] = z.days_in_month
    z['sdates'] = pd.to_datetime(z[['year', 'month', 'day']])
    return z['sdates']

This works, for now, but the dateoffset is a really heavy step.

df = pd.DataFrame(d,index=categories)
df['m_offsets'] = df.interval.apply(lambda x: list(range(0, 72, x)))
df = df.explode('m_offsets')
df['sdates'] = df.apply(lambda x: x.xdate + pd.DateOffset(months=x.m_offsets), axis=1)
pandas python
2 Answers

Here's one option. You're adding months, so we can actually calculate new year/month/day by only dealing with integers in a vectorized way, and then create datetime from these y/m/d combinations:

def f_proposed(df):
    z = df.copy()
    z = z.reset_index()

    # repeat xdate as many times as the number of periods
    z = z.loc[np.repeat(z.index, z['periods'])]
    
    # calculate k number of months to add
    z['k'] = z.groupby(level=0).cumcount() * z['interval']
    
    # calculate new year/month/day and convert to datetime
    z['year'] = (z['xdate'].dt.year * 12 + z['xdate'].dt.month - 1 + z['k']) // 12
    z['month'] = (z['xdate'].dt.month - 1 + z['k']) % 12 + 1
    
    # clip day to days_in_month
    z['days_in_month'] = pd.to_datetime(
        z['year'].astype(str)+'-'+z['month'].astype(str)+'-01').dt.days_in_month
    z['day'] = np.clip(z['xdate'].dt.day, 0, z['days_in_month'])
    
    z['sdates'] = pd.to_datetime(z[['year', 'month', 'day']])
    
    # drop temporary columns
    z = z.set_index('index').drop(columns=['k', 'year', 'month', 'day', 'days_in_month'])
    return z

To compare performance with the original, I've generated a test dataset with 10,000 rows.

Here's my timings (~23x speedup for 10K):

%timeit f_proposed(z)
82.7 ms ± 222 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit f_original(z)
1.92 s ± 2.75 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

P.S. For 170K it takes about 1.39s with f_proposed and 33.6 s with f_original on my machine

user_67927430

Semi-vectorized way

As I say below, I don't think there is a pure vectorized way to add a variable and general DateOffset to a Series of Timestamps. @perl solution works in the case where the DateOffset is an exact multiple of 1 month.

Now, adding a single constant DateOffset is vectorized, so we can use the following. It capitalizes on the fact that there is a limited set of distinct values for the date offset. It is also relatively fast, and it is correct for any DateOffset and dates:

n = df['periods'].values
period_no = np.repeat(n - n.cumsum(), n) + np.arange(n.sum())
z = pd.DataFrame(
    np.repeat(df.reset_index().values, repeats=n, axis=0),
    columns=df.reset_index().columns,
).set_index('index')
z = z.assign(madd=period_no * z['interval'])

z['sdates'] = z['xdate']

for madd in set(z['madd'].unique()):
    z.loc[z['madd'] == madd, 'sdates'] += pd.DateOffset(months=madd)

Timing:

# modified large dummy data:
N = 170_000
todays_date = datetime.datetime.now().date()
xdates = pd.date_range(todays_date-datetime.timedelta(10), periods=N, freq='H')
categories = np.random.choice(list('ABCDE'), N)
d = {'xdate': xdates, 'periods': np.random.randint(1,10,N), 'interval': np.random.randint(1,12,N)}
df = pd.DataFrame(d,index=categories)

%%time (the above)
CPU times: user 3.49 s, sys: 13.5 ms, total: 3.51 s
Wall time: 3.51 s

(Note: for 10K rows using the generation above, I see times of ~240ms, but of course it is dependent on how many distinct month offsets you have in your data).

Example result (for one draw of 170K rows as per above):

>>> z.tail()
                    xdate periods interval madd              sdates
index                                                              
B     2040-08-25 06:00:00       8        8   48 2044-08-25 06:00:00
B     2040-08-25 06:00:00       8        8   56 2045-04-25 06:00:00
D     2040-08-25 07:00:00       3        2    0 2040-08-25 07:00:00
D     2040-08-25 07:00:00       3        2    2 2040-10-25 07:00:00
D     2040-08-25 07:00:00       3        2    4 2040-12-25 07:00:00

Correction on the initial answer

I stand corrected: my original answer is not vectorized either. The first part, exploding the DataFrame and building the number of months to add, is vectorized and very fast. But the second part, adding a DateOffset of a variable number of months, is not.

I hope I am wrong, but I don't think there is currently a way to do that second part in a vectorized way.

Direct date-parts manipulation (e.g. month = (month - 1 + n_months) % 12 + 1, etc.) are bound to fail for corner cases (e.g. '2021-02-31'). Short of replicating the logic used in DateOffset, this is not going to work for certain cases.

Initial answer

Here is a vectorized way:

n = df.periods.values
period_no = np.repeat(n - n.cumsum(), n) + np.arange(n.sum())
z = pd.DataFrame(
    np.repeat(df.reset_index().values, repeats=n, axis=0),
    columns=df.reset_index().columns,
).set_index('index').assign(period_no=period_no)

z['sdates'] = z['period_no']  * z['interval'] * pd.DateOffset(months=1) + z['xdate']
user_7581740
Related