Улучшение нескольких INSERT с помощью Psycopg2

Ты не знаешь, пока не узнаешь

Я обнаружил, что всякий раз, когда я спешу, чтобы что-то работало, у меня возникает плохое настроение. Привычка просто писать код и доводить дело до конца. Чаще чем нет, это приводит к беспорядочному и/или медленному коду. И тогда, когда мне понадобится чтобы сделать что-то подобное еще раз, я просто повторно использую исходный код и в итоге получаю потенциальная сложная проблема.

Одна из наиболее распространенных задач, которые я постоянно выполняю, — это взять кучу файлов и вставьте данные в базу данных Postgres. До недавнего времени мне не приходилось иметь дело с достаточно большими наборами данных, поэтому мой плохо написанный код все еще был приемлемо с точки зрения времени выполнения. Обычно это означает выполнение вставки сценарий сначала локально в тестовой базе данных, а затем в рабочей базе. я использую Python 3 и драйвер postgres Psycopg2.

Psycopg2 execute и execute_values

Исходный код выглядел примерно так:

    def insert_data(filename, date):
    
        sql = """
        INSERT INTO test (a, b, c, d)
        VALUES (%s, %s, %s, %s)
        """
    
        with open(filename) as csvfile, get_cursor() as c:
            reader = csv.reader(csvfile)
            header = next(reader)
    
            for row in reader:
                n = row[0][2:]
                values = (row[1], date, n, row[2])
                c.execute(sql, values)

Это выглядит как довольно невинная функция вставки, она принимает файл перебирает каждую строку и вставляет ее в таблицу.

Рефакторинговая функция выглядит так:

    def insert_data(filename, date):
    
        sql = """
        INSERT INTO test (a, b, c, d)
        VALUES %s
        """
    
        with open(filename) as csvfile, get_cursor() as c:
            reader = csv.reader(csvfile)
            header = next(reader)
            values_list = []
    
            for row in reader:
                n = row[0][2:]
                values = (row[1], date, n, row[2])
                values_list.append(values)
    
            execute_values(c, sql, values_list)

Разница между этими двумя функциями заключается в execute и execute_values. Каждый раз, когда вы используете execute, psycopg2 выполняет полный возврат данных из базы данных на ваш компьютер, то есть строка будет выполнена как INSERT на сервере базы данных, а затем будет возврат. Функциональность с Postgres заключается в том, что вы можете вставлять несколько строк одновременно, и это то, что execute_values делает.

Вместо вставки запроса одной строки в переработанной версии создается запрос с несколькими строками для вставки. Это сокращает количество запросов туда и обратно, существенно повышает производительность сервера базы данных и приводит к гораздо более высокой производительности.

Время выполнения

Я запустил две разные функции на небольшом подмножестве данных, состоящем из 288 478 строк. что составляет 3% файлов, которые я вставлял.

код execute:

real    0m54.761s
user    0m12.752s
sys     0m4.876s

код execute_values:

real    0m7.545s
user    0m2.092s
sys     0m0.544s

Что ж, это увеличение на 700% при небольшом количестве строк. Я не стал сравнивать время, необходимое для полного набора данных, но запуск переработанной версии занял около 25 минут, поэтому исходная версия заняла бы часы!

Урок выучен

Сэкономил бы я время, если бы перед написанием кода более подробно изучил документацию? Вероятно. Я думаю, что подобные уроки — это то, что разделяет программистов старшего и младшего уровня, тех, кто способен понять масштаб проблемы и ее решения еще до того, как они начнут писать код. Между тем, джуниорам приходится тратить время на написание плохого кода, чтобы учиться.

Обновление, еще один тест

Итак, мой друг указал мне на другой пакет Python под названием dataset, сказав, что он его использует, потому что он ленивый пользователь Python, у которого аллергия на SQL. Он также сказал, что Python > SQL, поэтому я решил доказать его неправоту, а также потому, что я не верил, что другой пакет, использующий SQLAlchemy, будет быстрее, чем просто использование Psycopg2. (SQLAlchemy построен на Psycopg2)

    db = dataset.connect('postgresql://test@127.0.0.1:5432/testdb')
    def insert_demand_data(filename, date):
        with open(filename) as csvfile:
            reader = csv.reader(csvfile)
            header = next(reader)
            values_list = []
            for row in reader:
                n = row[0][2:]
                values = dict(node=row[1], date=date, n=n, r=row[2])
                values_list.append(values)
            table.insert_many(values_list)

И что же ты знаешь.

real    0m53.392s
user    0m17.512s
sys     0m3.112s

testdb=# select count(*) from test ;
count
--------
288478
(1 row)