数千万行网格配对表的高效多表内连接优化方案问询

我需要对数千万行规模的多张网格配对表执行全局内连接，常规的pandas/polars标准连接在数据量提升后性能瓶颈明显，但这类表存在特殊关联特性，应该有更高效的优化空间。

网格配对表的定义 #

这类表表示两个网格索引的匹配关系，其中一列存储单个索引，另一列存储与之匹配的多个索引列表。例如：

含义是：A[0]匹配B[1]、B[2]；A[1]匹配B[4]、B[5]；A[3]匹配B[7]。

当处理N个网格时，会生成C(N,2)张两两配对表，最终需要得到满足所有网格间两两匹配约束的全局索引组合。

示例与现有实现 #

以3个网格A、B、C为例，对应的3张配对表如下：

AB表 #

AC表 #

BC表 #

最终需要的全局内连接结果是同时满足AB、AC、BC匹配约束的索引组合，有两种等价呈现形式：

压缩格式

展开格式

现有Pandas实现代码 #

from functools import reduce
import pandas as pd 

def create_table(indices: dict, names: list) -> pd.MultiIndex:
    table = pd.DataFrame.from_dict(indices, orient='index').stack().astype(int)
    index = table.reset_index().set_index(['level_0', 0]).index
    return index.set_names(names)
    
AB = create_table({
    0: [1, 2],
    1: [4, 5],
    3: [7],
}, ['A', 'B'])

AC = create_table({
    0: [1, 2, 3],
    1: [2],
}, ['A', 'C'])

BC = create_table({
    1: [1, 2],
    2: [1, 2],
    5: [1, 2, 3],
}, ['B', 'C'])

join   = lambda df1, df2: df1.join(df2, how='inner')
result = reduce(join, [AB, AC, BC]).reorder_levels(['A', 'B', 'C'])
print(result)

运行输出：

MultiIndex([(0, 1, 1),
            (0, 1, 2),
            (0, 2, 1),
            (0, 2, 2),
            (1, 5, 2)],
           names=['A', 'B', 'C'])

性能瓶颈 #

小规模数据下上述方案表现正常，但处理数百万行的多张表时，问题凸显：

• 如果仅以单个网格为锚点（比如只做A-B、A-C连接，不校验B-C约束），我有一个几秒内完成、内存占用极低的方案；
• 但添加全局两两匹配约束后，必须展开所有列表执行内连接，资源消耗直接提升一个数量级以上。

优化思路尝试 #

我考虑保留表格的压缩格式（列元素为列表），避免展开带来的内存爆炸，但需要实现压缩格式下的高效表转置操作，例如：

| A | B | | B | A |
| - | ---- | => | - | ---- |
| 0 | 1, 2 | | 1 | 0 |
| 1 | 1 | | 2 | 0, 1 |

贴近真实场景的Polars压力测试代码 #

from sklearn.neighbors import BallTree
from itertools import combinations
from functools import reduce
import numpy as np 
import polars as pl 
import time

def create_table_pl(indices, names: list):
    c1 = np.repeat(np.arange(len(indices), dtype='int64'), list(map(len, indices)) )
    c2 = np.concatenate(indices).astype(dtype='int64')
    print(f'{names}: {len(c2):,} rows')
    return pl.DataFrame(dict(zip(names, [c1,c2]))).set_sorted(names[0]).lazy()

def match(grids): 
    """ Use a BallTree to find close elements between grids """
    return BallTree(grids[1][:, None], p=np.inf).query_radius(grids[0][:, None], 5)

n = 500000
grids = np.arange(n), np.arange(0,n,2), np.arange(0,n,3), np.arange(0,n,5)
names = list('abcd')
pairs = list(map(create_table_pl, map(match, combinations(grids, 2)), combinations(names, 2)))

start = time.time()
join  = lambda df1, df2: df1.join(df2, on=set(df2.columns).intersection(set(df1.columns)))
table = reduce(join, pairs).select(names).sort(names)

print(table.collect())
print(f'{time.time()-start:.2f} seconds')

运行输出：

('a', 'b'): 2,749,985 rows
('a', 'c'): 1,833,325 rows
('a', 'd'): 1,099,994 rows
('b', 'c'): 916,662 rows
('b', 'd'): 549,997 rows
('c', 'd'): 366,665 rows
shape: (11_183_244, 4)
┌────────┬────────┬────────┬───────┐
│ a      ┆ b      ┆ c      ┆ d     │
│ ---    ┆ ---    ┆ ---    ┆ ---   │
│ i64    ┆ i64    ┆ i64    ┆ i64   │
╞════════╪════════╪════════╪═══════╡
│ 0      ┆ 0      ┆ 0      ┆ 0     │
│ 0      ┆ 0      ┆ 0      ┆ 1     │
│ 0      ┆ 0      ┆ 1      ┆ 0     │
│ 0      ┆ 0      ┆ 1      ┆ 1     │
│ …      ┆ …      ┆ …      ┆ …     │
│ 499999 ┆ 249998 ┆ 166665 ┆ 99999 │
│ 499999 ┆ 249998 ┆ 166666 ┆ 99999 │
│ 499999 ┆ 249999 ┆ 166665 ┆ 99999 │
│ 499999 ┆ 249999 ┆ 166666 ┆ 99999 │
└────────┴────────┴────────┴───────┘
1.66 seconds

图结构等价思路 #

整个问题可以等价转换为图结构问题：每个网格的索引是节点，配对表中的匹配关系是边，最终的全局匹配组合对应图中的全连接环。但目前不确定这条思路能否带来更高效的解决方案。

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