DocArray HNSWLib索引中子文档向量搜索问题的解决方案

在使用DocArray的HNSWLib索引时，开发者可能会遇到一个常见问题：当尝试在子索引(subindex)中进行KNN(k-nearest neighbors)向量搜索时，系统会抛出AttributeError: 'NoneType' object has no attribute '_to_node_protobuf'错误。这个问题看似复杂，但其实有明确的解决方案。

问题背景

DocArray是一个用于处理多模态数据的Python库，它提供了HNSWLib索引作为向量数据库的后端。在复杂的数据结构中，我们经常会使用嵌套文档（如主文档包含子文档列表）。当这些子文档包含向量字段时，我们希望能够直接在子文档级别进行向量相似性搜索。

错误原因分析

核心问题在于DocArray预定义的文档类型（如TextDoc、ImageDoc）默认不包含向量维度的类型信息。虽然这些文档类型可以包含embedding字段，但系统无法自动识别这些向量的维度，这在构建向量索引时是必需的信息。

当尝试在子索引中执行KNN搜索时，系统需要明确知道向量维度来正确初始化HNSW索引。缺少这个信息会导致序列化/反序列化过程中出现类型错误，最终表现为NoneType异常。

解决方案

解决这个问题的正确方法是创建自定义文档类型，明确指定embedding字段的维度。以下是具体实现方案：

from docarray.typing import NdArray
from docarray.documents import TextDoc

class MyTextDoc(TextDoc):
    embedding: NdArray[512]  # 明确指定向量维度为512

然后在主文档中使用这个自定义类型：

class QuoteFile(BaseDoc):
    quote_file_id: int
    texts: DocList[MyTextDoc]  # 使用自定义文档类型

实现示例

完整的解决方案实现如下：

import numpy as np
from docarray import DocList, BaseDoc
from docarray.index import HnswDocumentIndex

# 自定义文本文档类型
class MyTextDoc(TextDoc):
    embedding: NdArray[512]  # 明确指定维度

# 主文档类型
class QuoteFile(BaseDoc):
    quote_file_id: int
    texts: DocList[MyTextDoc]

# 配置HNSW索引
hnsw_config = HnswDocumentIndex.DBConfig(
    index_name='quote_files',
    work_dir='./.cache',
    default_column_config={
        np.ndarray: {
            'dim': 512,
            'index': True,
            'space': 'l2',
            # 其他HNSW参数...
        }
    }
)

# 初始化索引
di = HnswDocumentIndex[QuoteFile](hnsw_config)

# 创建并索引文档
td = MyTextDoc(text="Hello World", embedding=np.zeros(512))
qf = QuoteFile(quote_file_id=109, texts=DocList[MyTextDoc]([td]))
di.index(qf)

# 执行子索引搜索
query = MyTextDoc(text="Query", embedding=np.zeros(512))
results = di.find_subindex(query, 'texts', 'embedding', limit=1)

最佳实践建议

1. 始终为向量字段指定维度：即使文档类型只是用于子文档，也应该明确定义embedding字段的维度。

2. 保持维度一致性：确保配置中的维度(dim)与类型注解中的维度一致。

3. 考虑使用Optional：如果embedding字段可能为None，可以使用Optional[NdArray[512]]类型提示。

4. 文档类型继承：从预定义文档类型继承时，记得添加必要的类型信息。

总结

在DocArray中使用HNSWLib索引进行子文档向量搜索时，关键是要确保文档类型包含完整的向量维度信息。通过创建自定义文档类型并明确指定embedding字段的维度，可以避免序列化错误并实现高效的向量搜索功能。这个解决方案不仅适用于HNSWLib后端，也同样适用于DocArray支持的其他向量数据库后端。

理解这个问题的本质有助于开发者更好地设计DocArray文档结构，特别是在处理复杂嵌套文档和向量搜索场景时。记住，明确的类型信息是构建高效向量索引的基础。