如何为PyArrow ChunkedArray指定索引位置设置值？

PyArrow的ChunkedArray是不可变对象，无法直接像NumPy那样原地执行ca[[1,4]] = [999,888]操作，需要通过创建新数组的方式实现需求，以下提供两种方案：

方案一：合并为单个数组修改（不在意分块时推荐） #

先将ChunkedArray合并为单个Array，修改指定位置后再转回ChunkedArray（也可直接使用修改后的Array）：

import pyarrow as pa

# 原ChunkedArray
ca = pa.chunked_array([[1,2,3], [4,5,6]])
# 待更新的全局索引与对应值
target_indices = [1, 4]
new_values = [999, 888]

# 合并所有chunk为单个Array
full_array = ca.combine_chunks()
# 创建新数组，替换指定索引的值
updated_full_array = full_array.setitem(target_indices, new_values)
# 转回ChunkedArray（可自定义分块，这里直接作为单个chunk）
updated_ca = pa.chunked_array([updated_full_array])

print(updated_ca)

输出结果：

<pyarrow.lib.ChunkedArray object at 0x...>
[
  [
    1,
    999,
    3,
    4,
    888,
    6
  ]
]

方案二：保留原分块结构修改 #

如果需要保持原有的分块结构，可以逐个处理每个chunk，计算全局索引对应的局部位置后修改：

import pyarrow as pa

ca = pa.chunked_array([[1,2,3], [4,5,6]])
target_indices = [1, 4]
new_values = [999, 888]

# 整理每个chunk需要更新的局部索引和对应值
chunk_update_map = {}
for global_idx, val in zip(target_indices, new_values):
    # 获取全局索引所属的chunk索引
    chunk_idx = ca.get_chunk_index(global_idx)
    # 计算chunk内的局部索引
    local_idx = global_idx - ca.chunk_offsets[chunk_idx]
    # 记录到更新映射中
    if chunk_idx not in chunk_update_map:
        chunk_update_map[chunk_idx] = ([], [])
    chunk_update_map[chunk_idx][0].append(local_idx)
    chunk_update_map[chunk_idx][1].append(val)

# 遍历所有chunk，生成更新后的chunk列表
updated_chunks = []
for i, chunk in enumerate(ca.chunks):
    if i in chunk_update_map:
        local_indices, vals = chunk_update_map[i]
        # 修改当前chunk的指定位置，生成新chunk
        updated_chunk = chunk.setitem(local_indices, vals)
        updated_chunks.append(updated_chunk)
    else:
        updated_chunks.append(chunk)

# 重新构建ChunkedArray
updated_ca = pa.chunked_array(updated_chunks)

print(updated_ca)

输出结果：

<pyarrow.lib.ChunkedArray object at 0x...>
[
  [
    1,
    999,
    3
  ],
  [
    4,
    888,
    6
  ]
]

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