将Python NumPy数组插入PostgreSQL数据库的最佳方法

Question

我们的团队使用了一种非常依赖于将NumPy数据转储到文件中的软件，这会严重拖慢我们的代码。如果我们可以直接将NumPy数组存储在PostgreSQL数据库中，将会获得很大的性能提升。

其他在任何数据库或类似可搜索数据库的结构中存储NumPy数组的高效方法也是可以考虑的，但我们更喜欢PostgresSQL。

我的问题与之前提出的一个问题非常相似，但我正在寻找更强大和高效的答案，并希望存储任意的NumPy数组。

- Daniel Marchand

---

看一下这个扩展对你是否有用。 https://pgxn.org/dist/pandapost/ - Raj Verma

任何任意的NumPy数组...关系型数据库，即使是免费的、开源的企业级数据库如Postgres，也应该是计划和设计项目，所有的模式、表、字段、用户和其他组件最好都提前知道和准备好。在运行时随意转储任何数据都是不可取的，可能会导致混乱的系统。问问任何DBA！也许考虑使用.npy二进制格式、pickle或其他非文本格式？ - Parfait

你能详细说明一下混乱的系统吗？这可能是对该帖子的一个很好的回答。 - Daniel Marchand

Answer 1

不确定这是否符合您的要求，但假设您可以读/写现有的Postgres数据库：

import numpy as np
import psycopg2 as psy
import pickle

db_connect_kwargs = {
    'dbname': '<YOUR_DBNAME>',
    'user': '<YOUR_USRNAME>',
    'password': '<YOUR_PWD>',
    'host': '<HOST>',
    'port': '<PORT>'
}

connection = psy.connect(**db_connect_kwargs)
connection.set_session(autocommit=True)
cursor = connection.cursor()

cursor.execute(
    """
    DROP TABLE IF EXISTS numpy_arrays;
    CREATE TABLE numpy_arrays (
        uuid VARCHAR PRIMARY KEY,
        np_array_bytes BYTEA
    )
    """
)

这种方法的核心是将任何numpy数组（无论形状和数据类型）存储为numpy_arrays表中的一行，其中uuid是一个唯一标识符，以便以后检索该数组。实际的数组将以字节的形式保存在np_array_bytes列中。

向数据库插入数据：

some_array = np.random.rand(1500,550)
some_array_uuid = 'some_array'

cursor.execute(
    """
    INSERT INTO numpy_arrays(uuid, np_array_bytes)
    VALUES (%s, %s)
    """,
    (some_array_uuid, pickle.dumps(some_array))
)

从数据库查询：

uuid = 'some_array'
cursor.execute(
    """
    SELECT np_array_bytes
    FROM numpy_arrays
    WHERE uuid=%s
    """,
    (uuid,)
)
some_array = pickle.loads(cursor.fetchone()[0])

性能？

如果我们可以直接在 PostgreSQL 中存储 NumPy 数组，那么将获得极大的性能提升。

我还没有以任何方式对这种方法进行基准测试，因此无法确认或反驳这一点...

磁盘空间？

我猜测这种方法所需的磁盘空间与使用 np.save('some_array.npy', some_array) 将数组转储到文件中所需的空间相同。如果这是一个问题，考虑在插入前压缩字节。

- Vlad

---

拆封允许任意的Python代码执行，因此您永远不应该拆封不可信的数据。如果允许不同的进程为此列保存任意值，则存在（虽然很小）可利用的安全漏洞。为了解决这个问题，您可以使用bf = io.BytesIO(); some_array = np.load(bf, allow_pickle=False)来加载数据，而不是使用拆封（同样适用于保存）。 - Multihunter

顺便提一下，使用pickle.dump而不是np.save会额外增加（最小）37字节的开销。 - Multihunter

Answer 2

您可以使用subprocess.run()执行shell命令，通过管道从csv文件批量复制到Postgressql服务器中使用COPY。我更熟悉的是mssql，它有bcp方法，虽然无法完全测试我的解决方案，但我想象它是一种类似于通过终端调用的方法。终端命令基于第三个链接，该解决方案使用subprocess.call()，但已更新为subprocess.run()。

https://docs.python.org/3/library/subprocess.html#subprocess.run https://ieftimov.com/post/postgresql-copy/

Python psql \copy CSV to remote server

的意思是使用Python和psql将CSV文件复制到远程服务器。

import subprocess


psql_command = "\"\copy table (col1, col2) FROM file_location CSV HEADER QUOTE '\\\"' NULL ''\""
# user, hostname, password, dbname all defined elsewhere above.
command = ["psql",
    "-U", user,
    "-h", hostname,
    "-d", dbname,
    "-w", password,
    "-c", psql_command,
]

subprocess.run(command)

- hSin