使用PyArrow后端的Pandas创建字符串列时遇ArrowCapacityError

错误原因 #

你遇到的ArrowCapacityError是PyArrow字符串数组的32位偏移量限制导致的：旧版本PyArrow默认使用32位整数存储字符串数组的字节偏移量，这使得单个字符串数组的总字节数上限为2^31 - 2（即2147483646字节，约2GB）。你的edge_id列总字节数达到35亿多字节，超过了这个阈值，因此触发错误。而NumPy后端（object类型字符串）没有这个限制，所以可以正常执行。

解决方案 #

1. 升级PyArrow并启用64位偏移量 #

PyArrow 15.0.0及以上版本支持64位偏移量的字符串数组，可以彻底突破2GB限制。

• 先升级PyArrow：

  pip install --upgrade pyarrow
  

• 然后在代码中启用64位偏移量配置：

  import pyarrow as pa
  # 全局配置64位偏移量
  pa.set_default_string_offset_width(64)
  
  # 之后再执行你的DataFrame操作
  df["edge_id"] = df.knot.shift(1).fillna("start") + "-" + df.knot
  

2. 分块处理DataFrame #

如果暂时无法升级PyArrow，可以将大DataFrame拆分成多个小批次处理，确保每个批次生成的edge_id列总字节数不超过2GB：

chunk_size = 1_000_000  # 根据单条edge_id的平均字节数调整chunk大小
chunks = []

for i in range(0, len(df), chunk_size):
    chunk = df.iloc[i:i+chunk_size].copy()
    chunk["edge_id"] = chunk.knot.shift(1).fillna("start") + "-" + chunk.knot
    chunks.append(chunk)

# 合并所有分块
df = pd.concat(chunks, ignore_index=True)

3. 临时切换到NumPy后端处理 #

在生成edge_id时，临时将knot列转换为object类型（NumPy后端），完成后可根据需求转回PyArrow类型：

# 临时转换为object类型
df["knot_obj"] = df["knot"].astype(object)
df["edge_id"] = df.knot_obj.shift(1).fillna("start") + "-" + df.knot_obj

# 可选：将edge_id转回pyarrow字符串类型（如果总字节数仍超2GB，这一步会再次触发错误，需谨慎）
# df["edge_id"] = df["edge_id"].astype("string[pyarrow]")

# 删除临时列
df.drop("knot_obj", axis=1, inplace=True)

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