如何仅使用PyArrow根据映射替换数组中的所有值？

我有如下PyArrow数组：

import pyarrow as pa

arr = pa.array([1, 3, 2, 2, 1, 3])

希望根据映射{1: 'one', 2: 'two', 3: 'three'}替换数组中的值，最终得到：

<pyarrow.lib.LargeStringArray object at 0x7f8dd0b3c820>
[
  "one",
  "three",
  "two",
  "two",
  "one",
  "three"
]

我可以通过Polars中转实现该需求：

import polars as pl

pl.from_arrow(arr).replace_strict({1: 'one', 2: 'two', 3: 'three'}, return_dtype=pl.String).to_arrow()

输出结果：

<pyarrow.lib.LargeStringArray object at 0x7f8dd0b3c820>
[
  "one",
  "three",
  "two",
  "two",
  "one",
  "three"
]

请问是否可以仅使用PyArrow完成这个操作？

当然可以直接用PyArrow实现，以下是几种实用方法：

方法1：使用pyarrow.compute.replace_with_mask #

通过创建掩码数组逐个替换对应值：

import pyarrow as pa
import pyarrow.compute as pc

arr = pa.array([1, 3, 2, 2, 1, 3])
mapping = {1: 'one', 2: 'two', 3: 'three'}

# 初始化字符串数组作为结果容器
result = pa.array([''] * len(arr), type=pa.string())

for num, text in mapping.items():
    # 生成当前数值对应的掩码
    mask = pc.equal(arr, num)
    # 用掩码定位并替换值
    result = pc.replace_with_mask(result, mask, text)

print(result)

输出：

<pyarrow.lib.StringArray object at 0x...>
[
  "one",
  "three",
  "two",
  "two",
  "one",
  "three"
]

方法2：借助pyarrow.DictionaryArray转换 #

利用字典数组的键值映射特性完成转换：

import pyarrow as pa

arr = pa.array([1, 3, 2, 2, 1, 3])
mapping = {1: 'one', 2: 'two', 3: 'three'}

# 创建字典类型，绑定原数值与目标字符串的映射关系
dict_type = pa.dictionary(pa.int64(), pa.string())
dict_arr = pa.DictionaryArray.from_arrays(arr, pa.array(list(mapping.values()), type=pa.string()))

# 将字典数组转为普通字符串数组
result = dict_arr.cast(pa.string())
print(result)

输出与目标结果一致。

方法3：使用pyarrow.compute.map_lookup（PyArrow 15.0+） #

如果你的PyArrow版本≥15.0，可直接用更简洁的映射查找方法：

import pyarrow as pa
import pyarrow.compute as pc

arr = pa.array([1, 3, 2, 2, 1, 3])
mapping = {1: 'one', 2: 'two', 3: 'three'}

# 将Python字典转为PyArrow的Map类型常量
map_const = pa.array([mapping], type=pa.map_(pa.int64(), pa.string()))
# 批量查找映射值
result = pc.map_lookup(map_const[0], arr)

print(result)

该方法无需循环，一步完成替换操作。

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