Pandas groupby применять медленные действия

Я работаю над программой, которая включает в себя большие объемы данных. Я использую модуль Python Pandas для поиска ошибок в моих данных. Это обычно работает очень быстро. Однако этот фрагмент кода, который я написал, кажется намного медленнее, чем должен быть, и я ищу способ ускорить его.

Чтобы вы, ребята, правильно его протестировали, я загрузил довольно большой кусок кода. Вы должны быть в состоянии запустить его как есть. Комментарии в коде должны объяснить, что я пытаюсь сделать здесь. Любая помощь будет принята с благодарностью.

# -*- coding: utf-8 -*-

import pandas as pd
import numpy as np

# Filling dataframe with data
# Just ignore this part for now, real data comes from csv files, this is an example of how it looks
TimeOfDay_options = ['Day','Evening','Night']
TypeOfCargo_options = ['Goods','Passengers']
np.random.seed(1234)
n = 10000

df = pd.DataFrame()
df['ID_number'] = np.random.randint(3, size=n)
df['TimeOfDay'] = np.random.choice(TimeOfDay_options, size=n)
df['TypeOfCargo'] = np.random.choice(TypeOfCargo_options, size=n)
df['TrackStart'] = np.random.randint(400, size=n) * 900
df['SectionStart'] = np.nan
df['SectionStop'] = np.nan

grouped_df = df.groupby(['ID_number','TimeOfDay','TypeOfCargo','TrackStart'])
for index, group in grouped_df:
    if len(group) == 1:
        df.loc[group.index,['SectionStart']] = group['TrackStart']
        df.loc[group.index,['SectionStop']] = group['TrackStart'] + 899

    if len(group) > 1:
        track_start = group.loc[group.index[0],'TrackStart']
        track_end = track_start + 899
        section_stops = np.random.randint(track_start, track_end, size=len(group))
        section_stops[-1] = track_end
        section_stops = np.sort(section_stops)
        section_starts = np.insert(section_stops, 0, track_start)

        for i,start,stop in zip(group.index,section_starts,section_stops):
            df.loc[i,['SectionStart']] = start
            df.loc[i,['SectionStop']] = stop

#%% This is what a random group looks like without errors
#Note that each section neatly starts where the previous section ended
#There are no gaps (The whole track is defined)
grouped_df.get_group((2, 'Night', 'Passengers', 323100))

#%% Introducing errors to the data
df.loc[2640,'SectionStart'] += 100
df.loc[5390,'SectionStart'] += 7

#%% This is what the same group looks like after introducing errors 
#Note that the 'SectionStop' of row 1525 is no longer similar to the 'SectionStart' of row 2640
#This track now has a gap of 100, it is not completely defined from start to end
grouped_df.get_group((2, 'Night', 'Passengers', 323100))

#%% Try to locate the errors
#This is the part of the code I need to speed up

def Full_coverage(group):
    if len(group) > 1:
        #Sort the grouped data by column 'SectionStart' from low to high

        #Updated for newer pandas version
        #group.sort('SectionStart', ascending=True, inplace=True)
        group.sort_values('SectionStart', ascending=True, inplace=True)

        #Some initial values, overwritten at the end of each loop  
        #These variables correspond to the first row of the group
        start_km = group.iloc[0,4]
        end_km = group.iloc[0,5]
        end_km_index = group.index[0]

        #Loop through all the rows in the group
        #index is the index of the row
        #i is the 'SectionStart' of the row
        #j is the 'SectionStop' of the row
        #The loop starts from the 2nd row in the group
        for index, (i, j) in group.iloc[1:,[4,5]].iterrows():

            #The start of the next row must be equal to the end of the previous row in the group
            if i != end_km: 

                #Add the faulty data to the error list
                incomplete_coverage.append(('Expected startpoint: '+str(end_km)+' (row '+str(end_km_index)+')', \
                                    'Found startpoint: '+str(i)+' (row '+str(index)+')'))                

            #Overwrite these values for the next loop
            start_km = i
            end_km = j
            end_km_index = index

    return group

#Check if the complete track is completely defined (from start to end) for each combination of:
    #'ID_number','TimeOfDay','TypeOfCargo','TrackStart'
incomplete_coverage = [] #Create empty list for storing the error messages
df_grouped = df.groupby(['ID_number','TimeOfDay','TypeOfCargo','TrackStart']).apply(lambda x: Full_coverage(x))

#Print the error list
print('\nFound incomplete coverage in the following rows:')
for i,j in incomplete_coverage:
    print(i)
    print(j)
    print() 

#%%Time the procedure -- It is very slow, taking about 6.6 seconds on my pc
%timeit df.groupby(['ID_number','TimeOfDay','TypeOfCargo','TrackStart']).apply(lambda x: Full_coverage(x))

Ответ 1

Проблема, я считаю, в том, что ваши данные имеют 5300 различных групп. В связи с этим, что-то медленное внутри вашей функции будет увеличено. Вероятно, для экономии времени вы могли бы использовать векторизованную операцию, а не цикл for, но гораздо более простой способ сэкономить несколько секунд - return 0, а не return group. Когда вы return group, pandas на самом деле создаст новый объект данных, объединяющий ваши отсортированные группы, которые вы не используете. Когда вы return 0, pandas будет комбинировать вместо 5300 нулей, что намного быстрее.

Например:

cols = ['ID_number','TimeOfDay','TypeOfCargo','TrackStart']
groups = df.groupby(cols)
print(len(groups))
# 5353

%timeit df.groupby(cols).apply(lambda group: group)
# 1 loops, best of 3: 2.41 s per loop

%timeit df.groupby(cols).apply(lambda group: 0)
# 10 loops, best of 3: 64.3 ms per loop

Простое объединение результатов, которые вы не используете, занимает около 2,4 секунд; остальное время - это фактическое вычисление в вашем цикле, которое вы должны попытаться провести в векторе.


Edit:

С быстрой дополнительной векторизованной проверкой перед циклом for и возвратом 0 вместо group я получил время до примерно ~ 2 секунд, что в основном связано с сортировкой каждой группы. Попробуйте эту функцию:

def Full_coverage(group):
    if len(group) > 1:
        group = group.sort('SectionStart', ascending=True)

        # this condition is sufficient to find when the loop
        # will add to the list
        if np.any(group.values[1:, 4] != group.values[:-1, 5]):
            start_km = group.iloc[0,4]
            end_km = group.iloc[0,5]
            end_km_index = group.index[0]

            for index, (i, j) in group.iloc[1:,[4,5]].iterrows():
                if i != end_km:
                    incomplete_coverage.append(('Expected startpoint: '+str(end_km)+' (row '+str(end_km_index)+')', \
                                        'Found startpoint: '+str(i)+' (row '+str(index)+')'))                
                start_km = i
                end_km = j
                end_km_index = index

    return 0

cols = ['ID_number','TimeOfDay','TypeOfCargo','TrackStart']
%timeit df.groupby(cols).apply(Full_coverage)
# 1 loops, best of 3: 1.74 s per loop

Отредактируйте 2: вот пример, который включает мое предложение о перемещении сортировки вне группы и удаление ненужных циклов. Удаление данного цикла не намного быстрее для данного примера, но будет быстрее, если есть много незавершенных:

def Full_coverage_new(group):
    if len(group) > 1:
        mask = group.values[1:, 4] != group.values[:-1, 5]
        if np.any(mask):
            err = ('Expected startpoint: {0} (row {1}) '
                   'Found startpoint: {2} (row {3})')
            incomplete_coverage.extend([err.format(group.iloc[i, 5],
                                                   group.index[i],
                                                   group.iloc[i + 1, 4],
                                                   group.index[i + 1])
                                        for i in np.where(mask)[0]])
    return 0

incomplete_coverage = []
cols = ['ID_number','TimeOfDay','TypeOfCargo','TrackStart']
df_s = df.sort_values(['SectionStart','SectionStop'])
df_s.groupby(cols).apply(Full_coverage_nosort)

Ответ 2

Я обнаружил, что команды pandas locate (.loc или .iloc) также замедляют ход. Перемещая сортировку из цикла и преобразовывая данные в массивы numpy в начале функции, я получил еще более быстрый результат. Я знаю, что данные больше не являются фреймворком данных, но индексы, возвращаемые в списке, могут использоваться для поиска данных в исходном df.

Если есть способ ускорить этот процесс, я буду благодарен за помощь. Что я до сих пор:

def Full_coverage(group):

    if len(group) > 1:
        group_index = group.index.values
        group = group.values

        # this condition is sufficient to find when the loop will add to the list
        if np.any(group[1:, 4] != group[:-1, 5]):
            start_km = group[0,4]
            end_km = group[0,5]
            end_km_index = group_index[0]

            for index, (i, j) in zip(group_index, group[1:,[4,5]]):

                if i != end_km:
                    incomplete_coverage.append(('Expected startpoint: '+str(end_km)+' (row '+str(end_km_index)+')', \
                                        'Found startpoint: '+str(i)+' (row '+str(index)+')'))               
                start_km = i
                end_km = j
                end_km_index = index

    return 0

incomplete_coverage = []
df.sort(['SectionStart','SectionStop'], ascending=True, inplace=True)
cols = ['ID_number','TimeOfDay','TypeOfCargo','TrackStart']
%timeit df.groupby(cols).apply(Full_coverage)
# 1 loops, best of 3: 272 ms per loop