如何将Numpy二维数组作为PyArrow Table单列条目存储？

问题分析 #

你代码中的核心错误是：arrays = [pa_arr, pa_arr]传入的是两个独立的Arrow数组（每个对应单个二维数组），但pa.Table.from_arrays要求每一列是一个包含所有行元素的Arrow数组，而非每行一个Arrow数组。

实现步骤 #

由于你的数组形状固定为(150,150)，使用嵌套的FixedSizeList类型是最优选择——它无需存储偏移量，空间效率更高，且能完美匹配固定二维结构，同时支持零拷贝转换。

1. 定义Arrow类型

针对(150,150)的二维数组，定义两层嵌套的固定大小列表类型：

import numpy as np
import pyarrow as pa

# 内层：对应数组的一行（150个元素）
inner_fixed_type = pa.fixed_size_list(pa.int64(), 150)
# 外层：对应整个二维数组（150行）
outer_fixed_type = pa.fixed_size_list(inner_fixed_type, 150)

注：如果你的数组是其他 dtype（如float32），需将pa.int64()替换为对应Arrow类型（如pa.float32()）。

2. 批量转换二维Numpy数组为Arrow列

将所有二维数组ravel后拼接成一个一维数组，再通过两层FixedSizeListArray构造完整列，全程零拷贝（依赖原Numpy数组的内存连续性）：

# 模拟10k个(150,150)的Numpy数组
numpy_arrays = [np.random.randint(0, 100, size=(150,150), dtype=np.int64) for _ in range(10000)]
numpy_masks = [np.random.choice([0,1], size=(150,150), dtype=np.int64) for _ in range(10000)]

# 将所有数组ravel后拼接成一维数组（零拷贝，前提是原数组为连续内存）
flat_data = np.concatenate([arr.ravel(order="C") for arr in numpy_arrays])
flat_masks = np.concatenate([mask.ravel(order="C") for mask in numpy_masks])

# 构造内层固定大小列表（对应每一行）
inner_data_array = pa.FixedSizeListArray.from_arrays(flat_data, 150)
inner_mask_array = pa.FixedSizeListArray.from_arrays(flat_masks, 150)

# 构造外层固定大小列表（对应整个二维数组，作为Table的列）
data_column = pa.FixedSizeListArray.from_arrays(inner_data_array, 150)
mask_column = pa.FixedSizeListArray.from_arrays(inner_mask_array, 150)

3. 构造完整PyArrow Table

将标量列与数组列组合成Table：

# 构造其他标量列
indices = pa.array(range(10000), type=pa.int64())
names = pa.array([f"sample_{i}" for i in range(10000)], type=pa.string())

# 创建Table
table = pa.Table.from_arrays(
    [indices, data_column, mask_column, names],
    names=["index", "data", "mask", "name"]
)

# 验证转换正确性（可选）
sample_data = table["data"][0].as_py()
assert np.array_equal(np.array(sample_data), numpy_arrays[0])

4. 写入Delta Lake

使用deltalake库写入Delta表：

from deltalake import write_deltalake

# 写入Delta表（支持 overwrite、append 等模式）
write_deltalake(
    path="./my_delta_table",
    data=table,
    mode="overwrite",
    overwrite_schema=True
)

关键注意事项 #

• 确保Numpy数组是连续内存（默认np.array创建的是C顺序连续数组），否则ravel可能产生拷贝，但大部分场景下无需额外处理。
• 如果需要支持可变形状的二维数组，可将嵌套FixedSizeList替换为嵌套List类型，但会牺牲部分性能和空间效率。

内容的提问来源于stack exchange，提问作者Matěj Frič