Как сделать генератор байтов вместо записи в файл из pyarrow для фастапи ⇐ Python

[Anonymous]

Скажем, в общих чертах моя установка выглядит следующим образом:

Код: Выделить всё

from fastapi import FastAPI, Response
import pyarrow as pa
import pyarrow.ipc as ipc
app = FastAPI()

@app.get("/api/getdata")
async def getdata():
table = pa.Table.from_pydict({
"name": ["Alice", "Bob", "Charlie"],
"age": [25, 30, 22]})

### Not really sure what goes here
## something like this...
sink = io.BytesIO()
with ipc.new_file(sink, table.schema) as writer:
for batch in table.to_batches():
writer.write(batch)
sink.seek(0)
return StreamingResponse(content=sink, media_type="application/vnd.apache.arrow.file")

Это работает, но я сначала копирую всю таблицу в BytesIO? Кажется, что мне нужно сделать генератор, который выдает все, что Writer.write(batch) записывает в буфер, вместо того, чтобы фактически записывать это, но я не знаю, как это сделать. Я пробовал использовать pa.BufferOutputStream вместо BytesIO, но не могу поместить его в качестве возвращаемого объекта для fastapi.
Моя цель — иметь возможность получать данные на стороне js вот так...

Код: Выделить всё

import { tableFromIPC } from "apache-arrow";
const table = await tableFromIPC(fetch("/api/getdata"));
console.table([...table]);

В моем подходе это работает, мне просто хотелось бы знать, есть ли способ сделать это без копирования.
Я попробовал сделать небольшими партиями, но это все равно не работало так:

Код: Выделить всё

@app.get("/")
def get_table():
table_or_ds = ds.dataset(file_or_files, format=format)

def gen_out():
with BytesIO() as sink:
with ipc.new_file(sink, table_or_ds.schema) as writer:
for batch in table_or_ds.to_batches():
writer.write_batch(batch)
yield sink.getvalue()
sink.seek(0)
sink.truncate()
yield sink.getvalue()

return StreamingResponse(gen_out(), media_type="application/vnd.apache.arrow.file")

Проблема в том, что мойка.truncate не очищает память, поэтому это не помогает.

Подробнее здесь: https://stackoverflow.com/questions/767 ... rrow-for-f