如何在PyArrow中复现Polars的zip_with功能？

Polars的zip_with逻辑为：掩码为True时取第一个序列的值，掩码为False时取第二个序列的值。你之前使用pyarrow.compute.replace_with_mask时参数逻辑匹配错误，以下是正确的实现方案：

正确代码示例 #

import pyarrow as pa
import pyarrow.compute as pc

a1 = pa.array([1,2,3])
mask = pa.array([True, False, False])
a2 = pa.array([4,5,6])

# 复现Polars zip_with的逻辑：mask为True取a1，False取a2
result = pc.replace_with_mask(a2, mask, a1)
print(result)

输出结果 #

<pyarrow.lib.Int64Array object at 0x...>
[
  1,
  5,
  6
]

逻辑说明 #

pc.replace_with_mask的参数规则是：

• 第一个参数：基础数组（默认保留其值的数组）
• 第二个参数：掩码数组，True表示该位置需要被替换
• 第三个参数：替换用的数组

对应Polars的zip_with(s1, mask, s2)，我们需要：

1. 将s2设为基础数组（掩码为False时保留它的值）
2. 使用原始mask作为替换掩码（掩码为True时，把s2对应位置的值替换为s1的值）
3. 将s1设为替换数组

这样就能完全复现Polars的zip_with功能。

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