Function main

tests/test-usv.py:489–621

()

def main():
    vector_db_name = "usv_test_db"
    num_vectors = 1_000_000
    dimension = 20_000_000
    max_non_zero_dims = 100
    num_queries = 100
    batch_size = 100
    top_k = 10
    early_terminate_threshold = 0.5

    Path("datasets/usv_dataset").mkdir(parents=True, exist_ok=True)

    # Generate or load dataset
    vectors = generate_dataset(num_vectors, dimension, max_non_zero_dims)

    # Select query vectors
    query_vectors = select_query_vectors(vectors, num_queries)

    # Compute brute force results
    brute_force_results = compute_brute_force_results(
        vectors, query_vectors, dimension, top_k
    )

    # Login to get access token
    print("Logging in to server...")
    create_session()
    print("Session established")
    insert_vectors = input("Insert vectors? (Y/n): ").strip().lower() in ["y", ""]

    # collection = None
    if insert_vectors:
        # Create collection
        try:
            print(f"Creating collection: {vector_db_name}")
            create_db(name=vector_db_name)
            print("Collection created")

            # Create explicit index
            create_explicit_index(vector_db_name)
            print("Explicit index created")
        except Exception as e:
            print(f"Collection may already exist: {e}")

        # Insert vectors into server using a single transaction
        print("Creating single transaction for all vectors")
        print(f"Inserting {num_vectors} vectors in batches of {batch_size}...")
        start = time.time()

        # Create single transaction and process all batches within it
        txn_id = create_transaction(vector_db_name)
        with ThreadPoolExecutor(max_workers=32) as executor:
            futures = []
            for batch_start in range(0, num_vectors, batch_size):
                batch = vectors[batch_start : batch_start + batch_size]
                futures.append(
                    executor.submit(upsert_in_transaction, vector_db_name, txn_id, batch)
                )