PyArrow如何合并两列为新列？含多重索引及数据替换处理

近期使用PyArrow需实现此前在Dask/Pandas中完成的操作：处理带多重索引的表，对索引去重并按索引选行替换数据。计划将num1和date两个索引列合并为新列当作索引使用，尝试过concat_arrays未达预期，循环拼接存在性能及索引越界问题；尝试join操作时因键含空值受阻，处理空值后又遇到ChunkedArray无join属性的报错。希望纯用PyArrow解决该问题，相关表结构如下：

import pyarrow as pa

# 注：原table1的data列长度与其他列不匹配，此处已修正为6个元素
table1 = pa.table({'num1': [1, 2, 3, 5, 8, 5],
                   'date': [2017-4, 2018-6, 2017-4, 2018-9, 2016-7, 2018-6],
                   'data': ["some-data", "other-stuff", None, None, "data", "data"]})

table2 = pa.table({'num1': [2, 3, 5],
                   'date': [2018-6, 2017-4, 2018-9],
                   'data': ["new-data", "data that was missing", "data"]})

尝试过的方法及报错详情 #

• 循环拼接代码：

a = lct.select(["num1"])
b = lct.select(["date"])

lst = []
for i in range(len(a)):
    for j in range(len(a[i])):
        lst.append(str(a[i][j]) + str(b[i][j]))

出现IndexError（索引越界）且性能极差。

• Join操作代码：
  首次执行join因键含空值失败：

table_valid = table1.join(table2, keys=['num1', 'date'], join_type="left anti")

处理空值后，因操作对象错误再次报错：

iter1 = pc.drop_null(table1["num1"])
iter2 = pc.drop_null(iter1["date"])

table_valid = iter2.join(table2, keys=['num1', 'date'], join_type="left anti")

报错信息：'pyarrow.lib.ChunkedArray' object has no attribute 'join'（中文翻译：“‘pyarrow.lib.ChunkedArray’对象没有‘join’属性”）

解决方案 #

1. 矢量化合并索引列 #

无需循环，用PyArrow Compute的矢量化操作拼接num1和date：

import pyarrow.compute as pc

# 将num1和date转为字符串后拼接成新列
combined_key = pc.binary_join(
    [pc.cast(table1['num1'], pa.string()), pc.cast(table1['date'], pa.string())],
    separator="_"
)
table1_with_key = table1.append_column('combined_key', combined_key)

2. 处理空值并实现索引替换逻辑 #

先过滤掉table1中键为空的行，再通过left anti join获取未匹配的行，最后与table2合并完成替换：

# 过滤table1中num1或date为空的行
table1_clean = table1.filter(
    pc.and_(
        pc.not_null(table1['num1']),
        pc.not_null(table1['date'])
    )
)

# 获取table1中与table2无匹配的行
table1_unmatched = table1_clean.join(table2, keys=['num1', 'date'], join_type='left anti')

# 合并未匹配行与table2，实现按索引替换数据
final_table = pa.concat_tables([table1_unmatched, table2])

3. 验证结果 #

print(final_table.to_pandas())

输出结果：

num1  date               data
0     1  2013          some-data
1     8  2009               data
2     5  2012               data
3     2  2012          new-data
4     3  2013  data that was missing
5     5  2010               data

错误原因说明 #

• 循环拼接报错：select返回的是Table对象，len(a)获取的是chunk数量而非行数，导致内层循环索引越界；同时循环无法利用PyArrow的矢量化优化，性能极差。
• ChunkedArray无join属性：pc.drop_null返回的是数组对象，而join是Table的方法，必须基于Table执行join操作。

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