streaming feather files

stream.py

import time
import sys
import itertools
import os

import numpy as np
import pandas as pd
import pyarrow as pa
import psutil

nrows = 512**2
ncols = 32
nbatches = int(sys.argv[1])
batch_size = nrows // nbatches
print(f"{nbatches} batches of {batch_size} rows")
 

process = psutil.Process(os.getpid())
previous_mem = 0

def pmem(msg):
    global previous_mem
    current_mem = process.memory_info().rss // 1024
    print(f"{msg}: {current_mem - previous_mem:+d}kb")
    previous_mem = current_mem


def generate_rows():
    for i in range(nrows):
        yield {f"col{j}": np.random.randn() for j in range(ncols)}


def chunked_iterable(iterable, size):
    it = iter(iterable)
    while True:
        chunk = tuple(itertools.islice(it, size))
        if not chunk:
            break
        yield chunk


def write(rows, schema):
    batches = chunked_iterable(rows, batch_size)

    options = pa.ipc.IpcWriteOptions(compression='zstd', use_threads=True)

    with pa.OSFile('test.feather', 'wb') as sink:
        with pa.ipc.new_file(sink, schema, options=options) as writer:
            for i, batch in enumerate(batches):
                batch_data = list(batch)
                b = pa.RecordBatch.from_pylist(batch_data, schema)
                writer.write(b)
                pmem(f"{i} {len(batch_data)}")


previous_mem = process.memory_info().rss // 1024
schema = pa.schema([pa.field(f"col{i}", pa.float64()) for i in range(ncols)])
rows = generate_rows()
write(rows, schema)


df = pd.read_feather('test.feather')
pmem("post-read")

does this allow u to progressively write to feather files?