使用批处理添加带有pyarrow的列答案

【问题标题】：Use batches to add a column with pyarrow使用批处理添加带有pyarrow的列
【发布时间】：2022-08-18 01:00:27
【问题描述】：

我目前正在加载一个表，计算一个新列，将列添加到表中并将表保存到磁盘，这一切都很好。问题：我尝试了这批，但收到错误消息：

AttributeError: \'pyarrow.lib.RecordBatch\' object has no attribute \'append_column\'

有谁知道是否有办法做到这一点？

有效的代码，但没有批处理：

import pyarrow.parquet as pq
import pyarrow as pa
 
candidates = pq.ParquetFile(\'input.parquet\').read()
result = []
for row in candidates.to_pylist():
    row_result = function(row)
    result.append(row_result)
candidates_with_result = candidates.append_column(\'new_column_name\', pa.array(result))
pq.write_table(candidates_with_result, \'output.parquet\')

代码不起作用，但总体思路：

candidates = pq.ParquetFile(\'input.parquet\').read()
for batch in candidates.to_batches():
    result = []
    for row in batch.to_pylist():
        row_result = function(row)
        result.append(row_result)
    batch_with_results = batch.append_column(\'new_column_name\', pa.array(result))
    pq.write_table(batch_with_results, \'output.parquet\')

因此，非常感谢有关如何将函数批量应用于箭头表的任何帮助！

谢谢

迪诺

标签： python pyarrow apache-arrow

【解决方案1】：

它不支持开箱即用，但您可以执行以下操作：

    new_column = pa.array(result)
    batch_with_results = pa.RecordBatch.from_arrays(
        batch.columns + [new_column],
        schema=batch.schema.append(pa.field("new_column_name", new_column.type))
    )

【讨论】：

【解决方案2】：

感谢 0x26res 的回答，在这里我添加了 ParquetWriter 以正确附加：

# before I know the schema, I need to calculate one batch
candidates = pq.ParquetFile('input.parquet').read()
writer = None
for batch in candidates.to_batches():
    result = []
    for row in batch.to_pylist():
        row_result = function(row)
        result.append(row_result)
    new_column = pa.array(result)
    batch_with_results = pa.RecordBatch.from_arrays(
        batch.columns + [new_column],
        schema=batch.schema.append(pa.field("new_column_name", new_column.type))
    )
    if not writer:
        writer = pq.ParquetWriter("output.parquet", batch_with_results.schema)
    writer.write_batch(batch_with_results)
writer.close()

【讨论】：