时间戳

首页 > 解决方案 > 使用 pandas 多索引进行搜索优化

python-3.x - 使用 pandas 多索引进行搜索优化

问题描述

我想知道是否有办法优化我正在做的搜索。我有一个多索引（3 级）数据框df，如下所示：

IndexID IndexDateTime IndexAttribute ColumnA ColumnB
   1      2015-02-05        8           A       B
   1      2015-02-05        7           C       D
   1      2015-02-10        7           X       Y

2015-02-10我的问题是我想知道给定日期是否有数据ColumnA，例如在给定日期（在这种情况IndexID下IndexAttribute为 5）之前的给定日期（在本例中为 5）是否有数据列，像这样：

IndexID IndexDateTime IndexAttribute ColumnA ColumnB NewColumn
   1      2015-02-05        8           A       B       -1
   1      2015-02-05        7           C       D       -1
   1      2015-02-10        7           X       Y        C

我想在我的数据框中的每一行上执行这个搜索，它有 1900 万行。我这样做的方式是：

df['NewColumn'] = df.apply(lambda r: get_data(df, r.IndexID, r.IndexDateTime , r.IndexAttribute , 5), axis=1)

在哪里get_data：

def get_data(df, IndexID, IndexDateTime , IndexAttribute , days_before):
    idx = pd.IndexSlice
    date = (IndexID - pd.to_timedelta(days_before, 'd'))
    try:
        res = df.loc[idx[IndexID, date, IndexAttribute ],'ColumnA']
        return res
    except KeyError:
        return -1

这非常慢，需要2个多小时。我想知道这是否可以是一种更快的方法。问题：

我不能换班，因为我不知道两行中间有多少数据。一些想法？谢谢！

标签： python-3.xpandassearchoptimizationmulti-index

解决方案

使用 numpy 可以非常快。您只需要将数据框中的列作为 numpy 数组进行迭代。希望能帮助到你：

%time
def myfunc(df, days_before=5):

     # Fill A column witH -1s
     result = -np.ones_like(df.values[:, -1:])

     # Slice the first 3 columns and shift the dates 
     # to get the index that we are looking for
     idx = np.array((df['IndexID'].values,
                     df['IndexDateTime'] - pd.to_timedelta(days_before, 'd'),
                     df['IndexAttribute'].values)).T

     # Look for days matching in the first 3 columns
     _idx_comp = df.values[:, :3][np.newaxis, :] == np.array(idx)[:, np.newaxis]

     # Get the index where there is a match
     # between the row of the dataframe and the desired searched rows
     idx_found = np.where(np.all(_idx_comp, axis=-1))

     # Assign the corresponding rows to its required value
     result[idx_found[0]] = df['ColumnA'].values[idx_found[-1]]

     return result

df.assign(NewColumn=myfunc(df))

CPU times: user 2 µs, sys: 1e+03 ns, total: 3 µs
Wall time: 5.96 µs

   IndexID IndexDateTime  IndexAttribute ColumnA ColumnB NewColumn
0        1    2015-02-05               8       A       B        -1
1        1    2015-02-05               7       C       D        -1
2        1    2015-02-10               7       X       Y         C

推荐阅读