基于Faiss中Community ID的向量动态过滤优化相似性搜索

需求背景 #

我需要在Faiss执行相似性搜索前，按特定条件动态排除部分向量。查阅Faiss官方仓库后得知，直接过滤向量会影响性能，官方推荐使用IDSelector替代方案，现尝试该方案落地。

具体场景 #

现有对应SQL表的1000条记录，每条包含messageId、message、communityId字段，可按communityId分组。所有记录已建立向量索引，当特定社区发起查询时，仅需在该社区范围内搜索，无需全量遍历。

官方提到可行方案为「基于向量ID进行过滤」，但不清楚如何将communityId整合到向量ID中，以及如何基于communityId有效过滤记录。

已实现代码 #

os.environ['OPENAI_API_KEY'] = 'key'

messages = ["Hello, world!", "How are you?", "Greetings!","How are you"]
community_ids = [1, 2, 3, 1]


class CommunityIDSelector(faiss.IDSelector):

    def __init__(self):
        pass

    def is_member(self, id):
        return community_ids[id] == 1

id_selector = CommunityIDSelector()


embs = []
embeddings = OpenAIEmbeddings()

for message in messages:
    embs.append(embeddings.embed_query(text=message))

vectors = np.array(embs)
metadata = np.array(community_ids)

concatenated_vectors = np.concatenate((vectors,metadata[:,np.newaxis]),axis=1)

index = faiss.IndexFlatL2(concatenated_vectors.shape[1])
index.add(concatenated_vectors)

target_community_id = 1

query_vector = np.array([embeddings.embed_query(text='I am good')])

# Prepare query vector with target communityId
query_metadata = np.array([[target_community_id]])
concatenated_query = np.concatenate((query_vector, query_metadata), axis=1)

k = 3
distances, indices = index.search(concatenated_query, k,id_selector)

print("Filtered Messages:")

技术问题 #

1. 如何将communityId整合到向量ID中？
2. 如何基于communityId过滤记录？

解决方案 #

1. 将communityId整合到向量ID中的两种方式 #

Faiss的向量ID支持自定义整数，无需局限于默认自增ID，可通过以下两种方式关联communityId：

• 映射表关联：维护一个数组/字典，将Faiss默认的自增ID（0、1、2...）映射到对应的communityId。这种方式灵活无取值范围限制，适合数据量不大的场景（如1000条记录）：

  # 直接用原community_ids数组作为映射表，索引对应Faiss的自增ID
  id_to_community = community_ids.copy()
  

• 位运算编码：若communityId取值范围较小（如不超过2^16），可将communityId放在ID的高段位，原始自增ID放在低段位，通过位运算拼接和解码：

  # 拼接：communityId占高16位，原始索引占低48位（64位整数）
  custom_id = (community_id << 48) | original_index
  # 解码：提取communityId
  extracted_community_id = (custom_id >> 48) & 0xFFFF
  

2. 基于communityId过滤记录的实现 #

核心是通过自定义IDSelector类，结合上述ID整合方式实现过滤：

方式一：基于映射表的过滤

修改CommunityIDSelector类，支持动态传入目标communityId和映射表，避免硬编码：

class CommunityIDSelector(faiss.IDSelector):
    def __init__(self, target_community, id_to_community):
        self.target_community = target_community
        self.id_to_community = id_to_community

    def is_member(self, faiss_id):
        # 判断当前Faiss ID对应的communityId是否匹配目标
        return self.id_to_community[faiss_id] == self.target_community

使用示例：

# 创建针对communityId=1的选择器
id_selector = CommunityIDSelector(target_community=1, id_to_community=id_to_community)
# 执行带过滤的搜索
distances, indices = index.search(query_vector, k, id_selector)

方式二：基于编码ID的过滤

如果用位运算自定义了ID，直接在is_member方法中解码判断：

class CommunityIDSelector(faiss.IDSelector):
    def __init__(self, target_community):
        self.target_community = target_community

    def is_member(self, custom_id):
        # 解码提取communityId
        extracted_community = (custom_id >> 48) & 0xFFFF
        return extracted_community == self.target_community

添加向量时需传入自定义ID：

index = faiss.IndexFlatL2(vectors.shape[1])
# 逐个添加向量并绑定自定义ID
for idx in range(len(vectors)):
    custom_id = (community_ids[idx] << 48) | idx
    index.add_with_ids(vectors[[idx]], np.array([custom_id], dtype=np.int64))

注意：你之前将communityId拼接到向量中的做法是错误的——元数据会干扰向量特征的语义，导致相似性搜索结果失真，应完全通过IDSelector实现过滤逻辑。

内容的提问来源于stack exchange，提问作者James K J