使用由列变量确定的块大小加载熊猫数据框

Posted 2023-02-16

【中文标题】使用由列变量确定的块大小加载熊猫数据框

【英文标题】：Load pandas dataframe with chunksize determined by column variable

【发布时间】：2017-07-02 21:32:15

【问题描述】：

如果我的 csv 文件太大而无法使用 pandas 加载到内存中（在本例中为 35gb），我知道可以使用 chunksize 分块处理文件。

但是我想知道是否可以根据列中的值更改块大小。

我有一个 ID 列，然后每个 ID 都有几行信息，如下所示：

ID,   Time,  x, y
sasd, 10:12, 1, 3
sasd, 10:14, 1, 4
sasd, 10:32, 1, 2
cgfb, 10:02, 1, 6
cgfb, 10:13, 1, 3
aenr, 11:54, 2, 5
tory, 10:27, 1, 3
tory, 10:48, 3, 5
ect...

我不想将 ID 分成不同的块。例如将处理大小为 4 的块：

ID,   Time,  x, y
sasd, 10:12, 1, 3
sasd, 10:14, 1, 4
sasd, 10:32, 1, 2
cgfb, 10:02, 1, 6
cgfb, 10:13, 1, 3 <--this extra line is included in the 4 chunk

ID,   Time,  x, y
aenr, 11:54, 2, 5
tory, 10:27, 1, 3
tory, 10:48, 3, 5
...

有可能吗？

如果不使用带有 for 循环的 csv 库，则如下：

for line in file:
    x += 1
    if x > 1000000 and curid != line[0]:
        break
    curid = line[0]
    #code to append line to a dataframe

虽然我知道这只会创建一个块，并且 for 循环需要很长时间来处理。

【问题讨论】：

您的文件存储在哪里？一种方法是预先过滤 ID（也许，如果您要提取数据库提取，使用 SELECT * FROM ... WHERE ID == BLAH），为每个唯一 ID 值创建一个不同的文件。

很遗憾，这是不可能的，我无权访问数据库

这可能是一个 PITA，但我认为它会起作用：如果您现在尝试使用 chunksize，流式传输整个 35gb 文件，并为每个唯一值创建一个单独的 CSV ID (set(df['ID']))？然后，对于较大文件中的每一行，您将该行写入（读取：附加）到与该行的 ID 对应的现有 ID 文件中？编码不会太难，尽管我怀疑它可能需要多长时间......只是一个想法！否则，恐怕我也帮不上什么忙了。

【参考方案1】：

如果您逐行遍历 csv 文件，您可以使用依赖于任何列的生成器 yield 块。

工作示例：

import pandas as pd

def iter_chunk_by_id(file):
    csv_reader = pd.read_csv(file, iterator=True, chunksize=1, header=None)
    first_chunk = csv_reader.get_chunk()
    id = first_chunk.iloc[0,0]
    chunk = pd.DataFrame(first_chunk)
    for l in csv_reader:
        if id == l.iloc[0,0]:
            id = l.iloc[0,0]
            chunk = chunk.append(l)
            continue
        id = l.iloc[0,0]
        yield chunk
        chunk = pd.DataFrame(l)
    yield chunk

## data.csv ##
# 1, foo, bla
# 1, off, aff
# 2, roo, laa
# 3, asd, fds
# 3, qwe, tre
# 3, tre, yxc   

chunk_iter = iter_chunk_by_id("data.csv")

for chunk in chunk_iter:
    print(chunk)
    print("_____")

输出：

   0     1     2
0  1   foo   bla
1  1   off   aff
_____
   0     1     2
2  2   roo   laa
3  2   jkl   xds
_____
   0     1     2
4  3   asd   fds
5  3   qwe   tre
6  3   tre   yxc
_____

【讨论】：

所以这将为每个 id 创建一个块，我将如何创建更大的块，大小为 1000000，然后附加相同 id 的剩余行，以确保没有相同 id 的行在单独的块中？

@JoshuaKidd 你理解函数iter_chunk_by_id() 吗？您可以轻松修改它以对应您的新请求：在 if id == l.iloc[0,0] 行中添加 or 条件，如果块的长度小于 1000000，则继续向块添加行（顺便说一句：你为​​什么要删除我的答案是否被接受？）

【参考方案2】：

我基于@elcombato 提供的答案来获取任何块大小。实际上我有一个类似的用例，并且逐行处理使我的程序慢得难以忍受

def iter_chunk_by_id(file_name, chunk_size=10000):
"""generator to read the csv in chunks of user_id records. Each next call of generator will give a df for a user"""

csv_reader = pd.read_csv(file_name, compression='gzip', iterator=True, chunksize=chunk_size, header=0, error_bad_lines=False)
chunk = pd.DataFrame()
for l in csv_reader:
    l[['id', 'everything_else']] = l[
        'col_name'].str.split('|', 1, expand=True)
    hits = l['id'].astype(float).diff().dropna().nonzero()[0]
    if not len(hits):
        # if all ids are same
        chunk = chunk.append(l[['col_name']])
    else:
        start = 0
        for i in range(len(hits)):
            new_id = hits[i]+1
            chunk = chunk.append(l[['col_name']].iloc[start:new_id, :])
            yield chunk
            chunk = pd.DataFrame()
            start = new_id
        chunk = l[['col_name']].iloc[start:, :]

yield chunk