【神行百里】pandas查询加速之行索引篇

news2024/9/20 19:39:30

最近进行大数据处理的时候，发现我以前常用的pandas查询方法太慢了，太慢了，真是太慢了，查阅资料，遂发现了一种新的加速方法，能助力我飞上天，和太阳肩并肩，所以记录下来。

1. 场景说明与数据集构建

首先，阐述一下我们DataFrame查询应用场景，有两个csv表格文件；

表1中有两列，一列为id号（index），一列为名字（name）；
表2中有两列，一列为id号（index），一列为名字对应的内容（content）；
此外，表1和表2中的顺序是混乱的，每一行的index都不相同。

此时，你有一个想法，把name和content对应起来，在一张表中，有名字（name）就知道内容（content）了，我们先用下列代码构建这样一个数据集：

import pandas as pd
from pandas.core.frame import DataFrame
import random
from tqdm import tqdm

name = ["一念关山", "奉上", "立剑"]

content = ["一念月落，一念身错，一念关山难涉过。棋逢过客，执子者不问因果。", 
           "奉上，唯愿去踏破善恶之叠浪，一心并肩赴跌宕。", 
           "刺予我伤口，重生了魂魄，立剑的时候，疾骤的嘶吼。"]

result1 = []
result2 = []
for i in tqdm(range(10000000)):
    
    temp = random.randint(0, 2)
    xx1 = name[temp] + "__" + str(i)
    xx2 = "YNGS" + "__" + str(i)
    xx3 = content[temp] + "__" + str(i)
    
    result1.append([xx2, xx1])
    result2.append([xx2, xx3])
    
Frame1 = DataFrame(result1, columns=["index", "name"])
Frame1.to_csv("./test01.csv", index=False)

Frame2 = DataFrame(result2, columns=["index", "content"])
shuffled_df = Frame2.sample(frac=1, random_state=42)   # 打乱数据
shuffled_df.to_csv("./test02.csv", index=False)

代码构建了一个1000万行的数据表，当然了表中数据没有实际意义，仅供测试使用。

来瞅瞅我们生成的数据表格吧~
test01.csv：
在这里插入图片描述

test02.csv：
在这里插入图片描述

2. loc加速查询测试

2.1 原始方法

有了生成好的数据表，怎么做其实很简单，常规的做法如下虚线中的代码所示，对每一行循环，查询content表中与name表中index相同的行，再取content内容就可以了。

import pandas as pd
from pandas.core.frame import DataFrame
from tqdm import tqdm

data_name = pd.read_csv("./test01.csv", encoding="utf-8")
data_content = pd.read_csv("./test02.csv", encoding="utf-8")

name = data_name["name"].values.tolist()
index = data_name["index"].values.tolist()

result = []
for i in tqdm(range(len(index))):
    
    temp_index = index[i]
    
    # -----------------------------------------------------------------------
    temp_data = data_content[data_content["index"] == temp_index]
    temp_content = temp_data["content"].values.tolist()[0]
    # -----------------------------------------------------------------------
    
    result.append([temp_index, temp_content])
    
Frame = DataFrame(result, columns=["name", "content"])
Frame.to_csv("./finish.csv", index=False)

我们来看看运行速度，tqdm显示为725小时，完成这1000万行的数据需要那么久啊！这还能忍，人生苦短，python在手，我要加速 /(ㄒoㄒ)/~~

在这里插入图片描述

2.2 加速方法

其实 test01.csv 和 test02.csv 有相同的一列，即索引列（index），一般的dataframe的行索引默认为1、2、3…，我们可以设置index列为行索引，采用loc方法查询给定索引名对应的内容，修改代码如下虚线中所示：

import pandas as pd
from pandas.core.frame import DataFrame
from tqdm import tqdm

data_name = pd.read_csv("./test01.csv", encoding="utf-8")
data_content = pd.read_csv("./test02.csv", encoding="utf-8")

# -----------------------------------------------------------------------
data_content = data_content.set_index("index")  # 行索引设置
# -----------------------------------------------------------------------

name = data_name["name"].values.tolist()
index = data_name["index"].values.tolist()

result = []
for i in tqdm(range(len(index))):
    
    temp_index = index[i]
    
    # ---------------------------------------------------------
    temp_data = data_content.loc[temp_index]
    temp_content = [temp_data["content"]]
    # ---------------------------------------------------------

    result.append([temp_index, temp_content])
    
Frame = DataFrame(result, columns=["name", "content"])
Frame.to_csv("./finish.csv", index=False)