partition search is not faster than collection search

[barzan-hayati]

I have a scenario and my data size is 2M vectors of size 512. Then I generate collection with 2 partitions male and female. They are approximately same size:

num male entities is 1001095
num female entities is 998905

Collection parameters:

    index_params = {
        "metric_type": "COSINE",
        "index_type": "HNSW",
        "params": {"M": 64, "efConstruction": 256},
    }

    search_params = {
        "metric_type": "COSINE",
        "params": {"ef": 128},
    }

code for search on partition:

male_partition.search(
            data=[vector],
            anns_field="vector",
            param=search_params,
            output_fields=["id"],
            limit=50,
            consistency_level="Strong",
        )

code for search on collection:

collection.search(
            data=[vector],
            anns_field="vector",
            partition_names=["male", "female"],
            param=search_params,
            output_fields=["id"],
            limit=50,
            consistency_level="Strong",
        )

I'd like to search in specific partition, as these partitions have half size of the collection, I think that search on male partition should be faster than search on collection. But results show another thing(search on male or female is not too faster than the collection):

2024/08/17 16:03:24.214 [INFO] [__main__] : search among all partitions
2024/08/17 16:03:45.343 [INFO] [__main__] : execution time collection is 0.21128979921340943
2024/08/17 16:03:45.343 [INFO] [__main__] : search among male partition
2024/08/17 16:04:06.744 [INFO] [__main__] : execution time for male partition is 0.21398016214370727
2024/08/17 16:04:06.745 [INFO] [__main__] : search among female partition
2024/08/17 16:04:27.344 [INFO] [__main__] : execution time for female partition is 0.20585099935531617

bellow I bring code for more inspection:

import sys
import random
import sqlite3
import time
import numpy as np
import pandas as pd
from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    utility,
)

sys.path.append("./src")
from _core.settings.development import get_logger

__logger = get_logger(__name__.rsplit(".", 1)[0])

num_inserted_vectors = 0


def fetch_features_from_sqlite_to_collection(
    collection, db_address, num_records_to_fetch, num_records_to_commit, sleeping_time
):
    try:
        """
        prepare to fecth data from database
        """
        while num_fetched < num_records_to_fetch:  ## + self.num_inserted_vectors
            """
            fetch 30000 records
            """
            ids = list(map(lambda x: x[0], records))
            features = list(
                np.reshape(
                    np.frombuffer(
                        np.array(list(map(lambda x: x[1], records))),
                        dtype=np.float32,
                    ),
                    (-1, 512),
                )
            )
            genders = np.random.choice([0, 1], size=len(ids), p=[0.5, 0.5])
            male = np.nonzero(genders)
            ids_man = list(np.asarray(ids)[male])
            features_man = list(np.asarray(features)[male])

            data_male = pd.DataFrame(
                {
                    "id": ids_man,
                    "vector": features_man,
                }
            )
            try:
                mutation_result = collection.insert(
                    data=data_male, partition_name="male"
                )
            except Exception as error:
                __logger.info(f"Error: {error}")
                status = False

            female = np.nonzero(genders == 0)
            ids_woman = list(np.asarray(ids)[female])
            features_woman = list(np.asarray(features)[female])

            data_female = pd.DataFrame(
                {
                    "id": ids_woman,
                    "vector": features_woman,
                }
            )
            try:
                mutation_result = collection.insert(
                    data=data_female, partition_name="female"
                )
            except Exception as error:
                __logger.info(f"Error: {error}")


    except Exception as error:
        __logger.info(f"Error: {error}")
    except sqlite3.Error as error:
        __logger.info(f"Failed to read data from sqlite table: {error}")


def fetch_sqlite_to_collection(collection, collection_name, n: int, sleeping_time: int):

    if not utility.has_collection(collection_name, timeout=6000):
        __logger.info(f"collection exists")
        status_each_sqlite.append(False)
    else:
        num_records_to_fetch = 2000000

        status = fetch_features_from_sqlite_to_collection(
            collection, db_path, row_count, n, sleeping_time
        )


def setup_milvus_from_sqlite(
    collection, collection_name, num_entities: int, gender_partitioning: bool = False
):
    __logger.info("function setup_milvus_from_sqlite")

    n = 30000
    sleeping_time = 1
    status_insert = fetch_sqlite_to_collection(
        collection,
        collection_name,
        n=n,
        sleeping_time=sleeping_time,
    )
    if np.sum(status_insert) == len(status_insert):
        __logger.info(f"{status_insert}=={status_insert}")


def test_collection():
    connections.connect(host="localhost", port="19530")
    collection_name = "test"
    dim = 512

    field_1 = FieldSchema(
        name="id",
        dtype=DataType.INT64,
        is_primary=True,
        auto_id=False,
        description="pk",
    )
    field_2 = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
    schema = CollectionSchema(
        fields=[
            field_1,
            field_2,
        ]
    )
    collection = Collection(name=collection_name, schema=schema)
    collection.load()
    __logger.info(
        f"entities male is {collection.partition(partition_name='male').num_entities}"
    )
    __logger.info(
        f"entities female is {collection.partition(partition_name='female').num_entities}"
    )

    __logger.info("collection loaded")

    search_params = {
        "metric_type": "COSINE",
        "params": {"ef": 128},
    }
    __logger.info("search among all partitions")
    start_time = time.time()

    for i in range(100):
        vector = [random.random() for _ in range(dim)]

        results = collection.search(
            data=[vector],
            anns_field="vector",
            partition_names=["male", "female"],
            param=search_params,
            output_fields=["id"],
            limit=50,
            consistency_level="Strong",
        )
    __logger.info(f"execution time without partition is {time.time() - start_time}")
    __logger.info("search among male partition")
    male_partition = collection.partition(partition_name="male")
    male_partition.load()
    start_time = time.time()

    for i in range(100):
        vector = [random.random() for _ in range(dim)]

        results = male_partition.search(
            data=[vector],
            anns_field="vector",
            # partition_names=["male"],
            param=search_params,
            output_fields=["id"],
            limit=50,
            consistency_level="Strong",
        )

    __logger.info(f"execution time with partition male is {time.time() - start_time}")
    __logger.info("search among female partition")
    female_partition = collection.partition(partition_name="female")
    female_partition.load()

    start_time = time.time()

    for i in range(100):
        vector = [random.random() for _ in range(dim)]
        results = female_partition.search(
            data=[vector],
            anns_field="vector",
            # partition_names=["female"],
            param=search_params,
            output_fields=["id"],
            limit=50,
            consistency_level="Strong",
        )

    __logger.info(f"execution time with partition female is {time.time() - start_time}")
    exit()


def setup_collection():
    connections.connect(host="localhost", port="19530")
    collection_name = "test"
    dim = 512

    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)

    field_1 = FieldSchema(
        name="id",
        dtype=DataType.INT64,
        is_primary=True,
        auto_id=False,
        description="pk",
    )
    field_2 = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim)
    schema = CollectionSchema(
        fields=[
            field_1,
            field_2,
        ]
    )
    collection = Collection(name=collection_name, schema=schema)
    __logger.info(f"collection {collection_name} created")
    collection.create_partition(partition_name="male")
    __logger.info(f"partition male created")
    collection.create_partition(partition_name="female")
    __logger.info(f"partition female created")

    gender_partitioning = True
    num_entities = 0
    setup_milvus_from_sqlite(
        collection, collection_name, num_entities, gender_partitioning
    )

    collection.flush()
    __logger.info(f"row count {collection.num_entities}")

    index_params = {
        "metric_type": "COSINE",
        "index_type": "HNSW",
        "params": {"M": 64, "efConstruction": 256},
    }

    collection.create_index(field_name="vector", index_params=index_params)
    __logger.info("index done")
    collection.load()
    __logger.info("collection loaded")

    search_params = {
        "metric_type": "COSINE",
        "params": {"ef": 128},
    }


if __name__ == "__main__":
    test_new_partition = False
    if test_new_partition is True:
        setup_collection()
    else:
        test_collection()

[xiaofan-luan]

the reason is that you are using strong consistency. So a search need to wait around 200ms.

Remove the strong consistency limit and you will see both search takes < 10ms

[barzan-hayati]

Remove the strong consistency limit and you will see both search takes < 10ms

Removing strong consistency leads to decreasing accuracy. And I would like to have accuracy 99% no lower than it.

[xiaofan-luan]

and use partiton search will help on QPS. since the data is not huge I don't believe there will be a big difference on latency.

[barzan-hayati]

partiton search

I did partition search.

[barzan-hayati]

since the data is not huge I don't believe there will be a big difference on latency

May you right.

[barzan-hayati]

and use partiton search will help on QPS. since the data is not huge I don't believe there will be a big difference on latency.

Does using partition key affect decreasing search time?

[xiaofan-luan]

no, since you data is too small. that only thing helps is to not use strong consistency search.
Strong consistency add 200 ms latency at least, while the search itself might only takes 10ms

[yhmo]

Make clear some points:

• "Removing strong consistency leads to decreasing accuracy" -------------------------- The reason is: some data is still in the message queue(Pulsar/Kafka) not yet consumed by query nodes, and not visible to search requests.
  When the client calls insert(), the data is passed to the message queue and wait for querynode/datanode to asynchronously consume, this is like the WAL(write-ahead-log) machinery. Only when data is consumed by querynode, the data is visible to search requests. The "Strong" consistency level means "wait until the querynode consumes the latest data". Based on the time-tick machinery, milvus check data consistency for each 200ms, this is why the "Strong" consistency level needs 200ms~400ms to wait.

• search in one of male/female partitions vs search both partitions, no much difference even if you remove the "Strong" consistency level
  Since you are using HSNW index, which is pretty fast, it is no much difference for 1 million data or 2 million data search.
  If you change the index type to be FLAT, you will see the difference.

• Does using a partition key affect decreasing search time?
  Partition key is a good choice for there are lots of possibilities of the partition key values. For example, a field named "color" which has the possible values: red, blue, green, purple, yellow, white, black, orange, etc, and your app constantly filters on this field, then it is better to define it as a partition key. When you call search() to filter the partition key with expression like "color==red", it only scans 1/N of the entities to get high performance.
  For you case there are only male/female, no need to define it as a partition key.