如何通过Python REST API流式返回Arrow Flight的RecordBatch分块数据

原有实现需要先把所有批次全量读取到内存、拼接为Arrow Table、整体序列化后才返回响应，首响应等待时间长且内存占用高。通过流式响应+逐批次序列化的方案即可解决该问题：

实现思路 #

• Arrow 流格式固定开头为Schema信息，后续可直接追加单个RecordBatch的序列化结果，无需依赖全量数据
• 借助Web框架的流式响应能力，使用生成器边读取Flight返回的批次、边序列化、边返回给客户端，无需等待全量处理完成

核心实现代码 #

流式数据生成器 #

def generate_arrow_stream():
    reader = client.do_get(flight_info.endpoints[0].ticket, options)
    print('[INFO] 从Dremio服务器读取查询结果')
    sink = pa.BufferOutputStream()
    writer = None
    try:
        while True:
            batch, metadata = reader.read_chunk()
            # 首次读取到批次时初始化写入器，先写入Schema头
            if not writer:
                writer = pa.RecordBatchStreamWriter(sink, batch.schema)
            # 写入当前批次到缓冲区
            writer.write_batch(batch)
            current_buffer = sink.getvalue()
            if current_buffer.size > 0:
                # 返回当前批次序列化后的字节
                yield current_buffer.to_pybytes()
                # 重置缓冲区避免重复发送已返回内容
                sink = pa.BufferOutputStream()
    except Exception:
        # 读取到流末尾自动跳出
        pass
    finally:
        # 关闭写入器，写入流结束标记
        if writer:
            writer.close()
            final_buffer = sink.getvalue()
            if final_buffer.size > 0:
                yield final_buffer.to_pybytes()

FastAPI 接口示例 #

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

@app.get("/query-arrow-stream")
def get_query_result():
    return StreamingResponse(
        generate_arrow_stream(),
        media_type="application/vnd.apache.arrow.stream",
        headers={
            "Content-Disposition": "attachment; filename=query_result.arrow"
        }
    )

Flask 接口示例 #

from flask import Flask, Response

app = Flask(__name__)

@app.route("/query-arrow-stream")
def get_query_result():
    return Response(
        generate_arrow_stream(),
        mimetype="application/vnd.apache.arrow.stream",
        headers={
            "Content-Disposition": "attachment; filename=query_result.arrow"
        }
    )

注意事项 #

• 必须设置正确的MIME类型 application/vnd.apache.arrow.stream，保证客户端可以直接识别并解析Arrow流
• 该方案首响应时间仅取决于第一个批次的读取+序列化耗时，无需等待全量数据处理完成，内存占用仅为单批次大小
• 客户端接收时可直接调用 pa.open_stream(响应字节流) 逐批次读取数据，无需等待全量下载完成

批次列表示例 #

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