Pandas 优化 interpolation/counting 算法

Question

我有一堆数据（1000 万条记录）分解为标识符、位置和日期。我想找到在整个日期集中任何标识符从某个位置 A 移动到其他位置 B 的次数。任何标识符都可能没有所有可能日期的位置。当标识符没有记录位置时，应将其视为该日期的实际 'unknown' 位置。

这是一些可复制的假数据...

import numpy as np
import pandas as pd
import datetime

base = datetime.date.today()
num_days = 50
dates = np.array([base - datetime.timedelta(days=x) for x in range(num_days-1, -1, -1)])
ids = np.arange(50)
mi = pd.MultiIndex.from_product([ids, dates])
locations = np.array([chr(x) for x in 97 + np.random.randint(26, size=len(mi))])

s = pd.Series(locations, index=mi)
mask = np.random.rand(len(mi)) > .5
s[mask] = np.nan
s = s.dropna()

我最初的想法是创建一个数据框并使用布尔 masking/vectorized 运算来解决这个问题

df = s.unstack(0).fillna('unknown')

显然我的数据稀疏到足以导致 MemoryError（来自所有因取消堆叠而产生的额外条目）。

我目前的工作解决方案如下

def series_fn(s):
    s = s.reindex(pd.date_range(s.index.levels[1].min(), s.index.levels[1].max()), level=-1).fillna('unknown')
    mask_prev = (s != s.shift(-1))[:-1]
    mask_next = (s != s.shift())[1:]
    s_prev = s[:-1][mask_prev]
    s_next = s[1:][mask_next]
    s_tup = pd.Series(list(zip(s_prev, s_next)))
    return s_tup.value_counts()

result_per_id = s.groupby(level=0).apply(series_fn)
result = result_per_id.sum(level=-1)

result看起来像

(a, b)    1
(a, c)    5
(a, e)    3
(a, f)    3
(a, g)    3
(a, h)    3
(a, i)    1
(a, j)    1
(a, k)    2
(a, l)    2
...

我的所有数据大约需要 5 个小时。有谁知道这样做有什么更快的方法吗？谢谢！

Answer 1

嗯，我想我应该转置数据...好吧，这是一个相对简单的修复。而不是使用 groupby 和应用，

s = s.reorder_levels(['date', 'id'])
s = s.sortlevel(0)

results = []

for i in range(len(s.index.levels[0])-1):
    t = time.time()
    s0 = s.loc[s.index.levels[0][i]]
    s1 = s.loc[s.index.levels[0][i+1]]
    df = pd.concat((s0, s1), axis=1)

    # Note: this is slower than the line above
    # df = s.loc[s.index.levels[0][0:2], :].unstack(0)

    df = df.fillna('unknown')
    mi = pd.MultiIndex.from_arrays((df.iloc[:, 0], df.iloc[:, 1]))
    s2 = pd.Series(1, mi)
    res = s2.groupby(level=[0, 1]).apply(np.sum)
    results.append(res)
    print(time.time() - t)

results = pd.concat(results, axis=1)

仍然不清楚为什么注释掉的部分的时间是上面三行的三倍。

Pandas 优化 interpolation/counting 算法

Pandas optimizing an interpolation/counting algorithm

python

optimization

pandas