如何加速基于两个DataFrames的2公里范围近邻点检索

问题背景 #

我有两个包含经纬度的DataFrame（df与locations_df），需要找出每个locations_df行数据2公里范围内的df点集。已尝试向量化实现，但当locations_df行数超过1000时，运行速度明显变慢，寻求高效优化方案。

原实现代码：

import pandas as pd
import numpy as np

def select_points_for_multiple_locations_vectorized(df, locations_df, radius_km):
    R = 6371  # Earth's radius in kilometers

    # Convert degrees to radians
    df_lat_rad = np.radians(df['latitude'].values)[:, np.newaxis]
    df_lon_rad = np.radians(df['longitude'].values)[:, np.newaxis]
    loc_lat_rad = np.radians(locations_df['lat'].values)
    loc_lon_rad = np.radians(locations_df['lon'].values)

    # Haversine formula (vectorized)
    dlat = df_lat_rad - loc_lat_rad
    dlon = df_lon_rad - loc_lon_rad
    a = np.sin(dlat/2)**2 + np.cos(df_lat_rad) * np.cos(loc_lat_rad) * np.sin(dlon/2)**2
    c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1-a))
    distances = R * c

    # Create a mask for points within the radius
    mask = distances <= radius_km

    # Get indices of True values in the mask
    indices = np.where(mask)

    result = pd.concat([df.iloc[indices[0]].reset_index(drop=True), locations_df.iloc[indices[1]].reset_index(drop=True)], axis=1)

    return result

def random_lat_lon(n=1, lat_min=-10., lat_max=10., lon_min=-5., lon_max=5.):
    """
    this code produces an array with pairs lat, lon
    """
    lat = np.random.uniform(lat_min, lat_max, n)
    lon = np.random.uniform(lon_min, lon_max, n)

    return np.array(tuple(zip(lat, lon)))

df = pd.DataFrame(random_lat_lon(n=10000000), columns=['latitude', 'longitude'])
locations_df = pd.DataFrame(random_lat_lon(n=20), columns=['lat', 'lon'])

result = select_points_for_multiple_locations_vectorized(df, locations_df, radius_km=2)

优化方案1：使用BallTree空间索引 #

原向量化方案的时间复杂度是O(N*M)（N是df行数，M是locations_df行数），当M超过1000、N达千万级时，计算量会爆炸。改用sklearn.neighbors.BallTree构建空间索引，可将复杂度降至O(M log N)，大幅提升速度。

实现代码 #

import pandas as pd
import numpy as np
from sklearn.neighbors import BallTree

def select_points_with_balltree(df, locations_df, radius_km):
    EARTH_RADIUS = 6371  # 地球半径（公里）
    # 转换经纬度为弧度，BallTree的haversine metric要求输入弧度值
    df_coords = np.radians(df[['latitude', 'longitude']].values)
    loc_coords = np.radians(locations_df[['lat', 'lon']].values)
    
    # 构建BallTree，使用haversine计算球面距离
    tree = BallTree(df_coords, metric='haversine')
    
    # 将查询半径转换为弧度：弧度 = 距离 / 地球半径
    radius_rad = radius_km / EARTH_RADIUS
    
    # 批量查询每个location的邻域点索引
    indices_list = tree.query_radius(loc_coords, r=radius_rad)
    
    # 拼接结果
    result_dfs = []
    for loc_idx, df_indices in enumerate(indices_list):
        if len(df_indices) == 0:
            continue
        # 获取当前location的信息并重复对应匹配点数量
        loc_row = locations_df.iloc[[loc_idx]].reset_index(drop=True)
        repeated_loc = pd.concat([loc_row]*len(df_indices), ignore_index=True)
        # 获取匹配的df点
        matched_df = df.iloc[df_indices].reset_index(drop=True)
        # 拼接成结果行
        result_dfs.append(pd.concat([matched_df, repeated_loc], axis=1))
    
    return pd.concat(result_dfs, ignore_index=True) if result_dfs else pd.DataFrame()

# 测试调用
df = pd.DataFrame(random_lat_lon(n=10000000), columns=['latitude', 'longitude'])
locations_df = pd.DataFrame(random_lat_lon(n=1000), columns=['lat', 'lon'])
result = select_points_with_balltree(df, locations_df, radius_km=2)

优化方案2：分块并行计算 #

如果locations_df行数极大（比如数万行），可以结合多进程并行处理，进一步压榨CPU性能。将locations_df拆分为多个块，每个块单独用BallTree查询，最后合并结果。

实现代码 #

import pandas as pd
import numpy as np
from sklearn.neighbors import BallTree
from multiprocessing import Pool, cpu_count

def process_chunk(args):
    chunk, df_coords, radius_rad, df = args
    tree = BallTree(df_coords, metric='haversine')
    indices_list = tree.query_radius(np.radians(chunk[['lat', 'lon']].values), r=radius_rad)
    
    chunk_results = []
    for loc_idx, df_indices in enumerate(indices_list):
        if len(df_indices) == 0:
            continue
        loc_row = chunk.iloc[[loc_idx]].reset_index(drop=True)
        repeated_loc = pd.concat([loc_row]*len(df_indices), ignore_index=True)
        matched_df = df.iloc[df_indices].reset_index(drop=True)
        chunk_results.append(pd.concat([matched_df, repeated_loc], axis=1))
    
    return pd.concat(chunk_results, ignore_index=True) if chunk_results else pd.DataFrame()

def select_points_parallel(df, locations_df, radius_km, num_chunks=None):
    EARTH_RADIUS = 6371
    radius_rad = radius_km / EARTH_RADIUS
    df_coords = np.radians(df[['latitude', 'longitude']].values)
    
    # 默认用CPU核心数作为分块数
    if num_chunks is None:
        num_chunks = cpu_count()
    
    # 分割locations_df为多个块
    chunks = np.array_split(locations_df, num_chunks)
    
    # 准备每个进程的参数
    args_list = [(chunk, df_coords, radius_rad, df) for chunk in chunks]
    
    # 并行处理
    with Pool(num_chunks) as pool:
        chunk_results = pool.map(process_chunk, args_list)
    
    return pd.concat(chunk_results, ignore_index=True) if chunk_results else pd.DataFrame()

# 测试调用
result = select_points_parallel(df, locations_df, radius_km=2)

优化方案3：使用GeoPandas空间查询 #

如果需要处理更复杂的地理空间操作（比如多边形、缓冲区分析），可以用GeoPandas结合R-tree索引，代码更直观，适合地理数据场景。

实现代码 #

import geopandas as gpd
from shapely.geometry import Point

def select_points_with_geopandas(df, locations_df, radius_km):
    # 将df转换为GeoDataFrame，使用WGS84坐标系（EPSG:4326）
    gdf = gpd.GeoDataFrame(
        df,
        geometry=gpd.points_from_xy(df.longitude, df.latitude),
        crs="EPSG:4326"
    )
    
    # 将locations_df转换为GeoDataFrame
    loc_gdf = gpd.GeoDataFrame(
        locations_df,
        geometry=gpd.points_from_xy(locations_df.lon, locations_df.lat),
        crs="EPSG:4326"
    )
    
    # 转换为米为单位的投影（需根据数据所在纬度选择对应UTM带，示例用EPSG:32631，对应北纬0-6度）
    utm_crs = "EPSG:32631"
    gdf_utm = gdf.to_crs(utm_crs)
    loc_gdf_utm = loc_gdf.to_crs(utm_crs)
    
    # 创建2公里缓冲区（2000米）
    loc_gdf_utm['buffer'] = loc_gdf_utm.geometry.buffer(2000)
    
    # 空间连接：找出落在缓冲区内的df点
    joined = gpd.sjoin(gdf_utm, loc_gdf_utm[['buffer', 'lat', 'lon']], predicate='within', how='inner')
    
    # 整理结果并转换回原坐标系
    result = joined[['latitude', 'longitude', 'lat', 'lon']].reset_index(drop=True)
    return result

# 测试调用
result = select_points_with_geopandas(df, locations_df, radius_km=2)

方案对比 #

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