PyArrow ParquetWriter如何限制输出文件大小并实现自动拆分？

报错原因 #

PyArrow的ParquetWriter默认创建的是只写模式的文件句柄，其内置的size()方法仅支持可读模式的文件对象，因此直接调用会触发OSError报错。

可行实现方案 #

方案1：查询磁盘文件实际大小（快速实现） #

该方案无需修改现有写入逻辑，仅需每次写入后查询当前输出文件的磁盘占用即可。考虑到Parquet文件关闭时会写入末尾元数据Footer（通常仅几KB到几十KB），可将拆分阈值设为270MB左右，预留Footer空间即可保证最终文件大小落在200~300MB区间。

import os
import pyarrow.parquet as pq

# 拆分阈值：270MB，预留30MB空间给Parquet Footer
MAX_WRITE_SIZE = 270 * 1024 * 1024
# 替换为你实际要写入的table对象
write_table = your_table_obj
file_seq = 0
# 初始化第一个写入文件
current_file = f"output_{file_seq}.parquet"
writer = pq.ParquetWriter(current_file, write_table.schema)

for i in range(300000):
    writer.write_table(write_table)
    # 获取当前文件已写入大小
    current_size = os.path.getsize(current_file)
    if current_size >= MAX_WRITE_SIZE:
        # 关闭当前writer，写入Footer完成文件
        writer.close()
        file_seq += 1
        # 创建新的写入文件
        current_file = f"output_{file_seq}.parquet"
        writer = pq.ParquetWriter(current_file, write_table.schema)
# 关闭最后一个文件的writer
writer.close()

方案2：自定义计数IO类（更精准高效） #

如果写入批次非常多，频繁调用系统接口查询文件大小会有微小性能损耗，可以自定义可统计写入字节数的IO类，直接获取已写入的压缩后字节数，统计精度更高。

import pyarrow as pa
import pyarrow.parquet as pq

class ByteCountWriter(pa.NativeFile):
    def __init__(self, file_path):
        super().__init__()
        self._fd = open(file_path, "wb")
        self._written_bytes = 0

    def write(self, data):
        write_len = self._fd.write(data)
        self._written_bytes += write_len
        return write_len

    def tell(self):
        return self._written_bytes

    def close(self):
        self._fd.close()
        super().close()

# 阈值可以设为290MB，无需额外预留空间
MAX_WRITE_SIZE = 290 * 1024 * 1024
write_table = your_table_obj
file_seq = 0
current_io = ByteCountWriter(f"output_{file_seq}.parquet")
writer = pq.ParquetWriter(current_io, write_table.schema)

for i in range(300000):
    writer.write_table(write_table)
    if current_io.tell() >= MAX_WRITE_SIZE:
        writer.close()
        current_io.close()
        file_seq += 1
        current_io = ByteCountWriter(f"output_{file_seq}.parquet")
        writer = pq.ParquetWriter(current_io, write_table.schema)
writer.close()
current_io.close()

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