reimplement parquet reader using ParquetFile about embedding-reader HOT 20 CLOSED

maybe just try to use the current code but cache usage of ParquetFile and close the file when all pieces have been read

from embedding-reader.

consider using https://arrow.apache.org/docs/python/dataset.html , it's really good

from embedding-reader.

along with https://arrow.apache.org/docs/python/filesystems.html#filesystem-fsspec

from embedding-reader.

import fsspec
import pyarrow.dataset as ds
from tqdm import tqdm
fs, p = fsspec.core.url_to_fs("https://mystic.the-eye.eu/public/AI/cah/laion5b/embeddings/laion1B-nolang/laion1B-nolang-metadata")
files = fs.ls(p, detail=False)
d = ds.dataset(files, filesystem=fs)
b = d.to_batches()
for _ in tqdm(b):
    pass

about 1M sample/s eg 200MB/s, saturates the external server connection

from embedding-reader.

idea:

• use numpy read with the current implementation for numpy
• use pyarrow dataset for parquet
• zip the 2 for the numpy parquet reader

this should fix the parquet speed and will solve the memory issue

from embedding-reader.

doesn't work due to no support of start/end in pyarrow.dataset

from embedding-reader.

numpy reader doing 300MB/s from https fsspec
numpy parquet is at 30MB/s ...
this needs to be improved

from embedding-reader.

maybe https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Scanner.html#pyarrow.dataset.Scanner can be used to support start stop

from embedding-reader.

no, can't

from embedding-reader.

ok best idea is to cache the files

from embedding-reader.

seems a little bit better but not much

@lru_cache(maxsize=None)
def open_parquet_file(fs, filename):
    return pq.read_table(fs.open(filename, "rb"), use_threads=False)

from embedding-reader.

actually lru cache doesn't work with multiple threads...

from embedding-reader.

r = Semaphore(1)

d = {}
def open_parquet_file(fs, filename):
    r.acquire()
    if filename in d:
        r.release()
        return d[filename]
    print(filename)
    t = pq.read_table(fs.open(filename, "rb"), use_threads=False)
    d[filename] = t
    r.release()
    return t

makes things much much faster

from embedding-reader.

estimated 325min total (laion1B-nolang) with use_threads=True ; probably the same with False

from embedding-reader.

150min for numpy alone

from embedding-reader.

1012 without this change

from embedding-reader.

seems to have solved the memleak too

from embedding-reader.

https://github.com/rom1504/embedding-reader/pull/21/files

faster but didn't solve memleak

from embedding-reader.

actually looks like it did. memory usage is a bit high (15GB with these settings), but memleak seems gone

from embedding-reader.

did a fix here

from embedding-reader.