Преобразование DataFrame Pandas в SQL-оператор VALUES

Question

Используя pandas в python, мне нужно иметь возможность генерировать эффективные запросы из кадра данных в postgresql.К сожалению, DataFrame.to_sql (...) выполняет только прямые вставки, и запрос, который я хочу сделать, довольно сложен.

В идеале я хотел бы сделать это:

WITH my_data AS (
  SELECT * FROM (
    VALUES 
    <dataframe data>
  ) AS data (col1, col2, col3)
)
UPDATE my_table 
SET
my_table.col1 = my_data.col1,
my_table.col2 = complex_function(my_table.col2, my_data.col2),
FROM my_data
WHERE my_table.col3 < my_data.col3;

Однако, чтобы сделать это, мне нужно было бы превратить мой dataframe в заявление простых значений.Конечно, я мог бы переписать свои собственные функции, но прошлый опыт научил меня, что написание функций для удаления и дезинфекции sql никогда не следует делать вручную.

Мы используем SQLAlchemy, но связанные параметры, похоже, работают только сограниченное число аргументов, и в идеале я хотел бы, чтобы сериализация данных в текст выполнялась на скорости C.

Итак, есть ли способ, с помощью панд или с помощью SQLAlchemy, эффективно повернутьмой фрейм данных в подзаголовок значений и вставить его в мой запрос?

Answer 1

Вы можете использовать psycopg2.extras.execute_values.Например, при данной настройке

CREATE TABLE my_table (
col1 int
, col2 text
, col3 int
);
INSERT INTO my_table VALUES 
(99, 'X', 1)
, (99, 'Y', 2)
, (99, 'Z', 99);

# | col1 | col2 | col3 |
# |------+------+------|
# |   99 | X    |    1 |
# |   99 | Y    |    2 |
# |   99 | Z    |   99 |

Код Python

import psycopg2
import psycopg2.extras as pge
import pandas as pd
import config

df = pd.DataFrame([
    (1, 'A', 10), 
    (2, 'B', 20),
    (3, 'C', 30)])

with psycopg2.connect(host=config.HOST, user=config.USER, password=config.PASS, database=config.USER) as conn:
    with conn.cursor() as cursor:
        sql = '''WITH my_data AS (
          SELECT * FROM (
            VALUES %s
          ) AS data (col1, col2, col3)
        )
        UPDATE my_table 
        SET
        col1 = my_data.col1,
        -- col2 = complex_function(col2, my_data.col2)
        col2 = my_table.col2 || my_data.col2
        FROM my_data
        WHERE my_table.col3 < my_data.col3'''

        pge.execute_values(cursor, sql, df.values)

обновляет my_table до

# SELECT * FROM my_table
| col1 | col2 | col3 |
|------+------+------|
|   99 | Z    |   99 |
|    1 | XA   |    1 |
|    1 | YA   |    2 |

---

В качестве альтернативы вы можете использоватьpsycopg2 до генерирует SQL.Код в format_values почти полностью скопирован из исходного кода для pge.execute_values.

import psycopg2
import psycopg2.extras as pge
import pandas as pd
import config

df = pd.DataFrame([
    (1, "A'foo'", 10), 
    (2, 'B', 20),
    (3, 'C', 30)])


def format_values(cur, sql, argslist, template=None, page_size=100):
    enc = pge._ext.encodings[cur.connection.encoding]
    if not isinstance(sql, bytes):
        sql = sql.encode(enc)
    pre, post = pge._split_sql(sql)
    result = []
    for page in pge._paginate(argslist, page_size=page_size):
        if template is None:
            template = b'(' + b','.join([b'%s'] * len(page[0])) + b')'
        parts = pre[:]
        for args in page:
            parts.append(cur.mogrify(template, args))
            parts.append(b',')
        parts[-1:] = post
        result.append(b''.join(parts))
    return b''.join(result).decode(enc)

with psycopg2.connect(host=config.HOST, user=config.USER, password=config.PASS, database=config.USER) as conn:
    with conn.cursor() as cursor:
        sql = '''WITH my_data AS (
          SELECT * FROM (
            VALUES %s
          ) AS data (col1, col2, col3)
        )
        UPDATE my_table 
        SET
        col1 = my_data.col1,
        -- col2 = complex_function(col2, my_data.col2)
        col2 = my_table.col2 || my_data.col2
        FROM my_data
        WHERE my_table.col3 < my_data.col3'''

        print(format_values(cursor, sql, df.values))

выходов

WITH my_data AS (
          SELECT * FROM (
            VALUES (1,'A''foo''',10),(2,'B',20),(3,'C',30)
          ) AS data (col1, col2, col3)
        )
        UPDATE my_table 
        SET
        col1 = my_data.col1,
        -- col2 = complex_function(col2, my_data.col2)
        col2 = my_table.col2 || my_data.col2
        FROM my_data
        WHERE my_table.col3 < my_data.col3