Function main

tests/test-sparse-vector.py:480–644

()

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

    Path("datasets/sparse_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}")
            collection = create_db(name=vector_db_name, dimension=dimension)
            print("Collection created")

            # Create explicit index
            create_explicit_index(collection)
            print("Explicit index created")
        except Exception as e:
            print(f"Collection may already exist: {e}")
            try:
                collection = client.get_collection(vector_db_name)
                print("Using existing collection")
            except Exception as get_error:
                print(f"Failed to get existing collection: {get_error}")
                raise Exception(f"Cannot create or access collection: {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
        with collection.transaction() as txn:
            with ThreadPoolExecutor(max_workers=32) as executor:
                futures = []