python - 使用 Python dataframe 高效地将数百万行写入文件

Question

我有以下代码片段，可将 CSV 读入数据帧，并以符合 Redis 协议(protocol)的方式将键值对写入文件，即 SET key1 value1。代码是零碎的，我尝试使用多处理，但我不确定它的性能( yield )。

CSV 有大约 600 万行，可以很快(不到 2 分钟)读入数据帧。输出文件有 1200 万行(输入文件每行 2 行)。这大约需要 50 分钟才能完成。可以优化/更改我的代码的任何部分以使其运行得更快吗？文件完成后，将其加载到 Redis 只需不到 90 秒的时间。真正的瓶颈在于写入文件。我将有几个这样的文件要写，每个文件花费 50-60 分钟确实不理想。这个特定的数据集有 600 万行和 10 列，主要由带有一些浮点列的字符串组成。 Redis 键是字符串，浮点值是键值对中的 Redis 值。其他数据集的大小与此类似，甚至更大(在行和列方面)。

我正在研究将我生成的所有字符串加载到一个数据框中，然后使用 to_csv() 函数将其转储到一个文件中，但我不确定它的性能如何。

filepath = '/path/to/file.csv'

def df_to_file:
    df = pd.read_csv(filepath)
    f = open('output_file', 'w')
    for i in range(len(df.index)):
        if df['col1'].iloc[i] != '':
            key1 = str(df['col1'].iloc[i])+str(df['col4'].iloc[i])+str(df['col5'].iloc[i])+...+str(df['col_n'].iloc[i])
            val1 = df['col_n+1'].iloc[i]

            key1a = str(df['col1'].iloc[i])+str(df['col4'].iloc[i])+str(df['col5'].iloc[i])+...+str(df['col_n'].iloc[i])
            val1a = df['col_n+2'].iloc[i]

            print('SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a), file = f)

        if df['col2'].iloc[i] != '':
            key1 = str(df['col2'].iloc[i])+str(df['col4'].iloc[i])+str(df['col5'].iloc[i])+...+str(df['col_n'].iloc[i])
            val1 = df['col_n+1'].iloc[i]

            key1a = str(df['col2'].iloc[i])+str(df['col4'].iloc[i])+str(df['col5'].iloc[i])+...+str(df['col_n'].iloc[i])
            val1a = df['col_n+2'].iloc[i]

            print('SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a), file = f)
        if df['col3'].iloc[i] != '':
            key1 = str(df['col3'].iloc[i])+str(df['col4'].iloc[i])+str(df['col5'].iloc[i])+...+str(df['col_n'].iloc[i])
            val1 = df['col_n+1'].iloc[i]

            key1a = str(df['col3'].iloc[i])+str(df['col4'].iloc[i])+str(df['col5'].iloc[i])+...+str(df['col_n'].iloc[i])
            val1a = df['col_n+2'].iloc[i]

            print('SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a), file = f)
    f.close()

p = Process(target = df_to_file)
p.start()
p.join() 

Answer 1

使用类似 df['col1'].loc[...] 的结构遍历单个行会很慢，iloc和 loc基于选择器用于选择整个数据帧，并做很多与索引对齐相关的事情，如果对每一行都做的话会产生很高的开销。相反，只需使用 df.itertuples()遍历行将明显更快。

def df_to_file:
    df = pd.read_csv(filepath)
    f = open('output_file', 'wb') # writing in binary mode should be faster, if it is possible without unicode problems
    for row in df.itertuples():
        if row.col1:
            key1, val1 = string1, string2
            key1a, val1a = string1a, string2a
            print('SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a), file = f)
        if row.col2:
            key1, val1 = string1, string2
            key1a, val1a = string1a, string2a
            print('SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a), file = f)
        if row.col3:
            key1, val1 = string1, string2
            key1a, val1a = string1a, string2a
            print('SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a), file = f)
    f.close()

这可能是您可以进行的最低限度的优化。如果您更详细地准确描述了您在做什么，也许可以找到矢量化解决方案。

此外，不要将以上内容与 multiprocessing 一起使用.

此外，如所写，'SET {0} {1}\nSET {0} {1}'.format(key1, val1, key1a, val1a)永远是一样的。如果这些参数没有改变，那么只需在循环外执行字符串连接 一次 并在循环中重新使用整个字符串。

编辑:似乎你不能在上面那样做但是，给定:

This particular dataset has 6 million rows and 10 columns, mostly comprised of strings with a few float columns. The Redis keys are the strings and the float values are the Redis values in the key-value pair.

然后简单地 key1 = ''.join(row.col1, row.col4, row.col5, ...)不要使用 str和 +运算符，这是非常低效的，因为你暗示那些列已经是字符串了。如果您必须调用str在所有这些列上，使用 map(str, ...)

最后，如果您真的需要提高性能，请注意 row将是 namedtuple对象，是元组，您可以使用基于整数的索引而不是基于属性的标签访问，即 row[1]而不是 row.col1 (请注意，row[0] 将是 row.index，即索引)`，它应该更快(并且它会有所不同，因为您在每次迭代中对元组进行索引数十次并进行数百万次迭代)。