1000万条4D空间数据的动态Bounding Box统计查询优化

我有一个包含1000万条记录的numpy数组，数组共5列：前4列对应x、y、z、t坐标，最后一列是各点的标量值。需要为数据集中的每个点，查询其指定4D包围盒（由x_min、x_max、y_min、y_max、z_min、z_max、t_min、t_max定义）内的所有点，计算这些点标量值的中位数和标准差并存储。

注意事项 #

• 每个点对应的包围盒范围差异大，部分范围小、部分范围大；
• x、y、z、t四个轴的分辨率与尺度均不相同；
• 数组仅沿x轴有序，其余三轴无序。

目前已利用x轴有序特性缩小搜索空间，但因需执行1000万次查询（每个点对应一个中心包围盒），希望借助树形数据结构实现更高效的包围盒内点检索。以下是当前的多进程实现代码：

import numpy as np
import time
from multiprocessing import Pool

num_points = 10_000_000
df_us = np.random.rand(num_points, 5)
df = df_us[df_us[:,0].argsort()] # 沿x轴排序，模拟真实数据特性

# 实验中暂不考虑各轴范围和采样分辨率差异
xmin = np.random.rand(num_points)
xmax = xmin + 0.2
ymin = np.random.rand(num_points)
ymax = ymin + 0.4
zmin = np.random.rand(num_points)
zmax = zmin + 0.4
tmin = np.random.rand(num_points)
tmax = tmin + 0.4

def bbox_stat(xr_df, xmin, xmax, ymin, ymax, zmin, zmax, tmin, tmax):
    # 预分配结果数组
    medians = np.empty(len(xmin))
    stds = np.empty(len(xmin))
    
    for i in range(len(xmin)):
        # 利用x轴有序缩小搜索范围
        x_si = np.searchsorted(xr_df[:,0], xmin[i], side='left')
        x_ei = np.searchsorted(xr_df[:,0], xmax[i], side='right')
        
        # 筛选y/z/t轴符合条件的点
        slice_data = xr_df[x_si:x_ei]
        y_mask = (slice_data[:,1] >= ymin[i]) & (slice_data[:,1] <= ymax[i])
        z_mask = (slice_data[:,2] >= zmin[i]) & (slice_data[:,2] <= zmax[i])
        t_mask = (slice_data[:,3] >= tmin[i]) & (slice_data[:,3] <= tmax[i])
        valid_mask = y_mask & z_mask & t_mask
        
        # 计算统计量（修正原代码取列错误，应取第4列标量值）
        valid_scalars = slice_data[valid_mask, 4]
        medians[i] = np.median(valid_scalars) if len(valid_scalars) > 0 else np.nan
        stds[i] = np.std(valid_scalars) if len(valid_scalars) > 0 else np.nan
    
    return medians, stds

arguments = []
num_cores = 50
orig_len = df.shape[0]
unit_len = orig_len // num_cores

for i in range(num_cores):
    start_idx = i * unit_len
    end_idx = (i+1)*unit_len if i != num_cores-1 else orig_len
    arguments.append((
        df,
        xmin[start_idx:end_idx],
        xmax[start_idx:end_idx],
        ymin[start_idx:end_idx],
        ymax[start_idx:end_idx],
        zmin[start_idx:end_idx],
        zmax[start_idx:end_idx],
        tmin[start_idx:end_idx],
        tmax[start_idx:end_idx]
    ))

start_time = time.time()
with Pool(processes=num_cores) as pool:
    results = pool.starmap(bbox_stat, arguments)

# 合并多进程结果
all_medians = np.concatenate([res[0] for res in results])
all_stds = np.concatenate([res[1] for res in results])
end_time = time.time()

print(f"耗时: {end_time - start_time} 秒") 

基于树形数据结构的优化方案 #

方案1：Rtree（轴对齐包围盒专属优化） #

Rtree专为轴对齐空间范围查询设计，无需处理尺度差异，4D场景下性能稳定，完全匹配需求：

实现步骤

1. 安装Rtree：pip install rtree
2. 基于x/y/z/t坐标构建空间索引
3. 针对每个包围盒执行范围查询，直接获取符合条件的点索引

import rtree
import numpy as np
from multiprocessing import Pool

# 预处理数据
coords = df[:, :4]
scalars = df[:, 4]

# 构建Rtree索引（单进程构建效率更高）
idx = rtree.index.Index()
for i in range(len(coords)):
    # 点的空间范围定义为(x,x,y,y,z,z,t,t)
    idx.insert(i, (coords[i,0], coords[i,0], coords[i,1], coords[i,1], coords[i,2], coords[i,2], coords[i,3], coords[i,3]))

def bbox_stat_rtree(idx, scalars, xmin, xmax, ymin, ymax, zmin, zmax, tmin, tmax):
    medians = np.empty(len(xmin))
    stds = np.empty(len(xmin))
    
    for i in range(len(xmin)):
        # 定义4D包围盒范围
        bbox = (xmin[i], xmax[i], ymin[i], ymax[i], zmin[i], zmax[i], tmin[i], tmax[i])
        # 获取范围内所有点的索引
        valid_indices = list(idx.intersection(bbox))
        
        if valid_indices:
            medians[i] = np.median(scalars[valid_indices])
            stds[i] = np.std(scalars[valid_indices])
        else:
            medians[i] = np.nan
            stds[i] = np.nan
    
    return medians, stds

# 多进程调用逻辑同原代码，略去参数切片部分

方案2：Scipy KDTree（多维空间查询） #

KDTree适合低维空间的近邻/范围查询，但需先对各轴归一化消除尺度差异：

from scipy.spatial import KDTree
import numpy as np

# 轴归一化：消除各轴尺度差异
axis_mins = coords.min(axis=0)
axis_maxs = coords.max(axis=0)
norm_coords = (coords - axis_mins) / (axis_maxs - axis_mins)

# 构建KDTree
kdtree = KDTree(norm_coords)

def bbox_stat_kdtree(kdtree, scalars, xmin, xmax, ymin, ymax, zmin, zmax, tmin, tmax):
    medians = np.empty(len(xmin))
    stds = np.empty(len(xmin))
    
    for i in range(len(xmin)):
        # 归一化包围盒边界
        norm_min = np.array([xmin[i], ymin[i], zmin[i], tmin[i]])
        norm_min = (norm_min - axis_mins) / (axis_maxs - axis_mins)
        norm_max = np.array([xmax[i], ymax[i], zmax[i], tmax[i]])
        norm_max = (norm_max - axis_mins) / (axis_maxs - axis_mins)
        
        # 无穷范数的范围查询（对应轴对齐包围盒）
        valid_indices = kdtree.query_ball_point(norm_min, np.inf, p=np.inf, upper_bound=norm_max)
        
        if valid_indices:
            medians[i] = np.median(scalars[valid_indices])
            stds[i] = np.std(scalars[valid_indices])
        else:
            medians[i] = np.nan
            stds[i] = np.nan
    
    return medians, stds

方案对比 #

额外优化建议 #

• 预分配结果数组：避免循环中动态分配内存，提升计算效率
• 批量查询：将多个包围盒的查询请求合并，减少索引IO次数
• 内存映射：超大数据集使用np.memmap加载，避免内存溢出

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