LanceDB Rust客户端查询阻塞问题分析与解决方案

问题背景

在使用LanceDB Rust客户端(v0.17.0)进行向量查询时，开发者遇到了一个查询阻塞的问题。具体表现为：当查询条件中包含表中不存在的值时，第一次查询可以正常返回，但第二次查询会在收集结果时无限期挂起。

问题现象

开发者构建了一个查询函数query_inner，该函数接受一个表引用、嵌入向量、限制条件、最大距离和过滤条件作为参数。当过滤条件中包含表中不存在的文档ID(如document IN (123, 42)，其中123不存在但42存在)时，第一次查询可以正常返回结果，但第二次查询会在以下代码处阻塞：

let record_batches = stream.try_collect::<Vec<_>>().await?;

技术分析

查询计划对比

通过对比Rust和Python版本的查询计划，发现Rust版本多出了一个ScalarIndexQuery节点：

Rust查询计划:

ProjectionExec: expr=[ident@2 as ident, created_at@4 as created_at, document@3 as document, _distance@0 as _distance]
  Take: columns="_distance, _rowid, (ident), (document), (created_at)"
    CoalesceBatchesExec: target_batch_size=1024
      GlobalLimitExec: skip=0, fetch=5
        SortExec: TopK(fetch=5), expr=[_distance@0 ASC NULLS LAST], preserve_partitioning=[false]
          ANNSubIndex: name=embedding_idx, k=5, deltas=1
            ANNIvfPartition: uuid=4dff5dfb-0ddf-4365-82f0-0dbf6081cb47, nprobes=20, deltas=1
            ScalarIndexQuery: query=OR(document = 42,document = 123)

Python查询计划:

ProjectionExec: expr=[ident@2 as ident, document@3 as document, created_at@4 as created_at, embedding@5 as embedding, _distance@0 as _distance]
  Take: columns="_distance, _rowid, (ident), (document), (created_at), (embedding)"
    CoalesceBatchesExec: target_batch_size=8192
      GlobalLimitExec: skip=0, fetch=5
        SortExec: TopK(fetch=5), expr=[_distance@0 ASC NULLS LAST], preserve_partitioning=[false]
          ANNSubIndex: name=embedding_idx, k=5, deltas=1
            ANNIvfPartition: uuid=4dff5dfb-0ddf-4365-82f0-0dbf6081cb47, nprobes=20, deltas=1

根本原因

通过GDB调试发现，阻塞实际上发生在Tokio运行时层面，而不是LanceDB本身。问题根源在于调用方式不当导致的Tokio运行时死锁：

1. 原始代码使用了嵌套的block_on调用，这在Tokio运行时中是不安全的
2. 当查询条件中包含不存在的值时，某些内部操作可能需要在Tokio上下文中执行
3. 第二次查询时，由于运行时状态已被破坏，导致死锁

解决方案

正确的做法是使用spawn_blocking来执行阻塞操作，确保Tokio运行时不会被阻塞。修改后的调用方式如下：

let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(async {
    // ...
    let result = tokio::task::spawn_blocking(move || {
         rt.block_on(query_inner(tbl, embedding, Some(limit), Some(max_dist), filter))
    }).await.unwrap();
    result
});

关键改进点

1. 分离阻塞和非阻塞操作：使用spawn_blocking明确标识出可能阻塞的操作
2. 正确的运行时使用：确保Tokio运行时不会被嵌套调用破坏
3. 线程安全：阻塞操作被分配到专门的线程池执行

技术建议

1. 避免嵌套Tokio运行时：在Tokio生态中，嵌套调用block_on是常见错误来源
2. 合理使用异步/同步边界：明确区分异步和同步代码的边界
3. 理解运行时行为：了解Tokio运行时的线程模型和工作原理
4. 错误处理：为查询操作添加适当的超时机制，避免无限期阻塞

总结

这个问题表面上是LanceDB查询的阻塞问题，实际上是由于Tokio运行时使用不当导致的死锁。通过正确使用spawn_blocking分离阻塞操作，可以有效解决此类问题。这也提醒我们在使用异步Rust库时，需要深入理解底层运行时的行为模式，才能编写出健壮可靠的代码。