时间戳

首页 > 解决方案 > 有没有办法将 pandas 中块的大小定义为可用内存的函数？

python-3.x - 有没有办法将 pandas 中块的大小定义为可用内存的函数？

问题描述

我知道我可以加载包含数据块的文件：

import pandas
for chunk in pandas.read_csv("path_to_my_csv.csv", chunksize=1e9):
    # Process

其中 的值chunksize对应于每个“块”包含的行数。我想要做的是：

import pandas
for chunk in pandas.read_csv("path_to_my_csv.csv", chunkmem="200GB"):
    # Process

我想这样做的原因是能够在不同的机器上处理数据（具有不同数量的可用 RAM），并使用psutil.virtual_memory或类似的方式以自动方式参数化我的分块。

这样做的一种方法是计算单行的内存占用（从每列的数据类型），并使用它来参数化 的值chunksize，但理想情况下我希望能够使用不同的数据集来做到这一点结构。

编辑（回应黄比尔）：

鉴于 Pandas API 中没有直接实现，我这样做的方法是首先估计数据帧的内存占用：

import pandas
numberOfRows = int(1e10)  # Known a-priori
firstRecord = pandas.read_csv("big_data.csv", chunksize=2).get_chunk()
firstTwoRecords = pandas.read_csv("big_data.csv", chunksize=3).get_chunk()
rowFootprint = (firstTwoRecords.memory_usage().sum() -
                firstRecord.memory_usage().sum())
estimatedFootprint = numberOfRows * rowFootprint
print(estimatedFootprint)

然后将可用内存（来自psutil.virtual_memory）除以该估计值以获得块大小。这个估计器只需要读取文件的前两行。

标签： python-3.xpandasdataframe

解决方案

简短的回答：不。文档显示没有这样的参数。

有用的答案：你可以估计合适的chunksize。例如，通过读取前 30 行左右来猜测行的平均大小。但失败的几率实际上取决于数据。

import psutil
from pathlib import Path

# https://www.kaggle.com/tmdb/tmdb-movie-metadata?select=tmdb_5000_movies.csv
file_path = Path("/mnt/ramdisk/tmdb_5000_movies.csv")
fill_rate = 0.1
n_rows = 10

def estimate_bpl(file_path, n_rows=10):
    """Return estimates of bytes per line using the first n lines"""
    with open(file_path) as f:
        length = 0
        for i, line in enumerate(f):
            if i == n_rows:
                break
            length += len(line.encode('utf8'))
    return length / n_rows

avail_mem = psutil.virtual_memory().available
bpl = estimate_bpl(file_path, n_rows)
chunksize = int(avail_mem * fill_rate / bpl)

print(f"avail_mem={avail_mem}, fill rate={fill_rate}, bytes per line={bpl}, chunksize={chunksize}")
avail_mem=11166822400, fill rate=0.1, bytes per line=1409.4, chunksize=792310

推荐阅读