Polars中如何为空间数据创建动态滚动窗口计算非空占比

需求说明 #

给定存储直线分布传感器读数的Polars DataFrame：

import polars as pl

df = pl.from_repr("""
┌─────────────────────┬───────────────────┬──────────┐
│ Location Start (KM) ┆ Location End (KM) ┆ Readings │
│ ---                 ┆ ---               ┆ ---      │
│ f64                 ┆ f64               ┆ f64      │
╞═════════════════════╪═══════════════════╪══════════╡
│ 1.0                 ┆ 1.1               ┆ 7.0      │
│ 1.1                 ┆ 1.23              ┆ null     │
│ 1.23                ┆ 1.3               ┆ 8.0      │
│ 1.3                 ┆ 1.34              ┆ null     │
│ 1.34                ┆ 1.4               ┆ null     │
│ 1.4                 ┆ 1.5               ┆ 5.0      │
│ 1.5                 ┆ 1.65              ┆ 6.0      │
└─────────────────────┴───────────────────┴──────────┘
""")

需要创建150米（0.15KM）的前瞻滚动窗口，计算窗口内读数的非空值占比，预期输出如下：

┌─────────────────────┬───────────────────┬──────────┬─────────────────────────────────┐
│ Location Start (KM) ┆ Location End (KM) ┆ Readings ┆ Rolling % Non Null Readings (%) │
│ ---                 ┆ ---               ┆ ---      ┆ ---                             │
│ f64                 ┆ f64               ┆ f64      ┆ i64                             │
╞═════════════════════╪═══════════════════╪══════════╪═════════════════════════════════╡
│ 1.0                 ┆ 1.1               ┆ 7.0      ┆ 67                              │
│ 1.1                 ┆ 1.23              ┆ null     ┆ 13                              │
│ 1.23                ┆ 1.3               ┆ 8.0      ┆ 47                              │
│ 1.3                 ┆ 1.34              ┆ null     ┆ 33                              │
│ 1.34                ┆ 1.4               ┆ null     ┆ 60                              │
│ 1.4                 ┆ 1.5               ┆ 5.0      ┆ 100                             │
│ 1.5                 ┆ 1.65              ┆ 6.0      ┆ 100                             │
└─────────────────────┴───────────────────┴──────────┴─────────────────────────────────┘

注：也可使用中心窗口，示例采用前瞻窗口。

Polars原生的group_by_dynamic仅支持时间类型数据，rolling_map为固定行数窗口，均不适合这种基于空间位置的动态窗口场景。目标是避免显式循环，利用Polars内置方法实现高效计算。

Pandas循环实现参考 #

现有基于Pandas的循环实现代码如下：

import pandas as pd
import numpy as np

df_pd = pd.DataFrame({
    'Location Start (KM)': [1, 1.1, 1.23, 1.3, 1.34, 1.4, 1.5],
    'Location End (KM)': [1.1, 1.23, 1.3, 1.34, 1.4, 1.5, 1.65],
    'Readings': [7, np.nan, 8, np.nan, np.nan, 5, 6]
}) # 示例DataFrame

def calculate_perc_non_null(row, window_length):
    '''
    计算前瞻滚动窗口内非空读数占比的朴素实现：
    row: DataFrame的单行数据（Pandas Series）
    window_length: 窗口长度，单位KM
    '''
    window_start = row['Location Start (KM)']
    window_end = window_start + window_length # 生成窗口边界
    # 筛选窗口内的读数行
    eligible_readings = df_pd.loc[(df_pd['Location Start (KM)'] >= window_start) & (df_pd['Location Start (KM)'] <= window_end)]
    # 计算空值覆盖的总长度
    nulls = eligible_readings.loc[eligible_readings['Readings'].isnull()].copy()
    # 截断超出窗口的空值段
    nulls.loc[nulls['Location End (KM)'] > window_end, 'Location End (KM)'] = window_end
    total_length_of_nulls = (nulls['Location End (KM)'] - nulls['Location Start (KM)']).sum()
    # 计算非空占比
    non_null_perc = 100 * (1 - total_length_of_nulls / window_length)
    return non_null_perc

df_pd['Rolling % Non Null Readings (%)'] = df_pd.apply(lambda x: calculate_perc_non_null(x, window_length=0.15), axis=1)

运行输出结果：

Location Start (KM)  Location End (KM)  Readings  Rolling % Non Null Readings (%)
0                 1.00               1.10       7.0                        66.666667
1                 1.10               1.23       NaN                        13.333333
2                 1.23               1.30       8.0                        46.666667
3                 1.30               1.34       NaN                        33.333333
4                 1.34               1.40       NaN                        60.000000
5                 1.40               1.50       5.0                       100.000000
6                 1.50               1.65       6.0                       100.000000

Polars高效实现方案 #

利用Polars的广播连接、条件过滤和聚合能力，避免显式循环，实现空间窗口的高效计算：

import polars as pl

window_length = 0.15

# 生成每个窗口的起始和结束位置
window_df = df.select(
    pl.col("Location Start (KM)").alias("window_start"),
    (pl.col("Location Start (KM)") + window_length).alias("window_end")
)

# 交叉连接匹配所有窗口与数据行，筛选符合空间范围的记录
result = window_df.join(df, how="cross") \
    .filter(
        pl.col("Location Start (KM)") >= pl.col("window_start"),
        pl.col("Location Start (KM)") <= pl.col("window_end")
    ) \
    # 截断超出窗口的空值段
    .with_columns(
        pl.when(pl.col("Location End (KM)") > pl.col("window_end"))
        .then(pl.col("window_end"))
        .otherwise(pl.col("Location End (KM)")).alias("truncated_end"),
        pl.col("Readings").is_null().alias("is_null")
    ) \
    # 按窗口分组计算空值覆盖的总长度
    .group_by("window_start", "window_end") \
    .agg(
        total_null_length=pl.when(pl.col("is_null"))
        .then(pl.col("truncated_end") - pl.col("Location Start (KM)"))
        .otherwise(0.0).sum()
    ) \
    # 计算非空占比并合并回原DataFrame
    .with_columns(
        pl.round(100 * (1 - pl.col("total_null_length") / window_length)).cast(pl.Int64)
        .alias("Rolling % Non Null Readings (%)")
    ) \
    .join(df, left_on="window_start", right_on="Location Start (KM)", how="left") \
    .select(
        "Location Start (KM)",
        "Location End (KM)",
        "Readings",
        "Rolling % Non Null Readings (%)"
    )

print(result)

性能对比结果 #

2022年11月30日补充：该Polars方案性能显著优于Pandas循环实现，具体对比数据如下：

• 小数据量（7行）：
  • Pandas循环实现：11.3 ms ± 1.95 ms/循环（7次运行，每次100循环）
  • Polars方案：527 µs ± 57.2 µs/循环（7次运行，每次1000循环），性能提升约25倍；
• 大数据量（原数据扩大1000倍）：
  • Pandas循环实现：10 s ± 580 ms/循环
  • Polars方案：12.2 ms ± 1.18 ms/循环，性能提升约820倍。

内容的提问来源于stack exchange，提问作者sjs