大数据量下如何基于两个DataFrame按条件高效新增列

• 需求：为df1新增标记列，标记规则为：df1中某条记录如果在df2中存在同year值、且df2的x落在df1x±10区间、df2的y落在df1y±10区间的记录，则标记为1，否则标记为0。
• 原有实现问题：基于同year全量连接的写法在小数据量下可正常运行，但处理210548行的df1、301661行的df2时，会因同year分组下笛卡尔积数据量爆炸导致内存溢出，触发内核崩溃自动重启。

原有低效实现代码如下：

merged = df1.merge(df2, on="year", how="left", suffixes=("1", "2"))
merged["f"] = (
    (merged.x2 >= merged.x1 - 10)
    & (merged.x2 < merged.x1 + 10)
    & (merged.y2 >= merged.y1 - 10)
    & (merged.y2 < merged.y1 + 10)
)
dataset = (
    merged.groupby(["x1", "y1", "year1", "a", "b", "f"])
    .f.any()
    .astype(int)
    .reset_index()
)

原实现先按year做全量左连接，本质是把所有year相同的行两两配对做笛卡尔积。如果单个year下df1有1万行、df2有1万行，连接后会直接生成1亿行数据，内存占用可达数十GB，远超常规运行环境的内存上限，必然触发OOM崩溃。

优先推荐基于KDTree空间索引的实现，计算效率最高、内存占用最低，20万+30万行规模的数据在普通消费级笔记本上仅需数秒即可跑完，无内存压力。

方案1：KDTree空间索引匹配（推荐） #

核心思路是按year分组后，对二维坐标(x,y)构建空间索引，用切比雪夫距离直接查询每个点10阈值范围内是否存在匹配点，完全避免笛卡尔积计算。

import pandas as pd
import numpy as np
from scipy.spatial import KDTree

# 预处理列名，避免字段冲突
df1 = df1.reset_index(drop=True).rename(columns={'x':'x1', 'y':'y1'})
df2 = df2.reset_index(drop=True).rename(columns={'x':'x2', 'y':'y2'})

match_result = []
# 按year分块处理，进一步缩小匹配范围
for year, group1 in df1.groupby('year'):
    group2 = df2[df2['year'] == year]
    # 当前year下df2无数据，直接标记0
    if len(group2) == 0:
        match_result.extend([0]*len(group1))
        continue
    
    # 用df2的坐标构建KD树
    kd_tree = KDTree(group2[['x2', 'y2']].values)
    # 切比雪夫距离（p=np.inf）下r=10，正好对应x、y差均不超过10的规则
    # 如果要和原代码左闭右开逻辑对齐，把r设为10 - 1e-8即可
    hit_count = kd_tree.query_ball_point(
        group1[['x1', 'y1']].values,
        r=10,
        p=np.inf,
        return_length=True
    )
    # 匹配到至少1条记录则标记1
    match_result.extend((hit_count > 0).astype(int).tolist())

# 结果写回df1
df1['f'] = match_result

方案2：pandas近似连接（无第三方依赖可选） #

如果不想安装scipy，可基于merge_asof做x维度的近似匹配，先过滤掉x差超过10的无效配对，再校验y维度条件，相比原全量笛卡尔积性能也有数十倍提升。

import pandas as pd

# 按x排序，满足merge_asof的前置要求
df1 = df1.sort_values('x').reset_index(drop=True)
df2 = df2.sort_values('x').reset_index(drop=True)

match_result = []
for year, group1 in df1.groupby('year'):
    group2 = df2[df2['year'] == year]
    if len(group2) == 0:
        match_result.extend([0]*len(group1))
        continue
    # 先做x维度的容差连接，仅保留x差在10以内的配对
    merged = pd.merge_asof(
        group1,
        group2,
        on='x',
        by='year',
        direction='nearest',
        tolerance=10,
        suffixes=('1', '2')
    )
    # 再校验y维度条件，聚合得到是否存在匹配
    merged['is_match'] = (merged.y2 >= merged.y1 -10) & (merged.y2 < merged.y1 +10)
    flag = merged.groupby(level=0)['is_match'].max().astype(int).values
    match_result.extend(flag.tolist())

df1['f'] = match_result

• 注意：如果需要保留原代码中x2 < x1+10、y2 < y1+10的左闭右开逻辑，只需要把对应阈值减去1e-8的极小值即可，和原逻辑结果完全一致。

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