Polars 窗口函数中的聚合 - 如何根据其他列的聚合选择顶部值

Question

我有一个关于海运的大型数据集，其中包括 bol、voyage_id、承运商和总集装箱 (teus) 列，与此类似：

lf = pl.LazyFrame({
    'bol_id':(1,2,3,4,5,6,7,8,9),
    'voyage_id':(1,1,1,2,2,2,3,3,3),
    'carrier_scac':('mscu', 'mscu', 'hpld', 'hpld', 'hpld', 'hpld', 'ever', 'mscu', 'ever'),
    'teus':(20, 40, 5, 10, 25, 20, 5, 45, 5)
})
print(lf.collect())
┌────────┬───────────┬──────────────┬──────┐
│ bol_id ┆ voyage_id ┆ carrier_scac ┆ teus │
│ ---    ┆ ---       ┆ ---          ┆ ---  │
│ i64    ┆ i64       ┆ str          ┆ i64  │
╞════════╪═══════════╪══════════════╪══════╡
│ 1      ┆ 1         ┆ mscu         ┆ 20   │
│ 2      ┆ 1         ┆ mscu         ┆ 40   │
│ 3      ┆ 1         ┆ hpld         ┆ 5    │
│ 4      ┆ 2         ┆ hpld         ┆ 10   │
│ 5      ┆ 2         ┆ hpld         ┆ 25   │
│ 6      ┆ 2         ┆ hpld         ┆ 20   │
│ 7      ┆ 3         ┆ ever         ┆ 5    │
│ 8      ┆ 3         ┆ mscu         ┆ 45   │
│ 9      ┆ 3         ┆ ever         ┆ 5    │
└────────┴───────────┴──────────────┴──────┘

对于每个航次，我都想找到 teus 总和最高的承运人。我可以通过 group_by 后跟 join 来完成此操作，但我想使用窗口函数来完成此操作，并且无法完全弄清楚 Polars (0.20) 中的语法/逻辑。

当前工作功能：

def add_primary_carrier(lf):
    lf2 = (
        lf
        #select relevant cols
        .select('voyage_id', 'carrier_scac', 'teus')
        #ignore bols with missing data
        .drop_nulls()
        #sum up TEUs by voyage and carrier
        .group_by('voyage_id', 'carrier_scac')
        .agg(pl.col('teus').sum().alias('sum_teus'))
        #choose the carrier with the most TEUs on each voyage
        .sort('sum_teus', descending=True)
        .group_by('voyage_id')
        .agg(pl.col('carrier_scac').first().alias('primary_scac'))
    )
    lf = (
        #add primary scac column to main lf
        lf.join(lf2, how='left', on='voyage_id')
    )

但是窗口函数似乎会更干净（而且可能占用更少的资源）。比如：

def add_primary_carrier_window(lf):
    lf = (
        lf.with_columns(
            pl.col('carrier_scac')
            .sort_by(pl.col('teus').sum().over('carrier_scac'), descending=True)
            .drop_nulls().first()
            .over('voyage_id')
            .alias('primary_scac')
        )
    )
    return lf

但是该函数会抛出“聚合中不允许的窗口表达式”操作错误。

预先感谢您的帮助！

预期输出：

┌────────┬───────────┬──────────────┬──────┬──────────────┬──────────────┐
│ bol_id ┆ voyage_id ┆ carrier_scac ┆ teus ┆ primary_scac ┆ shared_cargo │
│ ---    ┆ ---       ┆ ---          ┆ ---  ┆ ---          ┆ ---          │
│ i64    ┆ i64       ┆ str          ┆ i64  ┆ str          ┆ bool         │
╞════════╪═══════════╪══════════════╪══════╪══════════════╪══════════════╡
│ 1      ┆ 1         ┆ mscu         ┆ 20   ┆ mscu         ┆ false        │
│ 2      ┆ 1         ┆ mscu         ┆ 40   ┆ mscu         ┆ false        │
│ 3      ┆ 1         ┆ hpld         ┆ 5    ┆ mscu         ┆ true         │
│ 4      ┆ 2         ┆ hpld         ┆ 10   ┆ hpld         ┆ false        │
│ 5      ┆ 2         ┆ hpld         ┆ 25   ┆ hpld         ┆ false        │
│ 6      ┆ 2         ┆ hpld         ┆ 20   ┆ hpld         ┆ false        │
│ 7      ┆ 3         ┆ ever         ┆ 5    ┆ mscu         ┆ true         │
│ 8      ┆ 3         ┆ mscu         ┆ 45   ┆ mscu         ┆ false        │
│ 9      ┆ 3         ┆ ever         ┆ 5    ┆ mscu         ┆ true         │
└────────┴───────────┴──────────────┴──────┴──────────────┴──────────────┘

Answer 1

跟踪器上存在一些与此相关的问题，例如https://github.com/pola-rs/polars/issues/14361

您基本上必须在单独的

.over

调用中从每个

.with_columns

“聚合”创建一列，因为它们不能“嵌套”。

(df.with_columns( 
    pl.col('teus')
      .sum()
      .over('voyage_id', 'carrier_scac')
      .alias('sum_teus')
   )
   .with_columns(
      pl.col('carrier_scac') 
        .sort_by('sum_teus', descending=True)
        .first()
        .over('voyage_id')
        .alias('primary_scac')
   )
)

shape: (9, 6)
┌────────┬───────────┬──────────────┬──────┬──────────┬──────────────┐
│ bol_id ┆ voyage_id ┆ carrier_scac ┆ teus ┆ sum_teus ┆ primary_scac │
│ ---    ┆ ---       ┆ ---          ┆ ---  ┆ ---      ┆ ---          │
│ i64    ┆ i64       ┆ str          ┆ i64  ┆ i64      ┆ str          │
╞════════╪═══════════╪══════════════╪══════╪══════════╪══════════════╡
│ 1      ┆ 1         ┆ mscu         ┆ 20   ┆ 60       ┆ mscu         │
│ 2      ┆ 1         ┆ mscu         ┆ 40   ┆ 60       ┆ mscu         │
│ 3      ┆ 1         ┆ hpld         ┆ 5    ┆ 5        ┆ mscu         │
│ 4      ┆ 2         ┆ hpld         ┆ 10   ┆ 55       ┆ hpld         │
│ 5      ┆ 2         ┆ hpld         ┆ 25   ┆ 55       ┆ hpld         │
│ 6      ┆ 2         ┆ hpld         ┆ 20   ┆ 55       ┆ hpld         │
│ 7      ┆ 3         ┆ ever         ┆ 5    ┆ 10       ┆ mscu         │
│ 8      ┆ 3         ┆ mscu         ┆ 45   ┆ 45       ┆ mscu         │
│ 9      ┆ 3         ┆ ever         ┆ 5    ┆ 10       ┆ mscu         │
└────────┴───────────┴──────────────┴──────┴──────────┴──────────────┘