Могу ли я сделать этот код более эффективным? В настоящее время требуется ~ 6 часов для запуска на ~ 1 млн записей

Question

У меня есть 2 DataFrames:

1. trips_df всего записей = 1 048 568

weather_df всего записей = 2654

Я пытаюсь вычислить и прикрепить total_precipitation для каждой поездки какколонка.Я делаю это, просматривая дату start_timestamp и end_timestamp для каждой поездки из trips_df в weather_df, и суммируя precipitation_amount в те времена, затем добавляя это значение обратно в trips_dfпод новым столбцом.

Код, используемый для этого:

def sum_precipitation(datetime1, datetime2, weather_data):

    time1_rd = datetime1.replace(minute=0, second=0)
    time2_ru = datetime2.replace(minute=0, second=0) + dt.timedelta(hours=1)

    if time1_rd in set(weather_data['start_precipitation_datetime']):

        start_idx = weather_data.start_precipitation_datetime[
            weather_data.start_precipitation_datetime==time1_rd].index[0]

        if time2_ru in set(weather_data['end_precipitation_datetime']):

            end_idx = weather_data.end_precipitation_datetime[
                weather_data.end_precipitation_datetime==time2_ru].index[0]

            precipitation_sum = weather_data.iloc[start_idx:end_idx+1, 7].sum()

        else: precipitation_sum = 0
    else: precipitation_sum = 0

    return round(precipitation_sum, 3)

def join_weather_to_trips(trips_data, weather_data):

    trips_weather_df = trips_data.copy()

    fn = lambda row : sum_precipitation(row.start_timestamp, row.end_timestamp, weather_data)
    col = trips_data.apply(fn, axis=1)
    trips_weather_df = trips_weather_df.assign(total_precipitation=col.values)

    return trips_weather_df


trip_weather_df = join_weather_to_trips(trips_df, weather_df)

Я запустил код на подмножестве из 65 записей, и это заняло ~ 1,3 с.(CPU times: user 1.27 s, sys: 8.77 ms, total: 1.28 s, Wall time: 1.28 s).Экстраполируя эту производительность на все мои данные, это заняло бы (1,3 * 1048568) / 65 = 20971,36 секунд или 5,8 часа.

Может ли кто-то с большим опытом сказать мне, если я делаю это правильно, где я мог бы ускорить этот код, или если есть какие-либо альтернативы (например, более быстрые реализации)?

Answer 1

Это может быть не самым быстрым, но вы можете попробовать:

trips_df['precipitation_amount'] = 0

for s,e,p in zip(weather_df['start_precipitation_datetime'], 
               weather_df['end_precipitation_datetime'],
               weather_df.precipitation_amount):
    masks = trips_df.start_timestamp.between(s,e) | trips_df.end_timestamp.between(s,e)
    trips_df.loc[masks, 'precipitation_amount'] += p

На моем компьютере потребовалось 10 секунд, чтобы обработать 1M рейсов и 260 погодных условий. Так что около 100 с для фактических данных.

Обновление: я пробовал 1M поездки и 2600 погодных условий, Wall time: 1min 36s

Примечание : вам может понадобиться уменьшить weather_df['end_precipitation_datetime'] на одну минуту, чтобы избежать двойного счета в случае, если поездка начинается в час.

Answer 2

Я бы предложил использовать pip install DateTimeRange

start_1 = datetime.datetime(2016, 3, 16, 20, 30)
end_1 = datetime.datetime(2016, 3, 17, 20, 30)
start_2 = datetime.datetime(2016, 3, 14, 20, 30)
end_2 = datetime.datetime(2016, 3, 17, 22, 30)
dtr1 = datetimerange.DateTimeRange(start_1, end_1)
dtr2 = datetimerange.DateTimeRange(start_2, end_2)

Затем, если вы хотите проверить, содержится ли dtr1 в dtr2:

>>> dtr1.start_datetime in dtr2
True
>>> dtr1.end_datetime in dtr2
True

Таким образом, вы экономите много «если, тогда».

Кстати, я не уверен, стоит ли вам использовать «set», почему бы и нет:

weather_data['start_precipitation_datetime'].values