How to speed up the search with filtering condition like "id in [...]"? (id is the primary key)

[xiaobingxia-at]

As described in the title, How to speed up the search with filtering condition like "id in [...]"? (id is the primary key).
I tried to use STL_SORT key on id, but turns out the search takes 6+ seconds long. (1-million vector collection, id is from [0 to 999,999], search by id in [...80% of all ids...]).
Is there any way to speed it up?

[yhmo]

Is the primary key a varchar field or int64 field? auto-id or not?

import time

from pymilvus import MilvusClient, DataType, connections, utility
import random

collection_name = "XXX"
dim = 768

client = MilvusClient(uri='http://localhost:19530')

schema = client.create_schema(
    auto_id=False,
    enable_dynamic_schema=True,
)

schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=dim)

index_params = client.prepare_index_params()
index_params.add_index(field_name="vector", index_type="FLAT", metric_type="L2", params={})

client.drop_collection(collection_name=collection_name)
client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)

batch = 10000
for i in range(100):
    data = [
        {"id": i*batch+k, "vector": [random.random() for _ in range(dim)]} for k in range(batch)
    ]
    client.insert(collection_name=collection_name, data=data)
    print("insert", i)

print(client.query(collection_name=collection_name, filter="", output_fields=["count(*)"], consistency_level="Strong"))

ids = [i for i in range(1000000)]
filter = f"id in {ids}"
for k in range(10):
    start = time.time()
    res = client.search(
        collection_name=collection_name,
        data=[[random.random() for _ in range(dim)]],
        filter=filter,
        limit=10,
        search_params={},
    )
    end = time.time()
    print("tims cost", end-start, "seconds")
    print(res)

I use a sequence id list to test, average latency is no more than 1 second:

data: ["{'count(*)': 1000000}"] 
tims cost 0.8775138854980469 seconds
data: ["[{'id': 121321, 'distance': 103.90615844726562, 'entity': {}}, {'id': 240471, 'distance': 103.97779846191406, 'entity': {}}, {'id': 929079, 'distance': 104.03509521484375, 'entity': {}}, {'id': 205746, 'distance': 104.36090850830078, 'entity': {}}, {'id': 528637, 'distance': 104.53172302246094, 'entity': {}}, {'id': 869927, 'distance': 105.01641082763672, 'entity': {}}, {'id': 407393, 'distance': 105.03257751464844, 'entity': {}}, {'id': 186134, 'distance': 105.0593032836914, 'entity': {}}, {'id': 174085, 'distance': 105.09526062011719, 'entity': {}}, {'id': 283938, 'distance': 105.16744995117188, 'entity': {}}]"] 
tims cost 0.7235550880432129 seconds
data: ["[{'id': 336083, 'distance': 108.74720764160156, 'entity': {}}, {'id': 374788, 'distance': 108.86965942382812, 'entity': {}}, {'id': 161733, 'distance': 109.30277252197266, 'entity': {}}, {'id': 661507, 'distance': 109.3426742553711, 'entity': {}}, {'id': 900109, 'distance': 109.39573669433594, 'entity': {}}, {'id': 839975, 'distance': 109.45675659179688, 'entity': {}}, {'id': 305347, 'distance': 109.50810241699219, 'entity': {}}, {'id': 169959, 'distance': 109.55290222167969, 'entity': {}}, {'id': 97351, 'distance': 109.66510009765625, 'entity': {}}, {'id': 472761, 'distance': 109.77461242675781, 'entity': {}}]"] 
tims cost 0.7415025234222412 seconds
data: ["[{'id': 366638, 'distance': 104.05885314941406, 'entity': {}}, {'id': 747626, 'distance': 104.23770904541016, 'entity': {}}, {'id': 196435, 'distance': 105.67147827148438, 'entity': {}}, {'id': 58605, 'distance': 105.88360595703125, 'entity': {}}, {'id': 848847, 'distance': 106.00325012207031, 'entity': {}}, {'id': 993889, 'distance': 106.1747817993164, 'entity': {}}, {'id': 380731, 'distance': 106.48426818847656, 'entity': {}}, {'id': 391379, 'distance': 106.55317687988281, 'entity': {}}, {'id': 6527, 'distance': 106.7060775756836, 'entity': {}}, {'id': 952807, 'distance': 106.87200927734375, 'entity': {}}]"] 
tims cost 0.7749273777008057 seconds
data: ["[{'id': 262523, 'distance': 105.84440612792969, 'entity': {}}, {'id': 215244, 'distance': 106.00906372070312, 'entity': {}}, {'id': 757350, 'distance': 107.12776947021484, 'entity': {}}, {'id': 828163, 'distance': 107.17172241210938, 'entity': {}}, {'id': 490194, 'distance': 107.2332763671875, 'entity': {}}, {'id': 643646, 'distance': 107.39852905273438, 'entity': {}}, {'id': 793452, 'distance': 107.8629379272461, 'entity': {}}, {'id': 458728, 'distance': 107.87567138671875, 'entity': {}}, {'id': 654351, 'distance': 107.96870422363281, 'entity': {}}, {'id': 422476, 'distance': 107.99417114257812, 'entity': {}}]"] 
tims cost 0.71490478515625 seconds
data: ["[{'id': 243069, 'distance': 104.41525268554688, 'entity': {}}, {'id': 496631, 'distance': 104.57406616210938, 'entity': {}}, {'id': 489829, 'distance': 104.6295394897461, 'entity': {}}, {'id': 905860, 'distance': 105.23485565185547, 'entity': {}}, {'id': 721659, 'distance': 105.34721374511719, 'entity': {}}, {'id': 413784, 'distance': 105.62699127197266, 'entity': {}}, {'id': 236628, 'distance': 105.9835205078125, 'entity': {}}, {'id': 865910, 'distance': 106.13034057617188, 'entity': {}}, {'id': 64332, 'distance': 106.15592956542969, 'entity': {}}, {'id': 325841, 'distance': 106.37908935546875, 'entity': {}}]"] 
tims cost 0.7002861499786377 seconds
data: ["[{'id': 881694, 'distance': 104.80622863769531, 'entity': {}}, {'id': 187967, 'distance': 105.53680419921875, 'entity': {}}, {'id': 577358, 'distance': 105.6851806640625, 'entity': {}}, {'id': 877622, 'distance': 106.60737609863281, 'entity': {}}, {'id': 694395, 'distance': 106.65791320800781, 'entity': {}}, {'id': 963658, 'distance': 106.66700744628906, 'entity': {}}, {'id': 954151, 'distance': 106.91569519042969, 'entity': {}}, {'id': 64029, 'distance': 107.2932357788086, 'entity': {}}, {'id': 207100, 'distance': 107.29450225830078, 'entity': {}}, {'id': 931661, 'distance': 107.64480590820312, 'entity': {}}]"] 
tims cost 0.7734131813049316 seconds
data: ["[{'id': 948351, 'distance': 102.90803527832031, 'entity': {}}, {'id': 300589, 'distance': 104.72644805908203, 'entity': {}}, {'id': 55921, 'distance': 105.17906951904297, 'entity': {}}, {'id': 889585, 'distance': 105.393798828125, 'entity': {}}, {'id': 988807, 'distance': 105.74481201171875, 'entity': {}}, {'id': 557365, 'distance': 105.8902587890625, 'entity': {}}, {'id': 671415, 'distance': 106.39933776855469, 'entity': {}}, {'id': 227312, 'distance': 106.50597381591797, 'entity': {}}, {'id': 350251, 'distance': 106.61498260498047, 'entity': {}}, {'id': 937961, 'distance': 106.62309265136719, 'entity': {}}]"] 
tims cost 0.7585752010345459 seconds
data: ["[{'id': 441736, 'distance': 103.67152404785156, 'entity': {}}, {'id': 860219, 'distance': 103.96321868896484, 'entity': {}}, {'id': 835589, 'distance': 104.49474334716797, 'entity': {}}, {'id': 704369, 'distance': 105.14997863769531, 'entity': {}}, {'id': 445861, 'distance': 105.72573852539062, 'entity': {}}, {'id': 491230, 'distance': 105.90056610107422, 'entity': {}}, {'id': 590572, 'distance': 106.1396484375, 'entity': {}}, {'id': 777991, 'distance': 106.45188903808594, 'entity': {}}, {'id': 484905, 'distance': 106.60836029052734, 'entity': {}}, {'id': 602891, 'distance': 106.62232971191406, 'entity': {}}]"] 
tims cost 0.7811543941497803 seconds
data: ["[{'id': 770605, 'distance': 105.25657653808594, 'entity': {}}, {'id': 833413, 'distance': 105.91862487792969, 'entity': {}}, {'id': 517137, 'distance': 106.72706604003906, 'entity': {}}, {'id': 707418, 'distance': 106.7569351196289, 'entity': {}}, {'id': 902315, 'distance': 107.0676040649414, 'entity': {}}, {'id': 401690, 'distance': 107.07972717285156, 'entity': {}}, {'id': 963000, 'distance': 107.2469482421875, 'entity': {}}, {'id': 115231, 'distance': 107.27041625976562, 'entity': {}}, {'id': 976731, 'distance': 107.27787780761719, 'entity': {}}, {'id': 861548, 'distance': 107.53715515136719, 'entity': {}}]"] 
tims cost 0.7604734897613525 seconds
data: ["[{'id': 525477, 'distance': 102.885009765625, 'entity': {}}, {'id': 462150, 'distance': 104.77396392822266, 'entity': {}}, {'id': 174301, 'distance': 105.24190521240234, 'entity': {}}, {'id': 270370, 'distance': 105.2962646484375, 'entity': {}}, {'id': 737635, 'distance': 105.68049621582031, 'entity': {}}, {'id': 540677, 'distance': 105.9517822265625, 'entity': {}}, {'id': 923236, 'distance': 106.1259765625, 'entity': {}}, {'id': 841718, 'distance': 106.19412231445312, 'entity': {}}, {'id': 883271, 'distance': 106.43067932128906, 'entity': {}}, {'id': 961291, 'distance': 106.53208923339844, 'entity': {}}]"] 

[xiaobingxia-at]

The whole collection only has one base_id value. So the whole collection's data (1-million) will be fetched, no data will be filtered out.

[xiaobingxia-at]

Or maybe I'm fetching 100 closest entities, and you are fetching 10 closest entities only.

[yhmo]

You can use my script to test. id from 0 to 999999, base_id only has one value, limit = 100, the average time cost is still under 1 second. (the first search is 4 seconds, which is a "cold start")
What value is your embedding dimension?

import time

from pymilvus import MilvusClient, DataType,
import random

collection_name = "XXX"
num_dimensions = 768

client = MilvusClient(uri='http://localhost:19530')

schema = client.create_schema(
    auto_id=False,
    enable_dynamic_schema=True,
)

schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="base_embedding_data", datatype=DataType.FLOAT_VECTOR, dim=num_dimensions)
schema.add_field(field_name="base_id", datatype=DataType.VARCHAR, max_length=64)

index_params = client.prepare_index_params()
index_params.add_index(field_name="id", index_type="INVERTED")
index_params.add_index(field_name="base_embedding_data", index_type="AUTOINDEX", metric_type="COSINE", params={})

client.drop_collection(collection_name=collection_name)
client.create_collection(collection_name=collection_name, schema=schema, index_params=index_params)

batch = 10000
for i in range(100):
    data = [
        {"id": i*batch+k, "base_embedding_data": [random.random() for _ in range(num_dimensions)], "base_id": "base_id"}
        for k in range(batch)
    ]
    client.insert(collection_name=collection_name, data=data)
    print("insert", i)

print(client.query(collection_name=collection_name, filter="", output_fields=["count(*)"], consistency_level="Strong"))

sample_ids = [i for i in range(1000000)]
for k in range(10):
    base_id = "base_id"
    start = time.time()
    res = client.search(
        collection_name=collection_name,
        filter=f"base_id=='{base_id}' and id in {sample_ids}",
        data=[[random.uniform(-1, 1) for _ in range(num_dimensions)]],
        limit=100,
        search_params={"metric_type": "COSINE", "params": {}},
    )
    end = time.time()
    print("tims cost", end-start, "seconds")
    # print(res)

data: ["{'count(*)': 1000000}"] 
tims cost 4.291515350341797 seconds
tims cost 0.9483821392059326 seconds
tims cost 0.9502863883972168 seconds
tims cost 0.9802300930023193 seconds
tims cost 1.0291211605072021 seconds
tims cost 0.8999319076538086 seconds
tims cost 1.0621774196624756 seconds
tims cost 0.9756629467010498 seconds
tims cost 1.0818328857421875 seconds
tims cost 1.0828430652618408 seconds

[xiaobingxia-at]

Thank you again! May I know how large is your cluster? especially the query node.

[yhmo]

I just used a standalone, i7-8700 3.20GHz, 64GB RAM.

[xiaobingxia-at]

Yes, the 'id' is from 0 to 999999.

[xiaofan-luan]

1. You don't need to build index on it when the in list is huge
2. usually it's not the filtering itself slow but the parsing slow -> [Enhancement]: Optimize expression parsing efficiency #36672 is working on improving it.
3. Filter with a list more than 10M is not recommended

[xiaobingxia-at]

Thank you!