Process parquet files asynchronously and upsert vectors to the server. Collects random samples of vectors for testing purposes.
(
dataset_name,
collection,
brute_force_results,
batch_size=100,
matches_sample_size=100,
rps_sample_size=100000,
quick_test=False,
)
| 543 | |
| 544 | |
| 545 | def process_parquet_files( |
| 546 | dataset_name, |
| 547 | collection, |
| 548 | brute_force_results, |
| 549 | batch_size=100, |
| 550 | matches_sample_size=100, |
| 551 | rps_sample_size=100000, |
| 552 | quick_test=False, |
| 553 | ): |
| 554 | """ |
| 555 | Process parquet files asynchronously and upsert vectors to the server. |
| 556 | Collects random samples of vectors for testing purposes. |
| 557 | """ |
| 558 | # Adjust sample sizes for quick test |
| 559 | if quick_test: |
| 560 | matches_sample_size = 10 |
| 561 | rps_sample_size = 1000 |
| 562 | batch_size = 50 |
| 563 | |
| 564 | file_count = 0 |
| 565 | total_vectors_inserted = 0 |
| 566 | total_insertion_time = 0 |
| 567 | id_counter = 0 |
| 568 | matches_test_vectors = [] |
| 569 | rps_test_vectors = [] |
| 570 | |
| 571 | # For quick test, use the first 10 vectors as test vectors |
| 572 | if quick_test: |
| 573 | matches_test_vector_ids_set = set( |
| 574 | str(result["query_id"]) for result in brute_force_results[:10] |
| 575 | ) |
| 576 | else: |
| 577 | matches_test_vector_ids_set = set( |
| 578 | str(result["query_id"]) for result in brute_force_results |
| 579 | ) |
| 580 | |
| 581 | def get_next_file_path(count): |
| 582 | return os.path.join("datasets", dataset_name, f"test{count}.parquet") |
| 583 | |
| 584 | start_time = time.time() |
| 585 | |
| 586 | txns = [] |
| 587 | |
| 588 | with ThreadPoolExecutor(max_workers=3) as executor: |
| 589 | current_path = get_next_file_path(file_count) |
| 590 | if not os.path.exists(current_path): |
| 591 | print( |
| 592 | f"No parquet files found in {os.path.abspath(os.path.dirname(current_path))}" |
| 593 | ) |
| 594 | return matches_test_vectors, rps_test_vectors |
| 595 | |
| 596 | future = executor.submit( |
| 597 | read_single_parquet_file, |
| 598 | current_path, |
| 599 | dataset_name, |
| 600 | file_count, |
| 601 | id_counter, |
| 602 | quick_test, |
no test coverage detected