LanceDB项目中预过滤暴力搜索导致内存溢出问题分析

问题背景

在LanceDB数据库系统中，当用户执行带有预过滤条件的暴力向量搜索查询时，系统会出现内存使用量激增的情况，甚至导致内存溢出(OOM)。这个问题在数据集规模达到16GB向量数据、总数据量约22GB时尤为明显。

问题表现

用户在执行类似以下查询时遇到内存问题：

nearest = {
    "use_index": False,
    "column": "vector",
    "q": query_vector,
    "k": 500,
    "metric": "cosine",
}
result = dataset.to_table(
    columns=[],
    with_row_id=True,
    fast_search=True,
    nearest=nearest,
    prefilter=True,
    filter=filter,
)

当过滤条件匹配数据集中的大量记录(约40-70%)时，内存使用量会飙升至50-60GB，在64GB内存的机器上容易引发OOM错误。

技术分析

通过分析查询执行计划，发现问题出现在以下执行流程中：

ProjectionExec: expr=[_distance@2 as _distance, _rowid@0 as _rowid]
  FilterExec: _distance@2 IS NOT NULL
    SortExec: TopK(fetch=500), expr=[_distance@2 ASC NULLS LAST], preserve_partitioning=[false]
      KNNVectorDistance: metric=cosine
        Take: columns="_rowid, vector"
          CoalesceBatchesExec: target_batch_size=8192
            MaterializeIndex: query=<filter here>

关键问题点在于MaterializeIndex操作会将所有匹配过滤条件的记录作为一个大批次(batch)输出，而不是分批处理。这导致后续的TakeExec步骤需要一次性加载所有匹配行，造成内存峰值。

更严重的是，在TakeExec内部处理过程中，会调用concat_batches()函数，这会临时创建一个与原始数据大小相同的副本，导致内存使用量瞬间翻倍。对于大型数据集，这种双重内存分配极易触发OOM。

解决方案思路

1. 分批处理优化：修改MaterializeIndex的实现，使其能够分批输出结果，而不是一次性输出所有匹配记录。这样可以显著降低内存峰值。

2. 内存管理改进：优化TakeExec中的内存使用策略，避免不必要的concat_batches()调用，或者实现更高效的内存复用机制。

3. 查询计划调整：考虑在执行计划中加入显式的分批处理节点，强制将大数据集分割成可管理的小批次。

对用户的影响

这个问题主要影响以下场景的用户：

• 处理大规模向量数据集
• 执行预过滤暴力搜索(不使用索引)
• 过滤条件匹配大量记录

对于这类用户，建议暂时采取以下缓解措施：

• 增加查询的选择性，减少匹配记录数量
• 增加系统可用内存
• 考虑使用索引加速查询

总结

LanceDB中的预过滤暴力搜索内存问题揭示了在实现大规模数据处理系统时需要特别注意的内存管理挑战。通过分析执行计划和内存分配模式，我们能够准确识别问题根源，并提出了针对性的优化方向。这类问题的解决不仅能够提升系统稳定性，也为处理更大规模数据集奠定了基础。