使用 Python 和 Polars 对具有多个列的数据框进行排序

Question

我有一个数据排序问题，原始数据由三个“块”组成，其中包含一个“父”行和两个“子”行。最小工作示例如下所示：

import polars as pl
df_original = pl.DataFrame(
    {
        'Direction': ["Buy", "Sell", "Buy", "Sell", "Sell", "Buy"],
        'Order ID': [None, '123_1', '123_0', None, '456_1', '456_0'],
        'Parent Order ID': [123, None, None, 456, None, None],

    }
)

我想根据父行对这些进行排序。如果父订单是“买入”，那么下一行应该是“卖出”子订单，第三行应该是“买入”订单。

对于父“卖出”订单，需要先买入“买入”订单，然后再买入“卖出”订单。

我已经尝试过使用 Polars.sort()，但我缺少一段逻辑，无法找出它是什么。

最终结果应该是这样的：

df_sorted = pl.DataFrame(
    {
        'Direction': ["Buy", "Sell", "Buy", "Sell", "Buy", "Sell"],
        'Order ID': [None, '123_1', '123_0', None, '456_0', '456_1'],
        'Parent Order ID': [123, None, None, 456, None, None],

    }
)

Answer 1

如果我正确理解了问题，你想交换

"Buy"

/

"Sell"

的顺序。

此代码片段会产生您想要的输出。

df = pl.DataFrame(
    {
        'Direction': ["Buy", "Sell", "Buy", "Sell", "Sell", "Buy"],
        'Order ID': [None, '123_1', '123_0', None, '456_1', '456_0'],
        'Parent Order ID': [123, None, None, 456, None, None],

    }
)

consecutive = (pl.col("Direction").ne_missing(pl.col("Direction").shift()))

df.filter(consecutive).vstack(df.filter(~consecutive))

shape: (6, 3)
┌───────────┬──────────┬─────────────────┐
│ Direction ┆ Order ID ┆ Parent Order ID │
│ ---       ┆ ---      ┆ ---             │
│ str       ┆ str      ┆ i64             │
╞═══════════╪══════════╪═════════════════╡
│ Buy       ┆ null     ┆ 123             │
│ Sell      ┆ 123_1    ┆ null            │
│ Buy       ┆ 123_0    ┆ null            │
│ Sell      ┆ null     ┆ 456             │
│ Buy       ┆ 456_0    ┆ null            │
│ Sell      ┆ 456_1    ┆ null            │
└───────────┴──────────┴─────────────────┘