大规模地理数据集分类型最近邻计算加速与存储方案咨询

• 现有规模约200万行的地理数据集，包含4个字段：observation_id（观测ID）、latitude（纬度）、longitude（经度）、class_id（类别ID），数据样例如下：

• 数据集共覆盖18000个类别，类别样本分布不均，所有观测点均位于法国或美国境内。
• 核心需求：对每条观测数据，计算其到每个类别最近观测点的距离，生成维度为18000的距离向量：向量第i位对应该观测点到第i类最近样本的公里距离；若最近距离超过50km，统一填充固定值50即可，无需计算精确距离。
• 后续计划：将生成的距离特征输入MLP等神经网络模型，对比其与普通K近邻模型的效果，距离无需绝对精确，满足最近邻大致判定要求即可。

1. 计算效率极低：初步编写的基于haversine距离的计算脚本，单条观测计算耗时约8秒，全量200万条数据计算周期完全不可接受。现有实现代码如下：

lat1 = 45.705116 # 观测点10561949的纬度
lon1 = 1.424622 # 观测点10561949的经度
df = df[df.observation_id != 10561949] # 从数据集中移除当前待计算的观测点

list_obs = np.full(18000, 50) # 初始化长度18000的数组，默认填充值50

for observation_id, lat2, lon2 in zip(df['observation_id'], df['latitude'], df['longitude']):
  lon1, lat1, lon2, lat2 = map(radians, [lon1, lat1, lon2, lat2]) # 角度转弧度
  a = sin((lat2 - lat1)/2)**2 + cos(lat1) * cos(lat2) * sin((lon2 - lon1 )/2)**2 # Haversine公式上半部分
  dist = 2 * asin(sqrt(a)) * 6371 # Haversine公式计算球面距离，地球半径取6371km
  if list_obs[observation_id] >= dist:
    list_obs[observation_id] = dist

2. 结果存储困难：最终生成结果为2000000×18000规模的数组，体量极大，常规存储方案无法支撑。

计算效率优化 #

• 先做空间范围预过滤：50km的距离阈值对应纬度差约0.45度，经度差随纬度变化在0.45~0.9度区间，提前给所有点按1度经纬度间隔做网格索引，计算单条观测的距离时，仅需遍历自身所在网格+周边8邻域网格内的点，不需要遍历全量200万数据，单条待计算点需要比对的样本量能直接降到原来的1%以下。
• 按类别构建空间索引：不要对每个查询点遍历全量数据，提前给每个类别单独构建cKD树（直接调用scipy.spatial.cKDTree即可，原生支持半径查询，可直接设置50km查询上限，超出直接返回截断值），18000个类别的树可以提前批量构建完成，查询时直接对每个类的树做50km半径的最近邻查询即可，比逐点循环计算haversine距离快2~3个数量级。
• 简化距离计算：由于允许误差，且法国、美国均位于中纬度区域，可以直接将经纬度转换为Web墨卡托投影后的平面坐标，用欧氏距离代替haversine球面距离计算，50km尺度下误差不超过1%，完全满足最近邻判定要求，计算速度比三角函数实现的haversine公式快至少5倍。
• 批量并行计算：放弃单条串行处理逻辑，用numpy做向量化批量计算，配合多进程按CPU核心数拆分查询任务，速度还能再提升数倍到数十倍。

存储方案优化 #

• 利用数据稀疏性压缩：90%以上的距离值都是固定值50，属于极度稀疏特征，不要存储稠密数组，直接用稀疏矩阵格式存储，仅记录距离小于50km的非默认值，存储体积能直接压缩到原稠密数组的5%以下。
• 降低存储精度：距离值不需要双精度浮点精度，直接用float16或者uint8类型存储即可——用uint8存储时可以把0~50km的距离按0.2km步长量化，完全满足模型输入要求，存储体积比float64小87.5%。
• 分块存储：不要存成单个大文件，用numpy.memmap或者zarr、parquet分块格式存储，训练模型时可以按批次按需加载，不需要把全量数据读进内存，普通消费级硬件就能处理。
• 从根源减少存储需求：如果用MLP训练，完全没必要提前把所有距离向量算好存下来，可以在数据加载阶段做在线计算，配合前面的空间索引方案，单batch的计算延迟完全能跟上训练迭代速度，直接省去存储全量特征的成本。

内容的提问来源于stack exchange，提问作者César Leblanc