Читать большой файл из API по частям и записывать каждую часть в качестве файла Parquet с помощью PyArrow

Я экспериментирую с PyArrow, но у меня есть некоторые трудности в понимании некоторых моментов. Что я хочу достичь, так это читать большой файл (CSV в этом примере) частями и сохранять каждую часть в виде файла Parquet.

У меня есть следующий код:

import requests

import pyarrow as pa

def create_record_batches(lines, schema):
    sink = pa.BufferOutputStream()
    writer = pa.ipc.RecordBatchStreamWriter(sink, schema)
    for line in lines:
        if line:
            newline = [{col_name: val for col_name, val in zip(schema.names, line.decode().split(","))}]
            writer.write(pa.RecordBatch.from_pylist(newline))
        else:
            buffer_value = sink.getvalue()
            yield buffer_value.to_pybytes()
            # writer.close()
            sink = pa.BufferOutputStream()
            writer = pa.ipc.RecordBatchStreamWriter(sink, schema)

    writer.close()
    buffer_value = sink.getvalue()
    yield buffer_value.to_pybytes()

with requests.get(
    url="https://data.cityofnewyork.us/api/views/kxp8-n2sj/rows.csv?accessType=DOWNLOAD",
    stream=True
) as response:
    response.raise_for_status()

    lines = response.iter_lines(chunk_size=128*1024*1024, delimiter=b"\n")
    schema_str = next(lines).decode()
    schema = pa.schema([ (col_name, pa.string()) for col_name in schema_str.split(",") ])

    buffer = create_record_batches(lines, schema)

    counter = 0
    for buf in buffer:
        if counter == 3:
            break
        # print(buf)

        print("*")
        print("**********")
        print("**********")

        counter += 1

Я еще не дошел до стадии записи чего-либо в файл Parquet, но позвольте мне разобрать свой код и также подробно объяснить, что мне неясно.

• Я использую requests.get(url, stream=True), чтобы получить Response, и использую response.iter_lines(chunk_size=128*1024*1024), чтобы иметь возможность работать с частями по 128 МБ.
• Поскольку первая строка файла является заголовком, я использую следующее, чтобы создать схему Pyarrow:

schema_str = next(lines).decode()
schema = pa.schema([ (col_name, pa.string()) for col_name in schema_str.split(",") ])

• С оставшимися lines я вызываю функцию create_record_batches(lines, schema). Эта функция отвечает за итерацию по каждой строке каждой части, чтобы создать RecordBatch и записать его в объект RecordBatchStreamWriter.
  • В конце части, представленной пустой строкой, такой как b'', мы переходим в случай else, и я возвращаю байты из своего sink, так что он должен вернуть все RecordBatch конкретной части.
  • Все еще в случае else, я создаю совершенно новый sink и writer, моя цель здесь – очистить как sink, так и writer. Я пытался использовать sink.truncate() и sink.seek(0), но не смог этого сделать, потому что sink.getvalue() закрывает sink.
  • Наконец, я закрываю writer и возвращаю оставшееся содержимое в sink, но я не знаю, имеет ли это смысл, учитывая, что я использую iter_lines()..?

Что я не понимаю в моем текущем коде, так это то, что когда я проверяю потребление памяти (используя htop) своей программы во время отладки и имея точку останова на каждом объекте buf, я ожидал, что память будет колебаться вокруг 128 МБ, так как это размер моих частей, но память увеличивается, так что я что-то делаю не так, но не знаю, что. Мне кажется, это происходит потому, что sink и writer не действительно “сбрасываются”..

Что касается записи в файлы Parquet, то в моем плане использовать каждый объект buf, чтобы создать PyArrow.Table, а затем записать каждую таблицу в файл Parquet. Чтобы создать PyArrow.Table, я думаю, что я буду использовать следующее:

import pyarrow.parquet as pq

...

    for i_chunk, buf in enumerate(buffer):
        table = pa.ipc.open_stream(pa.py_buffer(buf)).read_all()
        file_path = f"some_file_name_chunk{i_chunk}.parquet"
        pq.write_table(table, file_path)

Я иду в правильном направлении с моим текущим кодом? И также с тем, что я планирую делать, чтобы записать свои части в файлы Parquet?

Большое спасибо, если вы можете помочь мне лучше освоить библиотеку PyArrow!

PS: Я решил не использовать такие библиотеки, как pandas или csv, моя цель – попробовать манипулировать только буферами и потоковыми объектами, если это возможно.