Custom Batch Insert not Working

[tasheer10]

I am using Milvus DB 2.3 hosted locally on a VM.
And pymilvus==2.3.2

def insert_data( records, BATCH_SIZE=2048):
        Logger.info("Inserting Data....")
        num_records = len(records)
        num_batches = (num_records + BATCH_SIZE - 1) // BATCH_SIZE
        Logger.debug(f"Number of Records: {num_records} ,Batch Size: {BATCH_SIZE} ,Number of batches: {num_batches}")
        for i in tqdm(range(num_batches)):
            start_idx = i * BATCH_SIZE
            end_idx = min((i + 1) * BATCH_SIZE, num_records)
            Logger.debug(f"Start Index: {start_idx} || End Index : {end_idx}")
            batch = records[start_idx:end_idx]
            collection.insert(batch)
            Logger.debug(f"Batch {i+1}/{num_batches} Completed")
        
        Logger.info("Insert Completed")
        
        if(new):
            Logger.info("Adding Indexes")
            index_params = {
                            "metric_type":"L2",
                            "index_type":"IVF_FLAT",
                            "params":{"nlist":1024}
                            }
            collection.create_index(field_name="Embedding", index_params=index_params)
            Logger.info("Added Index on Embeddings")
            collection.create_index(field_name="FilePath", index_name="scalar_index")
            Logger.info("Added Index on File Path")

I have introduced batching while inserting (as shown above) to know the progress of how much it is inserting,
But when I insert using batching, some batches are not getting inserted, if I insert everything at once, all the data is inserted.
Can you please help me with what I am doing wrong?

Example Logs:
Trying to inset records of length 4697.

When inserting in Batches of 1000:

insert_data(records,BATCH_SIZE=1000)
2023-11-24 19:09:38,880 110996 114764 VectorDB INFO Inserting Data....
2023-11-24 19:09:38,881 110996 114764 VectorDB DEBUG Number of Records: 4697 ,Batch Size: 1000 ,Number of batches: 5
  0%|          | 0/5 [00:00<?, ?it/s]2023-11-24 19:09:38,883 110996 114764 VectorDB DEBUG Start Index: 0 || End Index : 1000
2023-11-24 19:09:39,528 110996 114764 VectorDB DEBUG Batch 1/5 Completed
 20%|██        | 1/5 [00:00<00:02,  1.55it/s]2023-11-24 19:09:39,529 110996 114764 VectorDB DEBUG Start Index: 1000 || End Index : 2000
2023-11-24 19:09:40,576 110996 114764 VectorDB DEBUG Batch 2/5 Completed
 40%|████      | 2/5 [00:01<00:02,  1.13it/s]2023-11-24 19:09:40,577 110996 114764 VectorDB DEBUG Start Index: 2000 || End Index : 3000
2023-11-24 19:09:41,194 110996 114764 VectorDB DEBUG Batch 3/5 Completed
 60%|██████    | 3/5 [00:02<00:01,  1.31it/s]2023-11-24 19:09:41,195 110996 114764 VectorDB DEBUG Start Index: 3000 || End Index : 4000
2023-11-24 19:09:41,808 110996 114764 VectorDB DEBUG Batch 4/5 Completed
 80%|████████  | 4/5 [00:02<00:00,  1.42it/s]2023-11-24 19:09:41,810 110996 114764 VectorDB DEBUG Start Index: 4000 || End Index : 4697
2023-11-24 19:09:42,282 110996 114764 VectorDB DEBUG Batch 5/5 Completed
100%|██████████| 5/5 [00:03<00:00,  1.47it/s]
2023-11-24 19:09:42,285 110996 114764 VectorDB INFO Insert Completed
2023-11-24 19:09:42,285 110996 114764 VectorDB INFO Adding Indexes
2023-11-24 19:09:42,829 110996 114764 VectorDB INFO Added Index on Embeddings
2023-11-24 19:09:43,385 110996 114764 VectorDB INFO Added Index on File Path

print(vectordb_client_status())

Collection Stats:
Collection Name: ExplanationSearch
Collection Description: Explanation Search Data
Number of Entries in Collection : 3000

When inserting with batch 5000:

insert_data(records,BATCH_SIZE=5000)

2023-11-24 19:11:07,168 110996 114764 VectorDB INFO Inserting Data....
2023-11-24 19:11:07,169 110996 114764 VectorDB DEBUG Number of Records: 4697 ,Batch Size: 5000 ,Number of batches: 1
  0%|          | 0/1 [00:00<?, ?it/s]2023-11-24 19:11:07,171 110996 114764 VectorDB DEBUG Start Index: 0 || End Index : 4697
2023-11-24 19:11:10,402 110996 114764 VectorDB DEBUG Batch 1/1 Completed
100%|██████████| 1/1 [00:03<00:00,  3.23s/it]
2023-11-24 19:11:10,404 110996 114764 VectorDB INFO Insert Completed
2023-11-24 19:11:10,405 110996 114764 VectorDB INFO Adding Indexes
2023-11-24 19:11:10,970 110996 114764 VectorDB INFO Added Index on Embeddings
2023-11-24 19:11:11,524 110996 114764 VectorDB INFO Added Index on File Path

print(vectordb_client_status)

Collection Stats:
Collection Name: ExplanationSearch
Collection Description: Explanation Search Data
Number of Entries in Collection : 4697

[xiaofan-luan]

Number of Entries in Collection

Number of Entries in Collection : if you are using num_entities, this might be accurate.

Better way to do it is to query with count(*) as output field.

[tasheer10]

Yes I am using num_entities, but in the first way when I insert with batches of 1000, out of 4697 only 3000 are getting inserted.

[xiaofan-luan]

4697 seems to be a weird number.

1. did you have delete or duplicated keys?
2. could you try to count(*) using https://milvus.io/docs/query.md

[tasheer10]

I have checked for duplicates, all are unique entries. You can see in the second way when I provide the batch as 5000 all 4697 entries got inserted. 4697 is just a random number of entries. Nothing specific. It's just in this example.

[xiaofan-luan]

there might be some data not flushed yet.
the num_entities is not accurate so I recoomend you to use count(*) in order to make sure if your insert success or not.

[tasheer10]

vectordb_client.insert_data(records,BATCH_SIZE=1000)
2023-11-27 12:09:34,349 102632 24320 VectorDB INFO Inserting Data....
2023-11-27 12:09:34,350 102632 24320 VectorDB DEBUG Number of Records: 4697 ,Batch Size: 1000 ,Number of batches: 5
  0%|          | 0/5 [00:00<?, ?it/s]2023-11-27 12:09:34,352 102632 24320 VectorDB DEBUG Start Index: 0 || End Index : 1000
2023-11-27 12:09:35,029 102632 24320 VectorDB DEBUG Batch 1/5 Completed
 20%|██        | 1/5 [00:00<00:02,  1.47it/s]2023-11-27 12:09:35,031 102632 24320 VectorDB DEBUG Start Index: 1000 || End Index : 2000
2023-11-27 12:09:35,714 102632 24320 VectorDB DEBUG Batch 2/5 Completed
 40%|████      | 2/5 [00:01<00:02,  1.47it/s]2023-11-27 12:09:35,716 102632 24320 VectorDB DEBUG Start Index: 2000 || End Index : 3000
2023-11-27 12:09:36,331 102632 24320 VectorDB DEBUG Batch 3/5 Completed
 60%|██████    | 3/5 [00:01<00:01,  1.53it/s]2023-11-27 12:09:36,332 102632 24320 VectorDB DEBUG Start Index: 3000 || End Index : 4000
2023-11-27 12:09:36,994 102632 24320 VectorDB DEBUG Batch 4/5 Completed
 80%|████████  | 4/5 [00:02<00:00,  1.52it/s]2023-11-27 12:09:36,996 102632 24320 VectorDB DEBUG Start Index: 4000 || End Index : 4697
2023-11-27 12:09:37,506 102632 24320 VectorDB DEBUG Batch 5/5 Completed
100%|██████████| 5/5 [00:03<00:00,  1.58it/s]
2023-11-27 12:09:37,508 102632 24320 VectorDB INFO Insert Completed
2023-11-27 12:09:37,508 102632 24320 VectorDB INFO Adding Indexes
2023-11-27 12:09:38,073 102632 24320 VectorDB INFO Added Index on Embeddings
2023-11-27 12:09:38,628 102632 24320 VectorDB INFO Added Index on File Path

print(vectordb_client)
Collection Stats:
Collection Name: SemanticCodeSearch
Collection Description: Semantic Code Search Data
Number of Entries in Collection Using num_enities : 3000
Number of Entities using Count * : {'count(*)': 4697}

Thank you for the suggestion.
It shows the correct value when using count(*), just for my understanding, may I know what's the difference between count(*) and num_entities? And is there any way we can explicitly flush after each insert ?

[xiaofan-luan]

count * did actually process on all datas.

num_entities simply collect all the meta information and it is not accurate becuase:

1. it doesn't count for unflushed data
2. it does not contains any deleted data

[tasheer10]

Thank you for your support ! :)

[yhmo]

No need to call create_index() every time in the insert_data() method. You can call create_index() once after alll the entities are inserted.

I could not reproduce your issue with the following script. The row count is correct after insertions. I didn't change the logic of the insert_data() method.
milvus version: 2.3.2 standalone

import random

from pymilvus import (
    connections,
    FieldSchema, CollectionSchema, DataType,
    Collection,
    utility,
)

connections.connect(host='localhost', port='19530')
print(utility.get_server_version())

collection_name = "test"
dim = 128
metric_type = "L2"

# create collection
id_field = FieldSchema(
    name="ID",
    dtype=DataType.INT64,
    is_primary=True,
    auto_id=True,
)
vector_field = FieldSchema(
    name="Embedding",
    dtype = DataType.FLOAT_VECTOR,
    dim=dim
)

scalar_field = FieldSchema(
    name="FilePath",
    dtype = DataType.VARCHAR,
    max_length=1000,
)

schema = CollectionSchema(fields=[id_field, vector_field, scalar_field])

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

collection = Collection(name=collection_name, schema=schema)
print(f"Collection '{collection_name}' created")
print(collection.describe())

def insert_data(records, BATCH_SIZE=2048):
    print("=================== Inserting Data ========================")
    num_records = len(records)
    num_batches = (num_records + BATCH_SIZE - 1) // BATCH_SIZE
    print(f"Number of Records: {num_records} ,Batch Size: {BATCH_SIZE} ,Number of batches: {num_batches}")
    for i in range(num_batches):
        start_idx = i * BATCH_SIZE
        end_idx = min((i + 1) * BATCH_SIZE, num_records)
        print(f"Start Index: {start_idx} || End Index : {end_idx}")
        batch = records[start_idx:end_idx]
        collection.insert(batch)
        print(f"Batch {i + 1}/{num_batches} Completed")


count = 4697
data=[]
for i in range(count):
    data.append({"Embedding": [random.random() for _ in range(dim)], "FilePath": f"file_path_{i}"})

insert_data(data, BATCH_SIZE=1000)
insert_data(data, BATCH_SIZE=5000)

collection.flush() # call flush() once after all the entities are inserted

# create index
print("Adding Indexes")
index_params = {
    "metric_type": "L2",
    "index_type": "IVF_FLAT",
    "params": {"nlist": 1024}
}
collection.create_index(field_name="Embedding", index_params=index_params)
print("Added Index on Embeddings")
collection.create_index(field_name="FilePath", index_name="scalar_index")
print("Added Index on File Path")

# all entities are inserted, use count(*) to check the row count
collection.load()
res = collection.query(expr="", output_fields=["count(*)"], consistency_level="Strong")
print(res)

The outputs:

v2.3.2-dev
Collection 'test' created
{'collection_name': 'test', 'auto_id': True, 'num_shards': 1, 'description': '', 'fields': [{'field_id': 100, 'name': 'ID', 'description': '', 'type': 5, 'params': {}, 'element_type': 0, 'auto_id': True, 'is_primary': True}, {'field_id': 101, 'name': 'Embedding', 'description': '', 'type': 101, 'params': {'dim': 128}, 'element_type': 0}, {'field_id': 102, 'name': 'FilePath', 'description': '', 'type': 21, 'params': {'max_length': 1000}, 'element_type': 0}], 'aliases': [], 'collection_id': 445921815684915955, 'consistency_level': 2, 'properties': {}, 'num_partitions': 1}
=================== Inserting Data ========================
Number of Records: 4697 ,Batch Size: 1000 ,Number of batches: 5
Start Index: 0 || End Index : 1000
Batch 1/5 Completed
Start Index: 1000 || End Index : 2000
Batch 2/5 Completed
Start Index: 2000 || End Index : 3000
Batch 3/5 Completed
Start Index: 3000 || End Index : 4000
Batch 4/5 Completed
Start Index: 4000 || End Index : 4697
Batch 5/5 Completed
=================== Inserting Data ========================
Number of Records: 4697 ,Batch Size: 5000 ,Number of batches: 1
Start Index: 0 || End Index : 4697
Batch 1/1 Completed
Adding Indexes
Added Index on Embeddings
Added Index on File Path
[{'count(*)': 9394}]

[tasheer10]

Yes the issue was that I was not calling the flush method after insertion, hence it was not updating the metadata and num_entities showed partial results as told by @xiaofan-luan. I have modified the method as follows and it solves the issue.

def insert_data(records, BATCH_SIZE=2048):
        """
        Inserts data records into the active collection in batches.
        After insertion, this method adds indexes to the collection for efficient querying.

        Args:
            records (list): A list of dictionaries containing data records.
            BATCH_SIZE (int, optional): The batch size for inserting records. Defaults to 2048.
        """
        Logger.info("Inserting Data....")
        num_records = len(records)
        num_batches = (num_records + BATCH_SIZE - 1) // BATCH_SIZE
        Logger.debug(f"Number of Records: {num_records} ,Batch Size: {BATCH_SIZE} ,Number of batches: {num_batches}")
        for i in tqdm(range(num_batches)):
            start_idx = i * BATCH_SIZE
            end_idx = min((i + 1) * BATCH_SIZE, num_records)
            Logger.debug(f"Start Index: {start_idx} || End Index : {end_idx}")
            batch = records[start_idx:end_idx]
            collection.insert(batch)
            Logger.debug(f"Batch {i+1}/{num_batches} Completed")
            
        collection.flush() 
        Logger.info("Insert Completed")
        
        if(not(collection.indexes)):
            Logger.info("Adding Indexes")
            index_params = {
                            "metric_type":"L2",
                            "index_type":"IVF_FLAT",
                            "params":{"nlist":1024}
                            }
            collection.create_index(field_name="Embedding", index_params=index_params)
            Logger.info("Added Index on Embeddings")
            collection.create_index(field_name="FilePath", index_name="scalar_index")
            Logger.info("Added Index on File Path")