如何在FastAPI中用PyArrow生成字节流而非写入文件以避免内存拷贝

问题背景 #

在FastAPI服务中生成Arrow IPC格式数据返回前端时，现有实现是将整个表写入BytesIO缓冲区后再通过StreamingResponse返回，存在不必要的内存拷贝和内存占用过高问题。尝试过以下方案但未解决：

• 使用pa.BufferOutputStream无法直接适配FastAPI的返回类型
• 分批写入BytesIO并yield内容，但sink.truncate()无法有效释放内存，仍存在内存拷贝

前端依赖apache-arrow的tableFromIPC方法从fetch请求中解析数据。

解决方案：自定义PyArrow OutputStream实现流式输出 #

核心思路是实现一个自定义的PyArrow OutputStream子类，将writer写入的字节直接通过生成器输出，完全跳过中间缓冲区的内存拷贝。

实现代码 #

from fastapi import FastAPI, StreamingResponse
import pyarrow as pa
import pyarrow.ipc as ipc
from queue import Queue
from typing import Generator

app = FastAPI()

class GeneratorOutputStream(pa.OutputStream):
    def __init__(self):
        self.queue = Queue()
        self.closed = False

    def write(self, data: bytes) -> int:
        if self.closed:
            raise ValueError("Stream is closed")
        self.queue.put(data)
        return len(data)

    def close(self) -> None:
        self.closed = True
        self.queue.put(None)  # 标记流结束

    def generator(self) -> Generator[bytes, None, None]:
        while True:
            chunk = self.queue.get()
            if chunk is None:
                break
            yield chunk

@app.get("/api/getdata", response_class=StreamingResponse)
async def getdata():
    # 示例数据，实际可替换为从数据源加载的表或数据集
    table = pa.Table.from_pydict({
        "name": ["Alice", "Bob", "Charlie", "David"] * 1000,
        "age": [25, 30, 22, 35] * 1000
    })

    stream = GeneratorOutputStream()
    # 创建IPC File格式的writer（前端tableFromIPC支持File和Stream两种格式，这里用File格式）
    with ipc.new_file(stream, table.schema) as writer:
        for batch in table.to_batches(max_chunksize=1024):
            writer.write(batch)
    
    return StreamingResponse(
        content=stream.generator(),
        media_type="application/vnd.apache.arrow.file"
    )

代码说明 #

1. GeneratorOutputStream类：

   • 继承自pa.OutputStream，实现PyArrow writer要求的write和close方法
   • 用队列缓存每个写入的字节块，通过generator方法将队列内容转为可yield的生成器
   • 关闭流时放入None标记，告诉生成器停止输出

2. FastAPI接口实现：

   • 创建自定义流对象，初始化IPC writer
   • 分批写入表数据，每个batch的字节直接进入队列
   • 返回StreamingResponse，将生成器作为内容传入，媒体类型指定为Arrow File格式的MIME类型

前端验证代码 #

import { tableFromIPC } from "apache-arrow";

async function fetchArrowData() {
    const response = await fetch("/api/getdata");
    const table = await tableFromIPC(response);
    console.table([...table]);
}

fetchArrowData();

为什么之前的方案无效？ #

• BytesIO分批yield：每次调用sink.getvalue()会复制整个缓冲区的内容，即使truncate()后，BytesIO的底层缓冲区并不会立即释放内存，仍存在内存冗余
• BufferOutputStream：PyArrow的BufferOutputStream最终会生成一个pa.Buffer，本质还是一次性持有全部数据，无法直接转为流式输出

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