将镶木地板从 s3 拆分成块

Question

我正在使用以下代码从 s3 读取镶木地板文件。接下来，我想分块迭代它。我怎样才能实现它？

import s3fs
import fastparquet as fp

s3 = s3fs.S3FileSystem()
fs = s3fs.core.S3FileSystem()

bucket, path = 'mybucket', 'mypath'
root_dir_path = f'{bucket}/{path}'
s3_path = f"{root_dir_path}/*.parquet"
all_paths_from_s3 = fs.glob(path=s3_path)

fp_obj = fp.ParquetFile(all_paths_from_s3, open_with=s3.open, root=root_dir_path)
df = fp_obj.to_pandas()

一种方法是使用生成器：

def chunks(df, chunksize):
    for i in range(0, len(df), chunksize):
        yield df[i:i + chunksize]

for chunk in chunks(df, 1000):
    # dummy code to transform & operate on chunk
    print(len(chunk))
    # dummy code ends

有什么更节省空间和时间的方法？

Answer 1

就内存效率而言，使用 pyarrow 数据集可能会更好。

这样的事情应该可以工作，并且只会增量加载批次：

from pyarrow.dataset import dataset
ds = dataset(f"s3://{root_dir_path}", format="parquet")
batches = ds.to_batches()
for batch in batches:
   df = batch.to_pandas()
   transform(df)

如果需要 fastparquet，则使用 iter_row_groups 应该更节省内存，但您仍可能希望增量处理。

在这两种情况下，这都应该以增量方式加载数据，而不是创建一个大型数据帧然后对其进行迭代。