如何从PostgreSQL查询结果创建指定形状的二维NumPy数组

首先是创建测试表的SQL操作：

mydb=# CREATE TABLE yokes(id INT, name TEXT, size FLOAT);
CREATE TABLE
mydb=# INSERT INTO yokes(id, name, size) VALUES(1, 'foo', 0.5);
INSERT 0 1
mydb=# INSERT INTO yokes(id, name, size) VALUES(2, 'bar', 1.5);
INSERT 0 1
mydb=# SELECT * FROM yokes;
 id | name | size 
----+------+------
  1 | foo  |  0.5
  2 | bar  |  1.5
(2 rows)

问题分析 #

需求是将查询结果转成形状为(2,3)的NumPy数组，同时保留各列原始类型（int、str、float），但两次尝试都有问题：

1. 第一次指定dtype='i4, U8, f8'生成的是一维结构化数组，形状为(2,)，不符合二维要求。
2. 第二次转成列表再生成数组，形状符合但所有元素被转为字符串；再次指定dtype='i4, U8, f8'报错，是因为该写法仅适用于一维结构化数组，二维数组会尝试把整行转成第一个类型（int），遇到字符串就报错。

解决方案 #

方案1：使用结构化数组（推荐，保留类型且操作方便）

结构化数组虽然是一维，但能通过字段名精准访问各列，完全保留原始数据类型，是处理异构数据的标准方式：

import numpy as np
import psycopg2

connection = psycopg2.connect(dbname='mydb')
cursor = connection.cursor()
cursor.execute('SELECT id, name, size FROM yokes')
rows = cursor.fetchall()

# 定义结构化数据类型，每个字段对应一列的类型
arr = np.array(rows, dtype=[('id', 'i4'), ('name', 'U8'), ('size', 'f8')])
print(arr.shape)  # 输出：(2,)

# 按字段访问各列，保留原始类型
print(arr['id'])       # 输出：[1 2]，类型int32
print(arr['name'])     # 输出：['foo' 'bar']，类型U8
print(arr['size'])     # 输出：[0.5 1.5]，类型float64

方案2：生成二维object类型数组（满足形状要求，但性能有限）

如果必须得到(2,3)的二维形状，可以指定dtype=object，让数组元素保留原始Python类型：

import numpy as np
import psycopg2

connection = psycopg2.connect(dbname='mydb')
cursor = connection.cursor()
cursor.execute('SELECT id, name, size FROM yokes')
rows = cursor.fetchall()

arr = np.array(rows, dtype=object)
print(arr.shape)  # 输出：(2, 3)
print(arr.dtype)  # 输出：object

# 验证元素类型
print(type(arr[0][0]))  # <class 'int'>
print(type(arr[0][1]))  # <class 'str'>
print(type(arr[0][2]))  # <class 'float'>

注意：object类型数组无法发挥NumPy的数值计算优势，仅适合必须用二维结构的场景。

方案3：借助Pandas中转（简化操作）

如果允许引入Pandas库，能更便捷地处理类型转换，既可以生成结构化数组，也能生成二维object数组：

import numpy as np
import psycopg2
import pandas as pd

connection = psycopg2.connect(dbname='mydb')
# 直接读取SQL结果为DataFrame
df = pd.read_sql('SELECT id, name, size FROM yokes', connection)

# 转成NumPy结构化数组
structured_arr = df.to_records(index=False)
print(structured_arr.shape)  # 输出：(2,)

# 转成二维object数组
object_arr = df.to_numpy(dtype=object)
print(object_arr.shape)      # 输出：(2, 3)

内容的提问来源于stack exchange，提问作者Guido Flohr