Question

我现在将我的一个Spark项目迁移到Pandas，我遇到的一个问题是在每个组中的pandas中实现滚动和函数。

假设我有：

key time value
A   1    10
A   2    20
A   4    30
A   8    10
B   1    15
B   2    30
B   3    15

我想首先按key进行分组，然后计算类似滚动的函数，该函数填充value相对于time的中间时间段。例如，如果窗口大小为2，我所需的输出将为：

key time output
A   1    10    
A   2    30 (10+20)    
A   3    30 (10+20+0)
A   4    50 (20+0+30) 
A   5    30 (0+30+0)
A   6    30 (the same as above)
A   8    10 (7 is 0 so it is omitted)
A   9    10
A   10   10
B   1    15
B   2    45
B   3    60
B   4    45
B   5    15

我在group by和apply度过了整整一个下午。有一个聪明的方法来做到这一点？在火花中，我可以collect_list然后select使用udf到time和value来做到这一点，但是大熊猫的想法与Spark不同。

谢谢！

Answer 1

IIUIC，这是一种方法。

首先使用新时间窗创建dff

In [1458]: dff = (df.groupby('key')
                    .apply(lambda x: pd.Series(range(x.time.min(), x.time.max()+2)))
                    .reset_index(name='time').drop('level_1', 1))
In [1459]: dff
Out[1459]:
  key  time
0   A     1
1   A     2
2   A     3
3   A     4
4   A     5
5   B     1
6   B     2
7   B     3
8   B     4

然后，在早期df上合并，key上的groupby，使用value

添加当前和shift

In [1460]: dff.assign(ouput=dff.merge(df, how='left')
                               .fillna(0).groupby('key').value
                               .apply(lambda x: x+x.shift().fillna(0)))
Out[1460]:
  key  time  ouput
0   A     1   10.0
1   A     2   30.0
2   A     3   20.0
3   A     4   30.0
4   A     5   30.0
5   B     1   15.0
6   B     2   45.0
7   B     3   45.0
8   B     4   15.0

Answer 2

解决此问题的一种方法是在某个组中生成time的整个范围，然后再执行dropna

def make_rolling_sum(key, group, window):
    time_range = range(group['time'].min(), group['time'].max() + 1 + window, )
    df = group.set_index('time').reindex(time_range)
    result = df.rolling(window, min_periods=1).sum()
    result['key'] = key  # the reindex drops a lot of NaN's in this column
    return result.dropna().reset_index().reindex(columns=group.columns)

window = 3
pd.concat((make_rolling_sum(*group, window)for group in df.groupby('key')), ignore_index=True)

产量

    key     time    value
0   A   1   10.0
1   A   2   30.0
2   A   3   30.0
3   A   4   50.0
4   A   5   30.0
5   A   6   30.0
6   A   8   10.0
7   A   9   10.0
8   A   10  10.0
9   B   1   15.0
10  B   2   45.0
11  B   3   60.0
12  B   4   45.0
13  B   5   15.0

滚动总和填充分组DataFrame上的中间值

2 个答案: