Support Apache pyarrow, specifically parquet format

[IgnacioJPickering]

Currently, in order to release the memory held on by the mmapped numpy arrays, the repo uses madvise in linux / macOS. This reduces RSS ~40% (at least on Linux), but I thought it may be a good idea to support Apache pyarrow's parquet format, since it is possible to read from disk in chunks of "row groups", and iterate over those, and afterwards release the memory automatically. Parquet is super popular format and supporting it for fps should (hopefully?) not be much trouble.

I've made a few preliminary tests and it seems that generating fingerprint files with row groups of ~100k each, and iterating in batches of ~512-1024 rows is a pretty robust default that doesn't eat up a lot of memory (I got constant memory usage of around 100 MB in ubuntu 24.04, python 3.11). Further experimentation with this is needed.

The trickiest part of this would be extracting the clusters from disk for refinement purposes, from the pyarrow dataset, since the parquet format is not really designed for efficient random access. Probably the best thing is to pre-allocate numpy arrays, iterate over the row groups and filter the requierd rows to build the queried clusters, and then copy them to the arrays. Some benchmarks will be needed to see if there are more efficient ways to do this.

Another thing to note is that the latest pyarrow versions seem to have somewhat of a buggy interaction with the default memory allocator they use internally (mimalloc) (apache/arrow#47266, apache/arrow#47595) so maybe that is something to take into account. I've gotten very strange memory stats when running memray over simple read_table commands in pyarrow. Maybe setting the memory pool manually to use the system's malloc is the best choice here, or maybe I'm just being overly cautious and this is not an issue. In any case further tests would be nice.

In the end, there are quite a few headaches associated with this, so it may not be worth the extra complexity, idk (the madvice scheme does introduce a bit of complexity too bear in mind, and it only works on Unix).

This is an example script I've been using for testing:

import time
from copy import deepcopy

import numpy as np
from pyarrow import parquet  # requires pip install pyarrow

make_fps = True
if make_fps:
    rng = np.random.default_rng(12620509540149709235)
    fps = rng.integers(0, 256, (10_000_000, 256), dtype=np.uint8)

    ar = pa.FixedSizeListArray.from_arrays(fps.reshape(-1), fps.shape[1])
    table = pa.Table.from_arrays([ar], names=["fingerprints"])
    # I believe for our use case it is better to use a small row-group-size to avoid
    # large RAM spikes (at the expense of IO throughput)

    # (default for very large tables is 1024 ** 2, produces memory peaks of ~700B)
    # 500k produces memory peaks of ~400 MB
    # 250k produces memory peaks of ~300 MB
    # 100k seems a good default that avoids the memory peaks of ~700 MB

    # fingerprints are *practically incompressible* since they are random bits so better
    # to save as uncompressed, in principle should reduce latency, has no measurable
    # effect though. File is slightly larger than the corresponding *.npy file since
    # parquet stores some metadata, but the difference is negligible (~5%)
    parquet.write_table(
        table, "./fingerprints.pq", row_group_size=100_000, compression="NONE"
    )

    # Save also as npy file
    np.save("./fingerprints.npy", fps)

pq_file = parquet.ParquetFile("./fingerprints.pq")

# Note that iteration inside the *fit* functions should look somewhat like this:
rows = 512  # 512 seems a good default in my tests
_start = time.perf_counter()
for batch in pq_file.iter_batches(rows):
    # Fetch the numpy array from the pyarrow buffer, zero copy
    # The API for this is a bit awkward, hopefully the pyarrow team improves this in the
    # future?
    #
    # [0]: get the FixedSizeListArray from the RecordBatch
    # buffers(): get all Buffers associated with the array (just 1)
    # [-1]: Buffer is actually a tuple with some metadata, last elem of tuple holds buf
    # np.frombuffer(...): Use python's __buffer__ protocol to take ownership of the mem
    arr = np.frombuffer(batch[0].buffers()[-1], dtype=np.uint8).reshape(rows, -1)
    #
    # Do something with the elements in a loop.
    # Inside bitbirch iterations we need to at least do a copy to prevent bitbirch to
    # hold on to the memory
    for x in arr:
        _ = x.copy()

# This is actually pretty slow, but slightly faster than mmaped numpy, so not bad
print(f"Time elapsed: {time.perf_counter() - _start} s", flush=True)

# For reference this is the numpy iteration
arr = np.load("./fingerprints.npy", mmap_mode="r")
_start = time.perf_counter()
for x in arr:
    _ = x.copy()
print(f"Time elapsed: {time.perf_counter() - _start} s", flush=True)

I'm not planning on experimenting with this more for the time being, although I think it is interesting, I have too many other things going on right now 😅. @vicciv1623 let me know if you have questions, I'll do my best to help!

[IgnacioJPickering]

@ramirandaq I think the current numpy array + madvise is pretty good for Linux/macOS, but I wanted to keep my findings here, and maybe Vic would like to investigate some.

[ramirandaq]

Thanks for looking into this! @vicciv1623, take a look.