基于Polars高效计算大规模数据集的两点测地距离

核心优化思路 #

map_rows/逐行处理慢的根源是Python级别的循环开销，针对3000万行的规模，必须转向向量化操作或批次化处理，同时避免在循环中处理字符串解析这类可以提前向量化完成的步骤。

步骤1：先向量化解析坐标字符串 #

首先把接送点的坐标字符串列解析为单独的数值型经纬度列，这一步用Polars的内置字符串操作完成，速度远快于逐行处理：

import polars as pl

# 假设你的DataFrame有"pickup_coord"和"dropoff_coord"列，格式为"纬度,经度"
df = df.with_columns(
    # 拆分接送点坐标为结构化列
    pl.col("pickup_coord").str.split(",").list.to_struct(fields=["pickup_lat", "pickup_lng"]),
    pl.col("dropoff_coord").str.split(",").list.to_struct(fields=["dropoff_lat", "dropoff_lng"])
).unnest("pickup_coord", "dropoff_coord").with_columns(
    # 转换为浮点型
    pl.col(["pickup_lat", "pickup_lng", "dropoff_lat", "dropoff_lng"]).cast(pl.Float64)
)

步骤2：选择高效的距离计算方案 #

方案A：Haversine公式（速度优先，精度可接受）

如果可以接受Haversine公式的近似精度（误差约0.5%），用向量化的numpy实现，配合Polars的map_batches处理，速度比逐行快几十倍：

import numpy as np

def haversine_vec(lat1, lon1, lat2, lon2):
    # 转换为弧度
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    # 地球平均半径（公里）
    return c * 6371

# 用map_batches批量处理
df = df.with_columns(
    pl.struct(["pickup_lat", "pickup_lng", "dropoff_lat", "dropoff_lng"]).map_batches(
        lambda s: haversine_vec(
            s.struct.field("pickup_lat").to_numpy(),
            s.struct.field("pickup_lng").to_numpy(),
            s.struct.field("dropoff_lat").to_numpy(),
            s.struct.field("dropoff_lng").to_numpy()
        )
    ).alias("distance_km")
)

方案B：保留geopy的geodesic（精度优先）

如果必须用geopy.distance.geodesic的高精度计算，不要用map_rows，改用带batch_size参数的map_elements，减少Python与Polars的交互开销：

from geopy.distance import geodesic

def calc_distance_batch(batch):
    # batch是Polars的Struct列，转换为字典列表批量处理
    rows = batch.to_list()
    return [geodesic((row["pickup_lat"], row["pickup_lng"]), (row["dropoff_lat"], row["dropoff_lng"])).km for row in rows]

df = df.with_columns(
    pl.struct(["pickup_lat", "pickup_lng", "dropoff_lat", "dropoff_lng"]).map_elements(
        calc_distance_batch,
        return_dtype=pl.Float64,
        batch_size=10000  # 调整批次大小，根据内存情况优化
    ).alias("distance_km")
)

注：batch_size建议根据你的内存容量调整，越大批次开销越低，但不要超过内存承受范围。

进阶优化：用numba加速geodesic调用（可选）

如果需要进一步提升geodesic的计算速度，可以用numba包装底层计算逻辑（精度与geopy的geodesic接近）：

from numba import jit
import numpy as np

@jit(nopython=True)
def geodesic_numba(lat1, lon1, lat2, lon2):
    # 基于WGS84椭球的Vincenty公式实现
    a = 6378137.0
    f = 1/298.257223563
    b = a * (1 - f)
    
    lat1_rad = np.radians(lat1)
    lon1_rad = np.radians(lon1)
    lat2_rad = np.radians(lat2)
    lon2_rad = np.radians(lon2)
    
    dlon = lon2_rad - lon1_rad
    U1 = np.arctan((1 - f) * np.tan(lat1_rad))
    U2 = np.arctan((1 - f) * np.tan(lat2_rad))
    
    sinU1 = np.sin(U1)
    cosU1 = np.cos(U1)
    sinU2 = np.sin(U2)
    cosU2 = np.cos(U2)
    
    lambda_val = dlon
    lambda_prev = 0.0
    iter_limit = 100
    for _ in range(iter_limit):
        sinLambda = np.sin(lambda_val)
        cosLambda = np.cos(lambda_val)
        sinSigma = np.sqrt((cosU2 * sinLambda)**2 + (cosU1 * sinU2 - sinU1 * cosU2 * cosLambda)**2)
        if sinSigma == 0:
            return 0.0  # 同一点
        cosSigma = sinU1 * sinU2 + cosU1 * cosU2 * cosLambda
        sigma = np.arctan2(sinSigma, cosSigma)
        sinAlpha = cosU1 * cosU2 * sinLambda / sinSigma
        cosSqAlpha = 1 - sinAlpha**2
        cos2SigmaM = cosSigma - 2 * sinU1 * sinU2 / cosSqAlpha
        if np.isnan(cos2SigmaM):
            cos2SigmaM = 0.0  # 赤道附近情况
        C = f / 16 * cosSqAlpha * (4 + f * (4 - 3 * cosSqAlpha))
        lambda_prev = lambda_val
        lambda_val = dlon + (1 - C) * f * sinAlpha * (sigma + C * sinSigma * (cos2SigmaM + C * cosSigma * (-1 + 2 * cos2SigmaM**2)))
        if abs(lambda_val - lambda_prev) < 1e-12:
            break
    
    uSq = cosSqAlpha * (a**2 - b**2) / (b**2)
    A = 1 + uSq / 16384 * (4096 + uSq * (-768 + uSq * (320 - 175 * uSq)))
    B = uSq / 1024 * (256 + uSq * (-128 + uSq * (74 - 47 * uSq)))
    deltaSigma = B * sinSigma * (cos2SigmaM + B / 4 * (cosSigma * (-1 + 2 * cos2SigmaM**2) - B / 6 * cos2SigmaM * (-3 + 4 * sinSigma**2) * (-3 + 4 * cos2SigmaM**2)))
    s = b * A * (sigma - deltaSigma)
    return s / 1000  # 转换为公里

# 用map_batches调用numba加速函数
df = df.with_columns(
    pl.struct(["pickup_lat", "pickup_lng", "dropoff_lat", "dropoff_lng"]).map_batches(
        lambda s: geodesic_numba(
            s.struct.field("pickup_lat").to_numpy(),
            s.struct.field("pickup_lng").to_numpy(),
            s.struct.field("dropoff_lat").to_numpy(),
            s.struct.field("dropoff_lng").to_numpy()
        )
    ).alias("distance_km")
)

性能对比参考 #

• 原map_rows：3000万行约1.5小时
• 方案A（Haversine向量化）：预计几分钟内完成
• 方案B（geodesic+batch）：预计10-20分钟完成
• 方案C（numba高精度）：预计5-10分钟完成

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