用Polars模拟没有循环的多节点服务队列的惯用方法?

我正在模拟一个先进先出的服务队列,该队列具有参数化数量的服务节点。我想创建一个用户定义的函数或Polars表达式(在Python中或作为Rust插件也可以),用于计算每个对象的服务的模拟开始和结束时间。是否有使用Polars表达式的惯用方法或最推荐的方法?

我的polars DataFrame的每一行都对应于一个在特定时间到达队列的对象,一旦其服务时间开始,就需要特定分钟的服务时间。我的输入数据帧如下所示:

import polars as pl
df = pl.DataFrame({
    'id': [1,2,3,4],
    'servicing_time_requirement': [30, 5, 30, 5],
    'arrival_time': [0, 15, 16, 17],
})
df
shape: (5, 3)
┌─────┬────────────────────────────┬──────────────┐
│ id  ┆ servicing_time_requirement ┆ arrival_time │
│ --- ┆ ---                        ┆ ---          │
│ i64 ┆ i64                        ┆ i64          │
╞═════╪════════════════════════════╪══════════════╡
│ 1   ┆ 30                         ┆ 0            │
│ 2   ┆ 5                          ┆ 15           │
│ 3   ┆ 30                         ┆ 16           │
│ 4   ┆ 5                          ┆ 17           │
└─────┴────────────────────────────┴──────────────┘

例如,如果服务节点的数量=5(一次最多可以处理五个对象),则所需的输出为:

df.pipe(simulateQueue, nodes=5)
shape: (4, 5)
┌─────┬────────────────────────────┬──────────────┬──────────────────────┬────────────────────┐
│ id  ┆ servicing_time_requirement ┆ arrival_time ┆ servicing_start_time ┆ servicing_end_time │
│ --- ┆ ---                        ┆ ---          ┆ ---                  ┆ ---                │
│ i64 ┆ i64                        ┆ i64          ┆ i64                  ┆ i64                │
╞═════╪════════════════════════════╪══════════════╪══════════════════════╪════════════════════╡
│ 1   ┆ 30                         ┆ 0            ┆ 0                    ┆ 30                 │
│ 2   ┆ 5                          ┆ 15           ┆ 15                   ┆ 20                 │
│ 3   ┆ 30                         ┆ 16           ┆ 16                   ┆ 46                 │
│ 4   ┆ 5                          ┆ 17           ┆ 17                   ┆ 22                 │
└─────┴────────────────────────────┴──────────────┴──────────────────────┴────────────────────┘

但是,如果服务节点的数量=2(一次只处理两个对象),则在本例中需要等待,因此所需的输出为:

df.pipe(simulateQueue, nodes=2)
shape: (4, 5)
┌─────┬────────────────────────────┬──────────────┬──────────────────────┬────────────────────┐
│ id  ┆ servicing_time_requirement ┆ arrival_time ┆ servicing_start_time ┆ servicing_end_time │
│ --- ┆ ---                        ┆ ---          ┆ ---                  ┆ ---                │
│ i64 ┆ i64                        ┆ i64          ┆ i64                  ┆ i64                │
╞═════╪════════════════════════════╪══════════════╪══════════════════════╪════════════════════╡
│ 1   ┆ 30                         ┆ 0            ┆ 0                    ┆ 30                 │
│ 2   ┆ 5                          ┆ 15           ┆ 15                   ┆ 20                 │
│ 3   ┆ 30                         ┆ 16           ┆ 20                   ┆ 50                 │
│ 4   ┆ 5                          ┆ 17           ┆ 30                   ┆ 35                 │
└─────┴────────────────────────────┴──────────────┴──────────────────────┴────────────────────┘

我当前循环浏览的行 df ,计算 servicing_start_time 使用前几行的值 servicing_end_time ,但似乎必须有更有效的方法。Polars滚动表达式是我想到的一个想法,但它们似乎不太适合这个用例。

Python中的当前实现:

nodes = 2
arrival_times = df.get_column("arrival_times")
servicing_end_times = df.get_column("servicing_end_time")
servicing_time_requirements = df.get_column("servicing_time_requirement")
for i in range(0, servicing_end_times.len()):
    if servicing_end_times[i] is not None: continue
    next_done = servicing_end_times.filter(
        (servicing_end_times.is_not_null()) & 
        (servicing_end_times.rank(method='ordinal', descending = True).eq(nodes)))
    if next_done.len() == 0: next_done = arrival_times[i]
    else: next_done = max(next_done[0], arrival_times[i])
    servicing_end_times[i] = next_done + servicing_time_requirements[i]