Как быстро обработать большой .csv?

Question

Есть train.csv в 3млн+ строк состоящий из столбцов user_id, cart.

В зависимости от значений в двух этих столбцах необходимо сформировать третий столбец target - его значения 0 и 1.

Условие записи 1 - нахождение двух или более идентичных строк в таблице, соответственно 0 - отсутствие. Работаю в Jupyter Notebook и после использования вот такой вот нехитрой функции блокнот улетает в аут и огромное количество времени "думает":

training_set = pd.read_csv('train.csv')

def is_target(columns):
    user_id = columns[0]
    cart = columns[1]

    if len(training_set.loc[(training_set['user_id'] == user_id) & (training_set['cart'] == cart)]) > 1:
        return 1
    else:
        return 0
 
training_set['target'] = training_set[['user_id', 'cart']].apply(is_target, axis=1)

Как можно выполнить эту операцию наиболее эффективно? (моя функция вообще не может обработать такое количество данных).

Answer 1

Примерно так. Поскольку у меня нет ваших данных, я формирую свои, случайные (numpy тут нужен только для формирования случайных данных).

import pandas as pd
import numpy as np

columns = ['user_id','cart']

# создаём случайный датафрейм
n = 3_000_000
df = pd.DataFrame(np.random.randint(1, 100, size=(n, 2)), columns=columns)
df = df.sort_values(columns).reset_index(drop=True)

# формирование target
df_count = df.groupby(columns).size().reset_index(name='target')
df_count['target'] = df_count['target'] - 1
df_count.loc[df_count.target > 1, 'target'] = 1
df = df.merge(df_count, on=columns)

группируем данные по интересующим нас колонкам df.groupby(columns)
считаем количество элементов в группе .size()
записываем это число в колонку target группировки .reset_index(name='target')
преобразуем target, чтобы получилось 0 или 1 по вашей формуле (две следующие строки кода)
присоединяем колонку target из группировки к основному датафрейму df.merge(df_count, on=columns)

Работает в Google Colab практически моментально: 3 млн. строк - до секунды, 30 млн. - до 10 секунд, 100 млн. - до 30 секунд, а дальше уже в память упирается, не смог проверить на более больших фреймах.

БЛОГ НА HUSL

Как быстро обработать большой .csv?

Ответы (1 шт):