给定以下 Dataframe ：
|列_A|列_B|
| ------------ | ------------ |
| 1 | 1 |
| 1 | 2 |
| 1 | 3 |
| 2 | 4 |
| 2 | 5 |
| 2 | 6 |
| 3 | 7 |
| 3 | 8 |
| 3 | 9 |
如何将col_A中的一个组中的col_B的最大值分配给col_A下一个（按升序）组中的所有行？所需输出如下：
|col_A|col_B|max_col_B_from_prev_col_A_group|
| ------------ | ------------ | ------------ |
| 1 | 1 | |
| 1 | 2 | |
| 1 | 3 | |
| 2 | 4 | 3 |
| 2 | 5 | 3 |
| 2 | 6 | 3 |
| 3 | 7 | 6 |
| 3 | 8 | 6 |
| 3 | 9 | 6 |
我想我可以按组得到max，然后shift将结果转发，merge以某种方式将值返回到源 Dataframe 中。然而，我无法找到“某种方式”的解决方案：(
对如何做到这一点有什么想法吗？我一点也不执着于我的方法——只需要最快的解决方案，因为col_A中有1.5m行和80k离散值。
提前谢谢！
再现 Dataframe 的代码：

df = {
    'col_A': {0: 1, 1: 1, 2: 1, 3: 2, 4: 2, 5: 2, 6: 3, 7: 3, 8: 3},
    'col_B': {0: 1, 1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: 7, 7: 8, 8: 9},
}