Добавление серии в фрейм данных: проблемы с производительностью на больших фреймах данных

Question

У меня есть датафрейм с 1 000 000 строк.

Мне нужно разделить его на 2 кадра данных (обучение и проверка), используя собственный лог c. То, что я делаю, - это создание 2 пустых фреймов данных, зацикливание на оригинальном фрейме данных с помощью iterrows и добавление к каждому фрейму данных. По сути, если я вижу категорию более двух раз в проверочном фрейме данных, я больше не добавляю ее туда.

Это начинается достаточно быстро (100it / s), но через 10 минут и с увеличением размера фреймов данных оно падает до 4it / s.

def split_df(df: pd.DataFrame):
    train = pd.DataFrame(columns = df.columns)
    valid = pd.DataFrame(columns = df.columns)
    for data in tqdm(df.iterrows()):
        tmp = pd.Series(data[1])
        if (len(valid[valid['category_id'] == data[1]['category_id']]) > 2):
            train = train.append(tmp,ignore_index=True)
        else:
            train = train.append(tmp,ignore_index=True)
            valid = valid.append(tmp,ignore_index=True)
    return (train, valid)

Я не уверен, как улучшить этот

df sample:

Answer 1

Я думаю, вам нужно GroupBy.head для фильтра максимум 2 категории строк:

train = df.groupby('category_id').head(2)