本文介绍如何在 Polars 中高效地为多架飞机（按 plane_no 分组）计算每条 GPS 轨迹点的航向角，通过 shift(-1).over() 替代不可组合的 rolling().over()，实现分组内前后坐标对齐与向量化计算。
本文介绍如何在 polars 中高效地为多架飞机（按 `plane_no` 分组）计算每条 gps 轨迹点的航向角，通过 `shift(-1).over()` 替代不可组合的 `rolling().over()`，实现分组内前后坐标对齐与向量化计算。
在轨迹分析场景中，常需基于连续 GPS 坐标（经纬度）计算运动方向——即航向角（bearing），其定义为从当前点指向下一时刻点的正北顺时针夹角（单位：度，范围 0–360°）。Polars 原生不支持 rolling(...).over(...) 的嵌套操作，因此无法直接对每个 plane_no 分区执行“滚动取相邻两行”的操作。但借助 shift(-1).over("plane_no")，我们可优雅、高效地实现等效逻辑：为每个分组内的每一行，提取其下一行的经纬度值，从而构建 (lat₁, lon₁) → (lat₂, lon₂) 的坐标对。
✅ 推荐实现步骤（清晰、高效、可扩展）
-
排序保障时序性：首先按
plane_no和timestamp双重排序，确保同一飞机的轨迹点严格按时间升序排列； -
跨行取值（分组内）：使用
shift(-1).over("plane_no")为每行生成下一行的gps_lat和gps_lon，记为gps_lat_next/gps_lon_next； -
批量计算 bearing：将四列打包为
struct，用map_elements调用自定义函数（注意：生产环境建议改用apply+ 纯表达式或预编译函数提升性能）； - 清理临时列：删除中间辅助列，保留结果。
以下是完整可运行示例：
import polars as pl
import math
def calculate_bearing(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
"""计算两点间航向角（单位：度），返回 None 若任一坐标为 null"""
if None in (lat1, lon1, lat2, lon2):
return None
# 转弧度
lat1, lon1, lat2, lon2 = map(math.radians, [lat1, lon1, lat2, lon2])
dlon = lon2 - lon1
y = math.sin(dlon) * math.cos(lat2)
x = math.cos(lat1) * math.sin(lat2) - math.sin(lat1) * math.cos(lat2) * math.cos(dlon)
bearing = math.degrees(math.atan2(y, x))
return (bearing + 360) % 360
# 示例数据
df = pl.DataFrame({
"event_uid": [1, 2, 3, 4, 5, 6],
"plane_no": ["A", "A", "A", "B", "B", "B"],
"timestamp": [1, 2, 3, 1, 2, 3],
"gps_lat": [37.7749, 37.7750, 37.7751, 40.7128, 40.7130, 40.7132],
"gps_lon": [-122.4194, -122.4195, -122.4196, -74.0060, -74.0062, -74.0064]
})
# 核心计算链式表达式
result = (
df.sort(["plane_no", "timestamp"])
.with_columns([
pl.col("gps_lat").shift(-1).over("plane_no").alias("gps_lat_next"),
pl.col("gps_lon").shift(-1).over("plane_no").alias("gps_lon_next")
])
.with_columns(
pl.struct(["gps_lat", "gps_lon", "gps_lat_next", "gps_lon_next"])
.map_elements(
lambda row: calculate_bearing(
row["gps_lat"], row["gps_lon"],
row["gps_lat_next"], row["gps_lon_next"]
),
return_dtype=pl.Float64
)
.alias("bearing")
)
.drop(["gps_lat_next", "gps_lon_next"])
)
print(result)
⚠️ 注意事项与优化建议
-
map_elements性能提示：该函数会触发 Python 解释器循环，在大数据集上较慢。若性能关键，推荐：- 使用
pl.all().apply(...)（Polars 0.20+ 支持更优的 UDF 向量化）； - 或将
calculate_bearing重写为纯 Polars 表达式（利用pl.arctan2,pl.sin,pl.cos,pl.radians等），彻底避免 Python 回调；
- 使用
-
边界处理：每组末尾行的
bearing自动为null（因无后续点），符合语义，无需额外过滤； -
空值安全：
shift(-1)在组尾产生null，calculate_bearing内已做防御性检查，确保健壮； -
扩展性：如需同时计算距离、速度等指标，可复用相同
shift列，避免重复分组操作。
该方案摒弃了低效的 row_index + rolling 组合，充分利用 Polars 的分组窗口原语，兼具可读性、正确性与高性能，是处理多实体时序轨迹特征工程的典型范式。
