使用替代方法加载Pandas DataFrame时性能下降的原因排查

我需要处理约8GB的大型CSV文件，将其转为DataFrame后拆分为字典——以日期为键，每日数据为小型DataFrame，以此提升后续模拟的查询速度。

为了优化拆分速度，我重写了处理流程：不再先合并所有块为单个大DataFrame再拆分，而是将CSV分块存储在列表中，记录每个块的最小/最大日期，直接从包含目标日期的块中提取数据合并到字典。结果拆分时间从40分钟骤降到15秒，但模拟单周期耗时从2.5秒升至3.5秒。

这不是第一次出现类似问题：数月前我从原CSV中提取最近两年数据生成800MB的小型CSV，加载速度更快，但模拟耗时从12秒涨到20秒。两张表的列类型、数量完全一致，但小型表的模拟速度始终慢于原完整CSV的处理速度，且该现象早于拆分字典的操作。

新旧流程对比 #

旧流程 #

• 分块（20000行/块）加载CSV
• 合并所有块为单个DataFrame
• 拆分每日数据到字典
• 执行模拟

新流程 #

• 分块（20000行/块）加载CSV
• 直接从含目标日期的块中拆分数据到字典
• 执行模拟

关键差异 #

新流程生成的字典占用内存10.5GB，旧流程为12.2GB——内存占用更小，但模拟速度反而更慢。

旧方法代码 #

chunkList = [chunk for chunk in pandas.read_csv(f'{SOURCE}_candles.csv', chunksize=20000, low_memory=False)]
candles = pandas.concat(chunkList)
del chunkList

print("Got candles")

# Data type insurance...

candles['tm_uuid'] = candles['tm_uuid'].astype("string")
candles['ticker'] = candles['ticker'].astype("string")
candles['price_date'] = pandas.to_datetime(candles['price_date']).dt.date

# Specify the start date; don't want to waste time looking at days we don't care about

candles = candles.loc[(candles['price_date'] >= simul_date), :]

# Now get the active dates to use as keys for the dictionary.

activeDates = candles[CandleColumn.PRICE_DATE.value].unique().tolist()

print("Got active dates")

timerTest = time.time()
candlesDict = {}

# Now breakout into the dictionary

for date in activeDates:
    candlesDict[date] = candles.loc[(candles[CandleColumn.PRICE_DATE.value] == date), :]
    candlesDict[date].set_index(CandleColumn.TM_UUID.value, drop=False, inplace=True)
    print(f"Finished {date}")

print(f"Done in {time.time() - timerTest}")
print(f"Candles dictionary uses {sum([sys.getsizeof(candlesDict[date]) for date in activeDates]) / 1024}kb")

# We no longer need the candles table, so clear it from memory!
del candles

新方法代码 #

candles = [chunk for chunk in pandas.read_csv(f'{SOURCE}_candles.csv', chunksize=20000, low_memory=False)]

# Data type insurance...

for table in candles:
    table['tm_uuid'] = table['tm_uuid'].astype("string")
    table['ticker'] = table['ticker'].astype("string")
    table['price_date'] = pandas.to_datetime(table['price_date']).dt.date

# Specify the start date; don't want to waste time looking at days we don't care about

# todo this method is bad practice, create new list and copy. Maybe inplace drop of specified rows?
# Removing this does not improve processing speed...
for index, table in enumerate(candles):
    candles[index] = table.loc[(table['price_date'] >= simul_date), :]

# Now get the active dates to use as keys for the dictionary.

activeDatesUnflattened = [table[CandleColumn.PRICE_DATE.value].unique().tolist() for table in candles]
# Not much point to the .unique() call anymore eh?
activeDates = list(set(list(chain.from_iterable(activeDatesUnflattened))))
activeDates.sort()
del activeDatesUnflattened

timerTest = time.time()
candlesDict = {}
tableDates = []

# Capture the min and max dates for each table in the list

for table in candles:
    tableDates.append((table[CandleColumn.PRICE_DATE.value].min(), table[CandleColumn.PRICE_DATE.value].max()))

# Now breakout into the dictionary

for date in activeDates:
    for index, table in enumerate(candles):
        if date >= tableDates[index][0] and date <= tableDates[index][1]:
            # The date is within the bounds of the table
            if date in candlesDict.keys():
                # If there is a table there, we need to concat the two and make sure the index
                candlesDict[date] = pandas.concat([(table.loc[(table[CandleColumn.PRICE_DATE.value] == date), :]).set_index(CandleColumn.TM_UUID.value, drop=False, inplace=True), candlesDict[date]])
                # todo Faster way to preserve index?
            else:
                candlesDict[date] = table.loc[(table[CandleColumn.PRICE_DATE.value] == date), :]
                candlesDict[date].set_index(CandleColumn.TM_UUID.value, drop=False, inplace=True)
        else:
            continue
    print(f"Finished {date}")

print(f"Done in {time.time() - timerTest}")
print(f"Candles dictionary uses {sum([sys.getsizeof(candlesDict[date]) for date in activeDates]) / 1024}kb")

# We no longer need the candles table, so clear it from memory!
del candles

内容的提问来源于stack exchange，提问作者Allergic