回答问题

将 numpy 对象保存到 parquet 的传统方法是使用 Pandas 作为中间体。但是,我正在处理大量数据,这些数据不适合 Pandas 而不会破坏我的环境,因为在 Pandas 中,数据占用了大量 RAM。

我需要保存到 Parquet,因为我正在使用 numpy 中的可变长度数组,所以对于那个 parquet 实际上保存到比 .npy 或 .hdf5 更小的空间。

以下代码是一个最小示例,它下载我的一小部分数据,并在 pandas 对象和 numpy 对象之间进行转换以测量它们消耗了多少 RAM,并保存到 npy 和 parquet 文件以查看它们占用了多少磁盘空间。

# Download sample file, about 10 mbs

from sys import getsizeof
import requests
import pickle
import numpy as np
import pandas as pd
import os

def download_file_from_google_drive(id, destination):
    URL = "https://docs.google.com/uc?export=download"

    session = requests.Session()

    response = session.get(URL, params = { 'id' : id }, stream = True)
    token = get_confirm_token(response)

    if token:
        params = { 'id' : id, 'confirm' : token }
        response = session.get(URL, params = params, stream = True)

    save_response_content(response, destination)    

def get_confirm_token(response):
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            return value

    return None

def save_response_content(response, destination):
    CHUNK_SIZE = 32768

    with open(destination, "wb") as f:
        for chunk in response.iter_content(CHUNK_SIZE):
            if chunk: # filter out keep-alive new chunks
                f.write(chunk)

download_file_from_google_drive('1-0R28Yhdrq2QWQ-4MXHIZUdZG2WZK2qR', 'sample.pkl')

sampleDF = pd.read_pickle('sample.pkl')

sampleDF.to_parquet( 'test1.pqt', compression = 'brotli', index = False )

# Parquet file takes up little space 
os.path.getsize('test1.pqt')

6594712

getsizeof(sampleDF)

22827172

sampleDF['totalCites2'] = sampleDF['totalCites2'].apply(lambda x: np.array(x))

#RAM reduced if the variable length batches are in numpy
getsizeof(sampleDF)

22401764

#Much less RAM as a numpy object 
sampleNumpy = sampleDF.values
getsizeof(sampleNumpy)

112

# Much more space in .npy form 
np.save( 'test2.npy', sampleNumpy) 
os.path.getsize('test2.npy')

20825382

# Numpy savez. Not as good as parquet 
np.savez_compressed( 'test3.npy', sampleNumpy )
os.path.getsize('test3.npy.npz')

9873964

Answers

您可以使用 Apache Arrow (pyarrow) 直接读/写 numpy 数组到 parquet,这也是 pandas 中 parquet 的底层后端。请注意,parquet 是 tabular 格式,因此仍然需要创建一些表格。

import numpy as np
import pyarrow as pa

np_arr = np.array([1.3, 4.22, -5], dtype=np.float32)
pa_table = pa.table({"data": np_arr})
pa.parquet.write_table(pa_table, "test.parquet")

参考:numpy 到 pyarrow,pyarrow.parquet.write_table