从两个csv文件python中选择记录

Question

我有两个 csv 文件。 train_all.csv：

    msno                                         is_churn
0   waLDQMmcOu2jLDaV1ddDkgCrB/jl6sD66Xzs0Vqax1Y=    1
1   QA7uiXy8vIbUSPOkCf9RwQ3FsT8jVq2OxDr8zqa7bRQ=    1
2   fGwBva6hikQmTJzrbz/2Ezjm5Cth5jZUNvXigKK2AFA=    1
3   mT5V8rEpa+8wuqi6x0DoVd3H5icMKkE9Prt49UlmK+4=    1
4   XaPhtGLk/5UvvOYHcONTwsnH97P4eGECeq+BARGItRw=    1
5   GBy8qSz16X5iYWD+3CMxv/Hm6OPSrXBYtmbnlRtknW0=    1
6   lYLh7TdkWpIoQs3i3o6mIjLH8/IEgMWP9r7OpsLX0Vo=    1

它有 1963891 条记录，但只有 1082190 条唯一记录。这就是为什么我用 pandas 对独特的 msno-s 进行排序。 另一个 csv 大约 30GB，包含超过 9 亿条记录和用户日志。

msno    date    num_25  num_50  num_75  num_985 num_100 num_unq total_secs
rxIP2f2aN0rYNp+toI0Obt/N/FYQX8hcO1fTmmy2h34=    20150513    0   0   0   0   1   1   280.335
rxIP2f2aN0rYNp+toI0Obt/N/FYQX8hcO1fTmmy2h34=    20150709    9   1   0   0   7   11  1658.948
yxiEWwE9VR5utpUecLxVdQ5B7NysUPfrNtGINaM2zA8=    20150105    3   3   0   0   68  36  17364.956
yxiEWwE9VR5utpUecLxVdQ5B7NysUPfrNtGINaM2zA8=    20150306    1   0   1   1   97  27  24667.317
yxiEWwE9VR5utpUecLxVdQ5B7NysUPfrNtGINaM2zA8=    20150501    3   0   0   0   38  38  9649.029

我想遍历 train_all.csv 中唯一的 msno-s，并在 user_logs.csv 中为每个 msno 搜索 5 个 user_log 记录。我的代码在大约 20 分钟后停止，results.csv 中只有 104 条记录 - 所需的输出文件与匹配的 msno - user_logs。

import pandas as pd
import csv

reader = csv.reader(open('user_logs/user_logs.csv','r'))
writer = csv.writer(open('results.csv','w',newline=''))

data = pd.read_csv("train_all.csv")
unique_msnos = data["msno"].unique()

i = 0

for msno in range(len(unique_msnos)):
    counter = 0

    for row in reader:
        results_row = row

        if unique_msnos[msno] == row[0]:
            writer.writerow(results_row)
            counter+=1
            if counter == 5:
                i+=1
                break
            else:
                continue
        else:
            continue

Answer 1

它没有找到所有东西就停止的可能原因：

• 每次读取一行时，内部循环都会将文件中的位置向前移动。
• 当您跳出内部循环时，文件计数器/指针/位置（？）仍保留在它离开的位置
• 在外循环的下一次迭代期间，内循环从最后一个文件位置开始，而不是从开头开始。所以它在文件的前几行/行中找不到任何有效的 msno。您可以通过在内部循环之前打印文件的位置 (file.tell()) 来验证这一点（记住 ctrl-c 可以跳出它。）

您可以修复在内循环完成后将文件发送回开头（无论是通过 StopIteration 还是中断），以便内循环可以搜索整个文件以查找无。准备好等待更长的时间。

with open('user_logs/user_logs.csv','r') as infile:
    reader = csv.reader(infile)
    for msno in range(len(unique_msnos)):
        counter = 0
        #print(infile.tell())
        for row in reader:
            #inner_loop_stuff...
            #inner_loop_stuff...
            #inner_loop_stuff...
        infile.seek(0)

---

unique_msnos 有 1.1E6 条记录，reader 有 9E8 条记录。对于unique_msnos 中的每条记录，您将迭代 9e8 次 - 这将是 1e15 次迭代。你必须小心嵌套循环。

当checking for membership你应该使用一个集合。

unique_msnos = set(unique_msnos)

您想跟踪找到的记录数，因此可以将 unique_msnos 制作成一个字典，其中包含该计数的值。

unique_msnos = dict.fromkeys(unique_msnos, 0)

然后遍历 reader 检查成员资格并添加逻辑以查看是否已经有 5 个：

for row in reader:
    c = unique_msnos.get(row[0], None)
    if c in (5, None):
        continue
    writer.write(row)
    unique_msnos[row[0]] += 1
    ...

这应该会有所缓解 - 9e8 次迭代。可能还有其他优化。

---

如果保证 msno-s 的长度都相同，则使用切片而不是使用 csv.Reader 可能会获得增量改进。我真的不知道 csv 阅读器的速度有多快，但切片非常快。

类似

# unique_msnos is a dictionary
with open(open('user_logs/user_logs.csv','r') as f:
    for line in f:
        c = unique_msnos.get(line[:44], None)
        ....

它不会减少迭代次数，但如果你得到 (a_small_improvement * 9e8) 改进，它可能会很明显。

---

看起来您正在制作一个 DataFrame 然后将其丢弃。直接从文件中生成unique_msnos 可能会更快：

unique_msnos = {}
with open("train_all.csv") as f:
    for line in f:
        _, msno, churn = line.split()
        unique_msnos[msno] = 0

您需要尝试一下，看看它是否更快。

---

您可能希望自己制作一些较小的文件以供使用，同时尝试使其工作并进行优化。只需为每个文件执行一次，然后使用小文件进行调试。使用cprofile 或timeit 尝试找出瓶颈在哪里

limit = 10000  # or whatever
with open(file_path, 'r') as infile, open(small_filepath, 'w') as outfile:
    for n, line in enumerate(infile):
        outfile.write(line)
        if n == limit:
            break