code: import os
  import sys
  import shutil
  import platform
  import argparse
  import lancedb
  import time
  import uuid
  import toml
  import pyarrow as pa
  import pyarrow.parquet as pq
  from dotenv import load_dotenv, find_dotenv
  from pymilvus import (
      utility,
      connections,
      FieldSchema,
      CollectionSchema,
      DataType,
      Collection,
  )
  from qdrant_client import QdrantClient, models
  from qdrant_client.http.models import Distance, VectorParams, PointStruct
  
  # Constants
  DEFAULT_CONFIG_FILE = "configs/default.toml"
  MILVUS_MAX_BATCH_SIZE = 10000
  QDRANT_MAX_BATCH_SIZE = 1000
  
  # Import ChromaDB properly
  if platform.system() == "Linux":
      __import__('pysqlite3')
      import sys
      sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
  import chromadb
  
  # Load the environment variables
  load_dotenv(find_dotenv())
  
  
  def read_toml_file(file_path):
      with open(file_path, "r") as f:
          return toml.load(f)
  
  
  def read_parquet_file(file_path):
      table = pq.read_table(file_path)
      df = table.to_pandas()
      lst = df.values.tolist()
      return lst
  
  
  def create_batches(l, n): 
      # looping till length l 
      for i in range(0, len(l), n):  
          yield l[i:i + n]
  
  
  def init_db_collection(config):
      if config["database"] == "chroma":
          if os.path.exists("./chroma_db"):
              shutil.rmtree("./chroma_db")
          db = chromadb.PersistentClient(path="./chroma_db")
          collection = db.get_or_create_collection(config["table"])
      elif config["database"] == "lance":
          if os.path.exists("./lance_db"):
              shutil.rmtree("./lance_db")
          db = lancedb.connect("./lance_db")
          schema = pa.schema(
              [
                  pa.field("embedding", pa.list_(pa.float32(), list_size=config["dimension"])),
                  pa.field("token", pa.string()),
                  pa.field("id", pa.int64()),
              ])
          collection = db.create_table(config["table"], schema=schema)
      elif config["database"] == "milvus":
          connections.connect("default", host="localhost", port="19530")
          if utility.has_collection(config["table"]):
              utility.drop_collection(config["table"])
          fields = [
              FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
              FieldSchema(name="token", dtype=DataType.VARCHAR, max_length=16384),
              FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=config["dimension"])
          ]
          schema = CollectionSchema(fields, config["table"])
          collection = Collection(config["table"], schema)
      elif config["database"] == "qdrant":
          collection = QdrantClient("localhost", port=6333)
          collection.delete_collection(collection_name=config["table"])
          collection.create_collection(
              collection_name=config["table"],
              vectors_config=VectorParams(size=config["dimension"], distance=Distance.DOT),
              optimizers_config=models.OptimizersConfigDiff(
                  indexing_threshold=0,
              ),
          )
      return collection
  
  
  def init_client(config):
      if config["database"] == "qdrant":
          return QdrantClient("localhost", port=6333)
  
  
  def run_query(config, client, vector):
      if config["database"] == "qdrant":
          results = client.search(
              collection_name=config["table"],
              query_vector=vector,
              with_vectors=True,
              with_payload=True,
              limit=config["top_k"],
          )
          for idx, point in enumerate(results):
              print(f"[INFO] Score #{idx}: {point.score}")
  
  
  def insert_into_collection_bulk(collection, batch, config):
      if config["database"] == "milvus":
          mini_batches = list(create_batches(batch, MILVUS_MAX_BATCH_SIZE))
          for b in mini_batches:
              s = time.time()
              collection.insert([
                  # Add an Id to the row (probably uuid.uuid4() would be a good idea)
                  [str(uuid.uuid4()) for _ in b],
                  [row[config["payload_idx"]] for row in b],
                  [list(row[config["embedding_idx"]]) for row in b],
              ])
              print(f"[INFO] Inserted batch of size {len(b)} in {time.time() - s} seconds")
      elif config["database"] == "chroma":
          collection.add(
              ids=[str(idx) for idx, _ in enumerate(batch)],
              documents=[row[config["payload_idx"]] for row in batch],
              embeddings=[list(row[config["embedding_idx"]]) for row in batch],
          )
      elif config["database"] == "lance":
          collection.add([
              {
                  "id": idx, 
                  "token": row[config["payload_idx"]], 
                  "embedding": list(row[config["embedding_idx"]]), 
              } for idx, row in enumerate(batch)])
      elif config["database"] == "qdrant":
          mini_batches = list(create_batches(batch, QDRANT_MAX_BATCH_SIZE))
          for b in mini_batches:
              s = time.time()
              collection.upsert(
                  collection_name=config["table"],
                  wait=True,
                  points=[
                      PointStruct(
                          id=str(uuid.uuid4()), 
                          payload={"token": row[config["payload_idx"]]},
                          vector=list(row[config["embedding_idx"]]),
                      ) for row in b
                  ],
              )
              print(f"[INFO] Inserted batch of size {len(b)} in {time.time() - s} seconds")
  
  
  def get_collection_info(collection, config):
      if config["database"] == "milvus":
          collection.flush()
          print(collection.name)
          print(collection.schema)
          print(collection.num_entities)
  
      elif config["database"] == "chroma":
          print(collection.count())
      elif config["database"] == "lance":
          print(collection.schema)
          print(collection.count_rows())
      elif config["database"] == "qdrant":
          print("Collection info")
          print(collection.get_collection(collection_name=config["table"]))
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser() 
      parser.add_argument("--bench", type=str, default="qdrant-1M-1536", help="The benchmark setup to use")
      parser.add_argument("--debug", action="store_true", help="Whether to run the script in debug mode")
      parser.add_argument("--query", action="store_true", help="Whether to run a query on the collection")
      parser.add_argument("--ingest", action="store_true", help="Whether to ingest the embeddings into the collection")
      args = parser.parse_args()
  
      # Read the default config file
      config = read_toml_file(DEFAULT_CONFIG_FILE)
      config = {**config[args.bench]}
      print("[INFO] Running with config: ")
      print(config)
  
      # Ingest the dataset
      if args.ingest:
          total_rows_written = 0
          collection = init_db_collection(config)
  
          # Read the parquet files
          if args.debug:
              file_list = os.listdir(config["dataset"])[:2]
          else:
              file_list = os.listdir(config["dataset"])[config["train_start_idx"]:config["train_stop_idx"]]
  
          print(f"[INFO] Inserting {len(file_list)} files into collection {config['table']}")
  
          # Insert the embeddings into the collection
          for file in file_list:
              batch = read_parquet_file(os.path.join(config["dataset"], file))
              insert_into_collection_bulk(collection, batch, config)
              total_rows_written += len(batch)
              print(f"[INFO] Total rows written: {total_rows_written}")
          
          # Print out collection stats after insertion
          get_collection_info(collection, config)
      
      # Query the dataset
      if args.query:
          queries_ran = 0
          client = init_client(config)
  
          file_list = os.listdir(config["dataset"])[config["query_start_idx"]:config["query_stop_idx"]]
          print(f"[INFO] Running queries from {len(file_list)} files")
  
          for file in file_list:
              for row in read_parquet_file(os.path.join(config["dataset"], file)):
                  vector = row[config["embedding_idx"]]
                  print(f"[INFO] Running query #{queries_ran} for vector: [{vector[0]}, {vector[1]}, {vector[2]}, ...]")
                  run_query(config, client, vector)
                  queries_ran += 1
                  if queries_ran >= config["queries_to_run"]:
                      break
                  
      print("[INFO] Done!")
  
apis: lancedb.connect
extract_api: [((605, 622), 'platform.system', 'platform.system', ([], {}), '()\n', (620, 622), False, 'import platform\n'), ((707, 735), 'sys.modules.pop', 'sys.modules.pop', (['"""pysqlite3"""'], {}), "('pysqlite3')\n", (722, 735), False, 'import sys\n'), ((798, 811), 'dotenv.find_dotenv', 'find_dotenv', ([], {}), '()\n', (809, 811), False, 'from dotenv import load_dotenv, find_dotenv\n'), ((958, 982), 'pyarrow.parquet.read_table', 'pq.read_table', (['file_path'], {}), '(file_path)\n', (971, 982), True, 'import pyarrow.parquet as pq\n'), ((6025, 6050), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (6048, 6050), False, 'import argparse\n'), ((897, 909), 'toml.load', 'toml.load', (['f'], {}), '(f)\n', (906, 909), False, 'import toml\n'), ((1257, 1286), 'os.path.exists', 'os.path.exists', (['"""./chroma_db"""'], {}), "('./chroma_db')\n", (1271, 1286), False, 'import os\n'), ((1342, 1387), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': '"""./chroma_db"""'}), "(path='./chroma_db')\n", (1367, 1387), False, 'import chromadb\n'), ((3125, 3161), 'qdrant_client.QdrantClient', 'QdrantClient', (['"""localhost"""'], {'port': '(6333)'}), "('localhost', port=6333)\n", (3137, 3161), False, 'from qdrant_client import QdrantClient, models\n'), ((1300, 1328), 'shutil.rmtree', 'shutil.rmtree', (['"""./chroma_db"""'], {}), "('./chroma_db')\n", (1313, 1328), False, 'import shutil\n'), ((1505, 1533), 'os.path.exists', 'os.path.exists', (['"""./lance_db"""'], {}), "('./lance_db')\n", (1519, 1533), False, 'import os\n'), ((1588, 1617), 'lancedb.connect', 'lancedb.connect', (['"""./lance_db"""'], {}), "('./lance_db')\n", (1603, 1617), False, 'import lancedb\n'), ((3785, 3796), 'time.time', 'time.time', ([], {}), '()\n', (3794, 3796), False, 'import time\n'), ((7733, 7762), 'os.listdir', 'os.listdir', (["config['dataset']"], {}), "(config['dataset'])\n", (7743, 7762), False, 'import os\n'), ((1547, 1574), 'shutil.rmtree', 'shutil.rmtree', (['"""./lance_db"""'], {}), "('./lance_db')\n", (1560, 1574), False, 'import shutil\n'), ((1979, 2041), 'pymilvus.connections.connect', 'connections.connect', (['"""default"""'], {'host': '"""localhost"""', 'port': '"""19530"""'}), "('default', host='localhost', port='19530')\n", (1998, 2041), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2053, 2092), 'pymilvus.utility.has_collection', 'utility.has_collection', (["config['table']"], {}), "(config['table'])\n", (2075, 2092), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2459, 2500), 'pymilvus.CollectionSchema', 'CollectionSchema', (['fields', "config['table']"], {}), "(fields, config['table'])\n", (2475, 2500), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2522, 2557), 'pymilvus.Collection', 'Collection', (["config['table']", 'schema'], {}), "(config['table'], schema)\n", (2532, 2557), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((6903, 6932), 'os.listdir', 'os.listdir', (["config['dataset']"], {}), "(config['dataset'])\n", (6913, 6932), False, 'import os\n'), ((6975, 7004), 'os.listdir', 'os.listdir', (["config['dataset']"], {}), "(config['dataset'])\n", (6985, 7004), False, 'import os\n'), ((7272, 7309), 'os.path.join', 'os.path.join', (["config['dataset']", 'file'], {}), "(config['dataset'], file)\n", (7284, 7309), False, 'import os\n'), ((7957, 7994), 'os.path.join', 'os.path.join', (["config['dataset']", 'file'], {}), "(config['dataset'], file)\n", (7969, 7994), False, 'import os\n'), ((2106, 2146), 'pymilvus.utility.drop_collection', 'utility.drop_collection', (["config['table']"], {}), "(config['table'])\n", (2129, 2146), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2178, 2253), 'pymilvus.FieldSchema', 'FieldSchema', ([], {'name': '"""id"""', 'dtype': 'DataType.INT64', 'is_primary': '(True)', 'auto_id': '(True)'}), "(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True)\n", (2189, 2253), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2267, 2334), 'pymilvus.FieldSchema', 'FieldSchema', ([], {'name': '"""token"""', 'dtype': 'DataType.VARCHAR', 'max_length': '(16384)'}), "(name='token', dtype=DataType.VARCHAR, max_length=16384)\n", (2278, 2334), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2348, 2436), 'pymilvus.FieldSchema', 'FieldSchema', ([], {'name': '"""embedding"""', 'dtype': 'DataType.FLOAT_VECTOR', 'dim': "config['dimension']"}), "(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=config[\n    'dimension'])\n", (2359, 2436), False, 'from pymilvus import utility, connections, FieldSchema, CollectionSchema, DataType, Collection\n'), ((2620, 2656), 'qdrant_client.QdrantClient', 'QdrantClient', (['"""localhost"""'], {'port': '(6333)'}), "('localhost', port=6333)\n", (2632, 2656), False, 'from qdrant_client import QdrantClient, models\n'), ((1788, 1799), 'pyarrow.string', 'pa.string', ([], {}), '()\n', (1797, 1799), True, 'import pyarrow as pa\n'), ((1833, 1843), 'pyarrow.int64', 'pa.int64', ([], {}), '()\n', (1841, 1843), True, 'import pyarrow as pa\n'), ((4911, 4922), 'time.time', 'time.time', ([], {}), '()\n', (4920, 4922), False, 'import time\n'), ((1707, 1719), 'pyarrow.float32', 'pa.float32', ([], {}), '()\n', (1717, 1719), True, 'import pyarrow as pa\n'), ((2837, 2898), 'qdrant_client.http.models.VectorParams', 'VectorParams', ([], {'size': "config['dimension']", 'distance': 'Distance.DOT'}), "(size=config['dimension'], distance=Distance.DOT)\n", (2849, 2898), False, 'from qdrant_client.http.models import Distance, VectorParams, PointStruct\n'), ((2930, 2979), 'qdrant_client.models.OptimizersConfigDiff', 'models.OptimizersConfigDiff', ([], {'indexing_threshold': '(0)'}), '(indexing_threshold=0)\n', (2957, 2979), False, 'from qdrant_client import QdrantClient, models\n'), ((3934, 3946), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (3944, 3946), False, 'import uuid\n'), ((4165, 4176), 'time.time', 'time.time', ([], {}), '()\n', (4174, 4176), False, 'import time\n'), ((5404, 5415), 'time.time', 'time.time', ([], {}), '()\n', (5413, 5415), False, 'import time\n'), ((5119, 5131), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (5129, 5131), False, 'import uuid\n')]