如何创建与Numpy数组共享内存的Polars数组序列（偏移不可预测）

问题 #

我正在处理由约1000个样本数组构成的「记录」，这类记录有数百万条，内存占用极高，稍不注意就会超出PC可用内存。我希望使用Polars DataFrame，以数组类型列存储这些记录，且该列需与已有的Numpy数组共享内存——该Numpy数组中的记录分布在各种不可预测（非跨步）的偏移位置。

以下示例展示了如何创建与source_data共享内存的Python记录列表，但尝试创建同样共享内存的Polars Series时失败。请问如何创建与source_data共享内存的Polars数组序列？

import numpy as np
import polars as pl

record_len = 1000 # 记录长度
source_data = np.arange(1000000, dtype=np.int64) # 所有记录都存储在这个数组中，偏移位置随机
inds0 = np.arange(0,len(source_data), 10000) # 临时变量用于计算索引
inds = inds0+np.random.randint(0,1000, len(inds0)) # 模拟记录偏移的不可预测性
# 创建与source_data共享内存的Python记录列表
records = [source_data[ind:ind+record_len] for ind in inds]
# 验证所有记录都与source_data共享内存
assert all(np.shares_memory(records[i], source_data) for i in range(len(inds)))

# 尝试从列表创建Polars序列，目标是同样与source_data共享内存
series = pl.Series("record", records, dtype=pl.Array(pl.Int64, record_len))

# 尝试验证内存共享（可能方法不对，因为Polars底层用PyArrow）
# 测试失败
assert all(np.shares_memory(series[i].to_numpy(), source_data) for i in range(len(inds)))

解决方案 #

直接传入numpy切片列表创建Polars Array列会触发内存拷贝，因为Polars的Array类型基于PyArrow，而PyArrow的FixedSizeListArray要求内存连续且有固定步长。要实现零拷贝共享内存，需要直接基于source_data的内存创建PyArrow数组，再转换为Polars Series：

实现步骤 #

• 生成所有记录的完整索引集合：将每个记录的起始索引扩展为对应的索引范围，得到指向source_data中目标元素的一维数组
• 创建PyArrow Buffer：通过pyarrow.py_buffer直接包装source_data的内存，避免拷贝
• 构建FixedSizeListArray：用上述buffer创建固定长度的列表数组，每个元素长度为record_len
• 转换为Polars Series：将PyArrow数组转为Polars Series，实现内存共享

完整代码 #

import numpy as np
import polars as pl
import pyarrow as pa

record_len = 1000
source_data = np.arange(1000000, dtype=np.int64)
inds0 = np.arange(0, len(source_data), 10000)
inds = inds0 + np.random.randint(0, 1000, len(inds0))

# 生成所有记录的完整索引
flat_indices = np.concatenate([np.arange(ind, ind + record_len) for ind in inds])
# 确保索引不越界（实际场景建议保留该检查）
assert flat_indices.max() < len(source_data)

# 创建指向source_data内存的PyArrow Buffer
buffer = pa.py_buffer(source_data)

# 构建PyArrow FixedSizeListArray
int_array = pa.Int64Array.from_buffers(
    type=pa.int64(),
    length=len(flat_indices),
    buffers=[None, buffer],
    offset=flat_indices[0] * source_data.dtype.itemsize  # 偏移量：起始索引×每个元素字节数
)
fixed_size_list = pa.FixedSizeListArray.from_arrays(int_array, record_len)

# 转换为Polars Series
series = pl.from_arrow(fixed_size_list).rename("record")

# 验证内存共享：修改源数组，检查Polars中值是否同步变化
source_data[inds[0]] = 999999
assert series[0].to_numpy()[0] == 999999
print("内存共享验证成功")

关键说明 #

• 核心是直接基于source_data的内存构建PyArrow数组，全程无数据拷贝
• 偏移量需按字节计算：int64类型每个元素占8字节，所以偏移量为起始索引×8
• 验证方式改为修改源数组后检查Polars值的变化，避免to_numpy()生成新数组导致的误判

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