diff --git a/.gitignore b/.gitignore index 72aee6dc..e8d2bd86 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,4 @@ merge_times.csv merge_times_plot.png .DS_store +__pycache__ diff --git a/scripts/README.md b/scripts/README.md index 00734991..119d5778 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -2,6 +2,24 @@ This directory contains scripts for testing OpenSearch JVector functionality, particularly with large indices. +## Project Structure + +``` +scripts/ +├── jvector_index_and_search/ # JVector indexing and search testing +│ ├── README.md # Comprehensive documentation +│ ├── create_and_test_large_index.py # Main testing script +├── demo.sh # Demo script +├── requirements.txt # Python dependencies +└── README.md # This file +``` + +## Quick Links + +- **[JVector Index and Search Testing](jvector_index_and_search/README.md)** - Main testing framework +- **[Testing Guide](jvector_index_and_search/TESTING_RECALL.md)** - How to test recall measurement +- **[Package Documentation](jvector_index_and_search/jvector_utils/README.md)** - Utilities API reference + ## Installation ### Prerequisites @@ -49,118 +67,82 @@ pip install -r requirements.txt ## Usage -### Creating and Testing Large JVector Index +### JVector Index and Search Testing -The `create_and_test_large_index.py` script creates a large JVector index that exceeds 2GB after force merge, which is useful for testing large index handling capabilities. +The main testing framework is located in the `jvector_index_and_search/` directory. -```bash -python create_and_test_large_index.py [options] -``` - -#### Options: - -- `--host`: OpenSearch host:port (default: localhost:9200) -- `--index`: Index name (default: large-jvector-index) -- `--dimension`: Vector dimension (default: 768) -- `--num-vectors`: Number of vectors to index (default: 3,000,000) -- `--batch-size`: Batch size for indexing (default: 1,000) -- `--shards`: Number of shards (default: 1) +**See the [jvector_index_and_search/README.md](jvector_index_and_search/README.md) for complete documentation.** -#### Example: +#### Quick Start ```bash -# Create a large index with default settings -python create_and_test_large_index.py +cd jvector_index_and_search -# Create a larger index with custom settings -python create_and_test_large_index.py --dimension 1024 --num-vectors 5000000 --batch-size 2000 --shards 2 -``` - -#### What the script does: +# Run unit tests (no OpenSearch required) +python test_recall_measurement.py -1. Creates a knn_vector index with JVector engine -2. Indexes the specified number of vectors with the given dimension -3. Reports index stats before force merge -4. Performs a force merge to consolidate segments -5. Reports index stats after force merge -6. Tests search functionality on the large index +# Run integration tests (requires OpenSearch) +python test_recall_integration.py -#### Notes: +# Create and test a large index +python create_and_test_large_index.py --num-vectors 100000 -- The default settings (3M vectors with 768 dimensions) should create an index exceeding 2GB after force merge -- Adjust the parameters based on your available system resources -- The script requires sufficient memory and disk space to handle large indices +# With recall measurement +python create_and_test_large_index.py --measure-recall --num-vectors 100000 +``` -#### JVector Statistics +#### Key Features -The script collects and reports JVector-specific search and indexing statistics: +- **Large-scale indexing**: Index millions of vectors with configurable batch sizes +- **Force merge tracking**: Monitor graph merge and quantization times +- **Search testing**: Perform multiple searches with detailed JVector statistics +- **Recall measurement**: Memory-efficient ground truth tracking (~1000x less memory) +- **Performance visualization**: Generate plots of merge times vs document count -- `knn_query_visited_nodes`: Number of nodes visited during graph search -- `knn_query_expanded_nodes`: Number of nodes expanded during graph search -- `knn_query_expanded_base_layer_nodes`: Number of base layer nodes expanded -- `knn_query_graph_search_time`: Time spent on graph search (ms) -- `knn_quantization_training_time`: Time spent on quantization training (ms) -- `knn_graph_merge_time`: Time spent on graph merge (ms) +#### Common Commands -##### Search Testing -For each search iteration, the script: -1. Performs a kNN search -2. Collects the JVector stats -3. Reports the incremental changes for each metric +```bash +cd jvector_index_and_search -After all searches are complete, the script provides: -- Initial stats (before any searches) -- Final stats (after all searches) -- Total differences between initial and final stats -- Average values per search +# Basic large index test +python create_and_test_large_index.py --num-vectors 1000000 -This detailed reporting helps in understanding the search behavior and performance characteristics of the JVector engine on a per-query basis. +# With recall measurement +python create_and_test_large_index.py --measure-recall --num-vectors 100000 --num-recall-queries 20 -You can control the number of test searches with the `--num-searches` parameter: +# Performance analysis with visualization +python create_and_test_large_index.py --force-merge-frequency 100000 --csv-output merge_times.csv --plot -```bash -python create_and_test_large_index.py --num-searches 10 +# Search only (skip indexing) +python create_and_test_large_index.py --skip-indexing --index my-existing-index --dimension 768 ``` -#### CSV Output And Plotting +For complete documentation, options, and examples, see: +- **[jvector_index_and_search/README.md](jvector_index_and_search/README.md)** - Complete usage guide +- **[jvector_index_and_search/TESTING_RECALL.md](jvector_index_and_search/TESTING_RECALL.md)** - Testing and troubleshooting -You can save the merge time data to a CSV file using the `--csv-output` option. The CSV file will contain the following columns: +## Other Scripts -- `num_documents`: Number of documents indexed -- `graph_merge_time_ms`: Time taken for graph merge (in milliseconds) -- `quantization_training_time_ms`: Time taken for quantization training (in milliseconds) -- `force_merge_duration_sec`: Duration of force merge (in seconds) -- `index_size_bytes`: Size of the index after force merge (in bytes) +### demo.sh -```shell -# Run with CSV output -python create_and_test_large_index.py --batch-size 1000 --force-merge-frequency 1000 --num-vectors 100000 --csv-output merge_times.csv - -# Generate plots from existing CSV -python create_and_test_large_index.py --csv-output merge_times.csv --plot -``` +A demo script for quick testing (if available). -#### Important Note For Large Indices +## Profiling -When working with large indices, it's important to consider the point at which we will require quantization. -Quantization is becoming critical during index construction when we can't fit the full precision vectors in memory and are forced to use disk. -Therefore, we want to set the `minimum_batch_size_for_quantization` to a value high enough so we can avoid quantization during index construction. -Or alternatively, we can set it to a lower value and accept the additional compute cost of quantization during index construction, and thus avoid the disk access. +You can profile the OpenSearch Java process while running tests: -```shell -# Run with quantization disabled during index construction until we reach 10M documents -python create_and_test_large_index.py --batch-size 1000 --force-merge-frequency 1000 --num-vectors 100000 --min-batch-size-for-quantization 10000000 -``` - -For long running tests you would want to move the script to run in the background and redirect the output to a file: -```shell -nohup python create_and_test_large_index.py --batch-size 5000 --force-merge-frequency 100000 --num-vectors 10000000 --min-batch-size-for-quantization 10000000 > output.log 2>&1 & -``` - -You can also profile the java process while running the script: -```shell +```bash # Get the process id of the opensearch java process PID=$(jps | grep OpenSearch | awk '{print $1}') + # Start profiling jcmd $PID JFR.start name=OnDemand settings=profile duration=600s filename=/tmp/app_jfr_$(date +%s).jfr -``` \ No newline at end of file +``` + +## Contributing + +When adding new scripts or modifying existing ones: +1. Follow the modular structure established in `jvector_index_and_search/` +2. Add comprehensive documentation +3. Include unit tests where applicable +4. Update this README with links to new functionality \ No newline at end of file diff --git a/scripts/create_and_test_large_index.py b/scripts/create_and_test_large_index.py deleted file mode 100644 index 3e397295..00000000 --- a/scripts/create_and_test_large_index.py +++ /dev/null @@ -1,475 +0,0 @@ - -#!/usr/bin/env python3 - -import requests -import json -import time -import numpy as np -import argparse -import sys -import csv -import matplotlib.pyplot as plt -import os - -def create_index(host, index_name, dimension, shards=1, min_batch_size_for_quantization=1000000): - """Create a knn index with jvector engine""" - url = f"http://{host}/{index_name}" - - mapping = { - "settings": { - "index": { - "knn": True, - "knn.derived_source.enabled": True, - "number_of_shards": shards, - "number_of_replicas": 0 - } - }, - "mappings": { - "properties": { - "vector_field": { - "type": "knn_vector", - "dimension": dimension, - "method": { - "name": "disk_ann", - "space_type": "l2", - "engine": "jvector", - "parameters": { - "advanced.min_batch_size_for_quantization": min_batch_size_for_quantization - } - } - }, - "id": {"type": "keyword"} - } - } - } - - response = requests.put(url, json=mapping) - if response.status_code != 200: - print(f"Failed to create index: {response.text}") - sys.exit(1) - print(f"Successfully created index {index_name}") - return response.json() - -def index_vectors(host, index_name, num_vectors, dimension, batch_size=1000, force_merge_frequency=0, csv_file=None): - """Index vectors in batches""" - url = f"http://{host}/{index_name}/_bulk" - headers = {"Content-Type": "application/x-ndjson"} - - total_batches = (num_vectors + batch_size - 1) // batch_size - - # Initialize CSV file if provided - csv_writer = None - csv_file_handle = None - if csv_file: - csv_file_handle = open(csv_file, 'w', newline='') - csv_writer = csv.writer(csv_file_handle) - csv_writer.writerow(['num_documents', 'graph_merge_time_ms', 'quantization_training_time_ms', 'force_merge_duration_sec', 'index_size_bytes']) - - for batch in range(total_batches): - start_idx = batch * batch_size - end_idx = min((batch + 1) * batch_size, num_vectors) - current_batch_size = end_idx - start_idx - - bulk_data = [] - for i in range(current_batch_size): - doc_id = start_idx + i - # Create action line - action = {"index": {"_index": index_name, "_id": str(doc_id)}} - bulk_data.append(json.dumps(action)) - - # Create random vector with values between -1 and 1 - vector = np.random.uniform(-1, 1, dimension).tolist() - document = {"vector_field": vector, "id": str(doc_id)} - bulk_data.append(json.dumps(document)) - - bulk_body = "\n".join(bulk_data) + "\n" - - response = requests.post(url, headers=headers, data=bulk_body) - if response.status_code != 200: - print(f"Failed to index batch {batch+1}/{total_batches}: {response.text}") - sys.exit(1) - - print(f"Indexed batch {batch+1}/{total_batches} ({current_batch_size} vectors)") - - # Force merge if frequency is set and we've reached the threshold - if force_merge_frequency > 0 and (end_idx % force_merge_frequency < batch_size): - print(f"\nPerforming intermediate force merge after {end_idx} documents...") - merge_result = force_merge(host, index_name) - print(f"Index stats after intermediate force merge:") - stats = get_index_stats(host, index_name) - - # Write to CSV if enabled - if csv_writer and merge_result: - index_size = stats["indices"][index_name]["total"]["store"]["size_in_bytes"] if stats else 0 - csv_writer.writerow([ - end_idx, - merge_result['graph_merge_time'], - merge_result['quantization_time'], - merge_result['duration'], - index_size - ]) - csv_file_handle.flush() - - if csv_file_handle: - csv_file_handle.close() - - return True - -def force_merge(host, index_name, max_segments=1): - """Force merge the index to consolidate segments""" - # Get initial KNN stats - initial_stats = get_knn_stats(host) - initial_graph_merge_time = get_knn_stat_value(initial_stats, "knn_graph_merge_time") - initial_quantization_time = get_knn_stat_value(initial_stats, "knn_quantization_training_time") - - # Refresh first to ensure all documents are searchable - refresh_url = f"http://{host}/{index_name}/_refresh" - refresh_response = requests.post(refresh_url) - if refresh_response.status_code != 200: - print(f"Refresh failed: {refresh_response.text}") - - url = f"http://{host}/{index_name}/_forcemerge?max_num_segments={max_segments}&flush=true" - - print(f"Starting force merge to {max_segments} segments...") - start_time = time.time() - response = requests.post(url) - - if response.status_code != 200: - print(f"Force merge failed: {response.text}") - return False - - duration = time.time() - start_time - print(f"Force merge completed in {duration:.2f} seconds") - - # Get final KNN stats - final_stats = get_knn_stats(host) - final_graph_merge_time = get_knn_stat_value(final_stats, "knn_graph_merge_time") - final_quantization_time = get_knn_stat_value(final_stats, "knn_quantization_training_time") - - # Calculate and display the differences - graph_merge_diff = final_graph_merge_time - initial_graph_merge_time - quantization_diff = final_quantization_time - initial_quantization_time - - print(f"KNN Graph Merge Time: +{graph_merge_diff} ms") - print(f"KNN Quantization Training Time: +{quantization_diff} ms") - - return { - 'graph_merge_time': graph_merge_diff, - 'quantization_time': quantization_diff, - 'duration': duration - } - -def get_index_stats(host, index_name): - """Get index stats including size""" - url = f"http://{host}/{index_name}/_stats" - response = requests.get(url) - - if response.status_code != 200: - print(f"Failed to get index stats: {response.text}") - return None - - stats = response.json() - size_in_bytes = stats["indices"][index_name]["total"]["store"]["size_in_bytes"] - size_in_gb = size_in_bytes / (1024 * 1024 * 1024) - - print(f"Index size: {size_in_bytes} bytes ({size_in_gb:.2f} GB)") - return stats - -def get_knn_stats(host): - """Get KNN plugin stats including JVector search metrics""" - url = f"http://{host}/_plugins/_knn/stats" - response = requests.get(url) - - if response.status_code != 200: - print(f"Failed to get KNN stats: {response.text}") - return None - - stats = response.json() - return stats - -def print_jvector_search_stats(stats): - """Extract and print JVector-specific search statistics""" - if not stats or "nodes" not in stats: - print("No KNN stats available") - return - - # JVector-specific metrics we want to track - jvector_metrics = [ - "knn_query_visited_nodes", - "knn_query_expanded_nodes", - "knn_query_expanded_base_layer_nodes" - ] - - # Collect stats from all nodes - all_nodes_stats = {} - for node_id, node_stats in stats["nodes"].items(): - for metric in jvector_metrics: - if metric in node_stats: - if metric not in all_nodes_stats: - all_nodes_stats[metric] = 0 - all_nodes_stats[metric] += int(node_stats[metric]) - - # Print the stats - print("\nJVector Search Statistics:") - for metric, value in all_nodes_stats.items(): - print(f" {metric}: {value}") - -def test_search(host, index_name, dimension, k=10): - """Test kNN search on the index""" - url = f"http://{host}/{index_name}/_search" - - # Create random query vector - query_vector = np.random.uniform(-1, 1, dimension).tolist() - - query = { - "size": k, - "query": { - "knn": { - "vector_field": { - "vector": query_vector, - "k": k - } - } - } - } - - start_time = time.time() - response = requests.post(url, json=query) - duration = time.time() - start_time - - if response.status_code != 200: - print(f"Search failed: {response.text}") - return False - - results = response.json() - hits = results["hits"]["hits"] - - print(f"Search completed in {duration:.4f} seconds, found {len(hits)} results") - return True - -def test_search_with_stats(host, index_name, dimension, k=10, num_searches=5): - """Test kNN search on the index and report JVector stats for each iteration""" - # Get initial stats - current_stats = get_knn_stats(host) - print("\nInitial JVector Stats:") - print_jvector_search_stats(current_stats) - - # JVector-specific metrics we want to track - jvector_metrics = [ - "knn_query_visited_nodes", - "knn_query_expanded_nodes", - "knn_query_expanded_base_layer_nodes" - ] - - # Store initial values - prev_totals = {} - if current_stats and "nodes" in current_stats: - for node_id, node_stats in current_stats["nodes"].items(): - for metric in jvector_metrics: - if metric in node_stats: - if metric not in prev_totals: - prev_totals[metric] = 0 - prev_totals[metric] += int(node_stats[metric]) - - # Perform multiple searches and check stats after each - for i in range(num_searches): - print(f"\n--- Search Iteration {i+1}/{num_searches} ---") - - # Perform search - test_search(host, index_name, dimension, k) - - # Get stats after this search - new_stats = get_knn_stats(host) - - # Calculate current totals - current_totals = {} - if new_stats and "nodes" in new_stats: - for node_id, node_stats in new_stats["nodes"].items(): - for metric in jvector_metrics: - if metric in node_stats: - if metric not in current_totals: - current_totals[metric] = 0 - current_totals[metric] += int(node_stats[metric]) - - # Calculate and print differences for this iteration - print("\nJVector Stats for this iteration:") - for metric in jvector_metrics: - prev_val = prev_totals.get(metric, 0) - current_val = current_totals.get(metric, 0) - diff = current_val - prev_val - print(f" {metric}: +{diff}") - - # Update previous totals for next iteration - prev_totals = current_totals - - # Get final stats - final_stats = get_knn_stats(host) - - # Calculate initial and final totals - initial_totals = {} - if current_stats and "nodes" in current_stats: - for node_id, node_stats in current_stats["nodes"].items(): - for metric in jvector_metrics: - if metric in node_stats: - if metric not in initial_totals: - initial_totals[metric] = 0 - initial_totals[metric] += int(node_stats[metric]) - - final_totals = {} - if final_stats and "nodes" in final_stats: - for node_id, node_stats in final_stats["nodes"].items(): - for metric in jvector_metrics: - if metric in node_stats: - if metric not in final_totals: - final_totals[metric] = 0 - final_totals[metric] += int(node_stats[metric]) - - # Print summary - print("\n=== JVector Stats Summary ===") - print("\nInitial Stats:") - for metric, value in initial_totals.items(): - print(f" {metric}: {value}") - - print("\nFinal Stats:") - for metric, value in final_totals.items(): - print(f" {metric}: {value}") - - print("\nTotal Differences (Final - Initial):") - for metric in jvector_metrics: - initial_val = initial_totals.get(metric, 0) - final_val = final_totals.get(metric, 0) - diff = final_val - initial_val - print(f" {metric}: +{diff}") - - # Calculate per-search averages - print("\nAverage per Search:") - for metric in jvector_metrics: - initial_val = initial_totals.get(metric, 0) - final_val = final_totals.get(metric, 0) - diff = final_val - initial_val - avg = diff / num_searches if num_searches > 0 else 0 - print(f" {metric}: {avg:.2f}") - - return True - -def get_knn_stat_value(stats, stat_name): - """Extract a specific KNN stat value from all nodes""" - total = 0 - if stats and "nodes" in stats: - for node_id, node_stats in stats["nodes"].items(): - if stat_name in node_stats: - total += int(node_stats[stat_name]) - return total - -def plot_merge_times(csv_file): - """Plot graph merge time and quantization training time vs number of documents""" - if not os.path.exists(csv_file): - print(f"CSV file {csv_file} not found") - return - - # Read CSV data - num_docs = [] - graph_merge_times = [] - quantization_times = [] - - with open(csv_file, 'r') as f: - reader = csv.DictReader(f) - for row in reader: - num_docs.append(int(row['num_documents'])) - graph_merge_times.append(int(row['graph_merge_time_ms'])) - quantization_times.append(int(row['quantization_training_time_ms'])) - - if not num_docs: - print("No data found in CSV file") - return - - # Create plots - fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8)) - - # Graph merge time plot - ax1.plot(num_docs, graph_merge_times, 'b-o', linewidth=2, markersize=6) - ax1.set_xlabel('Number of Documents') - ax1.set_ylabel('Graph Merge Time (ms)') - ax1.set_title('JVector Graph Merge Time vs Number of Documents') - ax1.grid(True, alpha=0.3) - - # Quantization training time plot - ax2.plot(num_docs, quantization_times, 'r-o', linewidth=2, markersize=6) - ax2.set_xlabel('Number of Documents') - ax2.set_ylabel('Quantization Training Time (ms)') - ax2.set_title('JVector Quantization Training Time vs Number of Documents') - ax2.grid(True, alpha=0.3) - - plt.tight_layout() - - # Save plot - plot_file = csv_file.replace('.csv', '_plot.png') - plt.savefig(plot_file, dpi=300, bbox_inches='tight') - print(f"Plot saved to {plot_file}") - - # Show plot - plt.show() - -def main(): - parser = argparse.ArgumentParser(description="Create and test a large JVector index in OpenSearch") - parser.add_argument("--host", default="localhost:9200", help="OpenSearch host:port") - parser.add_argument("--index", default="large-jvector-index", help="Index name") - parser.add_argument("--dimension", type=int, default=768, help="Vector dimension") - parser.add_argument("--num-vectors", type=int, default=3000000, - help="Number of vectors to index (3M vectors with dim=768 should exceed 2GB)") - parser.add_argument("--batch-size", type=int, default=1000, help="Batch size for indexing") - parser.add_argument("--shards", type=int, default=1, help="Number of shards") - parser.add_argument("--num-searches", type=int, default=5, help="Number of searches to perform for stats testing") - parser.add_argument("--skip-indexing", action="store_true", help="Skip index creation and indexing, only run searches") - parser.add_argument("--force-merge-frequency", type=int, default=0, - help="Force merge after every N documents (0 to disable intermediate merges)") - parser.add_argument("--csv-output", type=str, help="CSV file to save merge time data") - parser.add_argument("--plot", action="store_true", help="Generate plots from CSV data") - parser.add_argument("--min-batch-size-for-quantization", type=int, default=1000000, - help="Minimum batch size for quantization (default: 1M)") - - args = parser.parse_args() - - if args.plot and args.csv_output: - plot_merge_times(args.csv_output) - return - - if not args.skip_indexing: - print(f"Creating large JVector index with {args.num_vectors} vectors of dimension {args.dimension}") - print(f"Estimated size: ~{args.num_vectors * args.dimension * 4 / (1024*1024*1024):.2f} GB (raw vectors only)") - - # Create index - create_index(args.host, args.index, args.dimension, args.shards, args.min_batch_size_for_quantization) - - # Index vectors - index_vectors(args.host, args.index, args.num_vectors, args.dimension, args.batch_size, args.force_merge_frequency, args.csv_output) - - # Get stats before final force merge - print("\nIndex stats before final force merge:") - get_index_stats(args.host, args.index) - - # Force merge - force_merge(args.host, args.index) - - # Get stats after force merge - print("\nIndex stats after final force merge:") - get_index_stats(args.host, args.index) - else: - print(f"Skipping index creation and indexing. Using existing index: {args.index}") - # Get current index stats - print("\nCurrent index stats:") - get_index_stats(args.host, args.index) - - # Test search with JVector stats - print("\nTesting search with JVector stats:") - test_search_with_stats(args.host, args.index, args.dimension, k=10, num_searches=args.num_searches) - - print("\nTest completed successfully!") - - # Generate plots if CSV output was specified - if args.csv_output and os.path.exists(args.csv_output): - print(f"\nCSV data saved to {args.csv_output}") - plot_merge_times(args.csv_output) - -if __name__ == "__main__": - main() diff --git a/scripts/jvector_index_and_search/README.md b/scripts/jvector_index_and_search/README.md new file mode 100644 index 00000000..b5b48e23 --- /dev/null +++ b/scripts/jvector_index_and_search/README.md @@ -0,0 +1,303 @@ +# JVector Index and Search Testing + +A comprehensive testing framework for JVector indices in OpenSearch, including indexing, search testing, and memory-efficient recall measurement. + +## Overview + +This directory contains tools for creating, testing, and benchmarking JVector indices with features including: +- Large-scale vector indexing with batch operations +- Force merge with detailed statistics tracking +- Search performance testing with JVector-specific metrics +- Memory-efficient recall measurement using heap-based ground truth tracking +- Performance visualization + +## Directory Structure + +``` +jvector_index_and_search/ +├── README.md # This file +├── create_and_test_large_index.py # Main testing script +├── test_recall_measurement.py # Unit tests for recall measurement +├── test_recall_integration.py # Integration tests with OpenSearch +├── TESTING_RECALL.md # Comprehensive testing guide +├── jvector_utils/ # Modular utilities package +│ ├── __init__.py +│ ├── index_operations.py # Index creation and management +│ ├── search_operations.py # Search testing and statistics +│ ├── recall_measurement.py # Ground truth tracking +│ ├── stats_utils.py # KNN statistics utilities +│ ├── visualization.py # Performance plotting +│ └── README.md # Package documentation +├── merge_times.csv # Example output data +└── merge_times_plot.png # Example visualization +``` + +## Quick Start + +### Prerequisites + +1. OpenSearch running with JVector plugin installed +2. Python 3.7+ with required dependencies + +### Installation + +```bash +cd scripts/jvector_index_and_search + +# Install dependencies +pip install -r ../requirements.txt +``` + +### Basic Usage + +```bash +# Create and test a large index +python create_and_test_large_index.py --num-vectors 100000 + +# With recall measurement +python create_and_test_large_index.py --measure-recall --num-vectors 100000 + +# With performance tracking +python create_and_test_large_index.py --csv-output merge_times.csv --plot +``` + +## Main Script: create_and_test_large_index.py + +### Features + +- **Large-scale indexing**: Index millions of vectors with configurable batch sizes +- **Force merge tracking**: Monitor graph merge and quantization times +- **Search testing**: Perform multiple searches with detailed JVector statistics +- **Recall measurement**: Memory-efficient ground truth tracking (uses ~1000x less memory) +- **Performance visualization**: Generate plots of merge times vs document count + +### Common Options + +```bash +# Basic options +--host localhost:9200 # OpenSearch host:port +--index large-jvector-index # Index name +--dimension 768 # Vector dimension +--num-vectors 3000000 # Number of vectors to index +--batch-size 1000 # Batch size for indexing + +# Search testing +--num-searches 5 # Number of searches to perform +--skip-indexing # Skip indexing, only run searches + +# Recall measurement +--measure-recall # Enable recall measurement +--num-recall-queries 10 # Number of query vectors for recall + +# Performance tracking +--force-merge-frequency 100000 # Force merge every N documents +--csv-output merge_times.csv # Save merge statistics to CSV +--plot # Generate plots from CSV data +``` + +### Examples + +#### 1. Basic Large Index Test +```bash +python create_and_test_large_index.py \ + --num-vectors 1000000 \ + --dimension 768 \ + --batch-size 1000 +``` + +#### 2. With Recall Measurement +```bash +python create_and_test_large_index.py \ + --measure-recall \ + --num-vectors 100000 \ + --num-recall-queries 20 \ + --num-searches 20 +``` + +#### 3. Performance Analysis +```bash +python create_and_test_large_index.py \ + --num-vectors 500000 \ + --force-merge-frequency 100000 \ + --csv-output merge_times.csv \ + --plot +``` + +#### 4. Search Only (Existing Index) +```bash +python create_and_test_large_index.py \ + --skip-indexing \ + --index my-existing-index \ + --dimension 768 \ + --num-searches 10 +``` + +## Testing + +### Unit Tests (No OpenSearch Required) + +Test the core recall measurement logic: + +```bash +python test_recall_measurement.py +``` + +**Expected output:** All 6 tests pass ✅ + +### Integration Tests (Requires OpenSearch) + +Test end-to-end with a real OpenSearch instance: + +```bash +python test_recall_integration.py + +# Customize parameters +python test_recall_integration.py --dimension 256 --num-vectors 5000 +``` + +**Expected output:** Recall values >0.9 ✅ + +See [TESTING_RECALL.md](TESTING_RECALL.md) for comprehensive testing documentation. + +## Memory-Efficient Recall Measurement + +### How It Works + +Traditional recall measurement requires storing all vectors in memory to compute ground truth: +- **Memory usage**: O(num_vectors × dimension) +- **Example**: 1M vectors × 768 dims ≈ 5,859 MB (~6 GB) + +Our heap-based approach tracks ground truth incrementally: +- **Memory usage**: O(num_queries × k) +- **Example**: 100 queries × k=10 ≈ 6 MB + +**Result**: ~1000x memory reduction! 🎉 + +### Usage + +```bash +python create_and_test_large_index.py \ + --measure-recall \ + --num-recall-queries 50 \ + --num-vectors 1000000 +``` + +The script will: +1. Pre-generate 50 query vectors +2. Track ground truth incrementally during indexing (using heaps) +3. Perform 50 searches and measure recall for each +4. Report average, min, max, and std dev of recall values + +## JVector Statistics Tracked + +The script collects detailed JVector-specific metrics: + +- `knn_query_visited_nodes` - Total nodes visited during search +- `knn_query_expanded_nodes` - Nodes expanded during search +- `knn_query_expanded_base_layer_nodes` - Base layer nodes expanded +- `knn_graph_merge_time` - Time spent merging graphs +- `knn_quantization_training_time` - Time spent on quantization + +Statistics are reported: +- Per search iteration +- As aggregate totals +- As per-search averages + +## Performance Visualization + +Generate plots showing how merge times scale with document count: + +```bash +# During indexing +python create_and_test_large_index.py \ + --force-merge-frequency 100000 \ + --csv-output merge_times.csv + +# Generate plots from existing CSV +python create_and_test_large_index.py \ + --plot \ + --csv-output merge_times.csv +``` + +Output: `merge_times_plot.png` with two subplots: +1. Graph merge time vs number of documents +2. Quantization training time vs number of documents + +## Package: jvector_utils + +The `jvector_utils` package provides reusable utilities that can be imported in other scripts: + +```python +from jvector_utils.index_operations import create_index, index_vectors +from jvector_utils.search_operations import test_search_with_stats +from jvector_utils.recall_measurement import GroundTruthTracker +from jvector_utils.visualization import plot_merge_times + +# Use in your own scripts +create_index("localhost:9200", "my-index", dimension=768) +index_vectors("localhost:9200", "my-index", num_vectors=10000, dimension=768) +``` + +See [jvector_utils/README.md](jvector_utils/README.md) for detailed package documentation. + +## Troubleshooting + +### Connection Errors + +**Problem**: "Failed to connect to OpenSearch" + +**Solution**: +```bash +# Verify OpenSearch is running +curl http://localhost:9200 + +# Check JVector plugin is installed +curl http://localhost:9200/_cat/plugins +``` + +### Low Recall Values + +**Problem**: Recall consistently <0.8 + +**Solutions**: +1. Run unit tests to verify recall measurement: `python test_recall_measurement.py` +2. Increase index size (larger indices generally have better recall) +3. Check distance metric matches (L2 vs cosine) +4. Review JVector configuration parameters + +### Memory Issues + +**Problem**: Out of memory during indexing + +**Solutions**: +1. Reduce batch size: `--batch-size 500` +2. Reduce number of recall queries: `--num-recall-queries 10` +3. Disable recall measurement if not needed +4. Use smaller vector dimension for testing + +## Best Practices + +1. **Start small**: Test with 10K-100K vectors before scaling to millions +2. **Use recall measurement**: Validate search quality with `--measure-recall` +3. **Monitor performance**: Use `--csv-output` to track merge times +4. **Run tests first**: Verify setup with `test_recall_measurement.py` +5. **Document results**: Save CSV outputs and plots for analysis + +## Additional Resources + +- **Testing Guide**: [TESTING_RECALL.md](TESTING_RECALL.md) +- **Package Documentation**: [jvector_utils/README.md](jvector_utils/README.md) +- **Main Scripts README**: [../README.md](../README.md) + +## Contributing + +When making changes: +1. Run unit tests: `python test_recall_measurement.py` +2. Run integration tests: `python test_recall_integration.py` +3. Verify backward compatibility +4. Update documentation as needed + +## License + +See the main repository LICENSE file. + diff --git a/scripts/jvector_index_and_search/TESTING_RECALL.md b/scripts/jvector_index_and_search/TESTING_RECALL.md new file mode 100644 index 00000000..99ff094e --- /dev/null +++ b/scripts/jvector_index_and_search/TESTING_RECALL.md @@ -0,0 +1,293 @@ +# Testing Recall Measurement + +This guide explains how to test and verify that the recall measurement functionality is working correctly. + +## Quick Start + +### 1. Unit Tests (No OpenSearch Required) + +Run the comprehensive unit test suite to verify the core recall measurement logic: + +```bash +cd scripts +python test_recall_measurement.py +``` + +**Expected Result:** +``` +====================================================================== +TEST SUMMARY +====================================================================== +✅ PASS: Basic Functionality +✅ PASS: Large Vector Set +✅ PASS: Recall Calculation +✅ PASS: Multiple Queries +✅ PASS: Cosine Distance +✅ PASS: Edge Cases + +Total: 6/6 tests passed + +🎉 All tests passed! Recall measurement is working correctly. +``` + +### 2. Integration Test (Requires OpenSearch) + +Test recall measurement end-to-end with a real OpenSearch instance: + +```bash +cd scripts +python test_recall_integration.py +``` + +**Expected Result:** +- Index created successfully +- Vectors indexed with ground truth tracking +- Searches performed with recall measurement +- Recall values typically >0.9 for small indices +- Test index cleaned up automatically + +## What Gets Tested + +### Unit Tests (`test_recall_measurement.py`) + +#### Test 1: Basic Functionality +- Creates a simple 3D test case with known vectors +- Verifies that the k-nearest neighbors are correctly identified +- Tests with vectors at known distances from the query + +#### Test 2: Large Vector Set +- Tests with 100 random vectors +- Compares tracker results against brute-force computation +- Ensures heap-based approach matches exact computation + +#### Test 3: Recall Calculation +- Tests perfect recall (100%) +- Tests partial recall (80%) +- Tests zero recall (0%) +- Verifies the `calculate_recall()` function + +#### Test 4: Multiple Query Vectors +- Tests tracking ground truth for multiple queries simultaneously +- Verifies each query has correct ground truth +- Tests with 3 queries and 50 vectors + +#### Test 5: Cosine Distance +- Tests with cosine distance metric (not just L2) +- Verifies correct distance calculations +- Tests with vectors at known angles + +#### Test 6: Edge Cases +- k larger than number of vectors +- Identical vectors (tie-breaking) +- Boundary conditions + +### Integration Test (`test_recall_integration.py`) + +#### End-to-End Workflow +1. **Connection Verification**: Checks OpenSearch is accessible +2. **Index Creation**: Creates a test index with JVector configuration +3. **Ground Truth Setup**: Pre-generates query vectors +4. **Indexing with Tracking**: Indexes vectors while tracking ground truth +5. **Search Testing**: Performs searches and measures recall +6. **Cleanup**: Removes test index + +#### Customization Options + +```bash +# Test with different parameters +python test_recall_integration.py \ + --host localhost:9200 \ + --dimension 256 \ + --num-vectors 5000 \ + --num-queries 10 \ + --k 20 +``` + +## Understanding Recall Values + +### What is Recall@k? + +Recall@k measures the fraction of true k-nearest neighbors found by approximate search: + +``` +Recall@k = (Number of true neighbors found) / k +``` + +### Expected Recall Values + +| Index Size | Expected Recall | Notes | +|------------|----------------|-------| +| Small (<10K) | >0.95 | Very high accuracy expected | +| Medium (10K-100K) | >0.90 | Good accuracy | +| Large (>100K) | >0.85 | Acceptable for approximate search | + +**Factors affecting recall:** +- Index size (larger = potentially lower recall) +- Vector dimension +- Graph construction parameters +- Query vector distribution +- Distance metric (L2 vs cosine) + +### Interpreting Results + +#### Good Results ✅ +``` +Average Recall@10: 0.9450 +Min Recall@10: 0.9000 +Max Recall@10: 1.0000 +``` +- High average recall (>0.9) +- Consistent across queries (low std dev) +- Minimum recall still acceptable + +#### Concerning Results ⚠️ +``` +Average Recall@10: 0.6500 +Min Recall@10: 0.3000 +Max Recall@10: 0.9000 +``` +- Low average recall (<0.8) +- High variance between queries +- Some queries have very poor recall + +## Troubleshooting + +### Unit Tests Fail + +**Problem:** `test_recall_measurement.py` fails + +**Possible Causes:** +1. Bug in `GroundTruthTracker` implementation +2. Incorrect distance calculation +3. Heap logic error + +**Solution:** +- Check the specific test that failed +- Review the implementation in `jvector_utils/recall_measurement.py` +- Verify distance calculations match expected metric (L2 or cosine) + +### Integration Test Fails to Connect + +**Problem:** "Cannot connect to OpenSearch" + +**Solution:** +```bash +# Check if OpenSearch is running +curl http://localhost:9200 + +# Start OpenSearch if needed +# (depends on your installation method) +``` + +### Low Recall Values + +**Problem:** Recall is consistently <0.8 + +**Possible Causes:** +1. Index too small (not enough vectors for good graph) +2. Incorrect distance metric +3. Graph construction parameters need tuning +4. Bug in ground truth tracking + +**Debugging Steps:** + +1. **Verify ground truth is correct:** + ```bash + # Run unit tests first + python test_recall_measurement.py + ``` + +2. **Test with larger index:** + ```bash + python test_recall_integration.py --num-vectors 10000 + ``` + +3. **Check distance metric matches:** + - Index uses L2 distance + - Tracker should use `space_type='l2'` + +4. **Inspect individual queries:** + - Look at per-query recall values + - Check if some queries have much lower recall + +### Memory Issues + +**Problem:** Out of memory during testing + +**Solution:** +```bash +# Reduce test size +python test_recall_integration.py --num-vectors 500 --dimension 64 + +# Or reduce number of queries +python test_recall_integration.py --num-queries 3 +``` + +## Manual Verification + +You can manually verify recall measurement with a small example: + +```python +import numpy as np +from jvector_utils.recall_measurement import GroundTruthTracker, calculate_recall + +# Create a simple test +query = np.array([0.0, 0.0, 0.0]) +tracker = GroundTruthTracker([query], k=3, space_type='l2') + +# Add vectors at known distances +tracker.update("close", [1.0, 0.0, 0.0]) # distance = 1.0 +tracker.update("medium", [2.0, 0.0, 0.0]) # distance = 2.0 +tracker.update("far", [3.0, 0.0, 0.0]) # distance = 3.0 + +# Get ground truth +ground_truth = tracker.get_ground_truth(0) +print(f"Ground truth: {ground_truth}") +# Expected: ['close', 'medium', 'far'] + +# Test recall calculation +approximate = ["close", "medium", "wrong"] +recall = calculate_recall(approximate, ground_truth) +print(f"Recall: {recall}") +# Expected: 0.6667 (2 out of 3 correct) +``` + +## Best Practices + +### For Development + +1. **Always run unit tests first** before integration tests +2. **Start with small indices** (1K vectors) for faster iteration +3. **Use consistent random seeds** for reproducible tests +4. **Monitor memory usage** when testing with large indices + +### For Production Use + +1. **Validate recall on representative data** before deploying +2. **Set appropriate k values** (typically 10-100) +3. **Pre-generate enough queries** for statistical significance (>20) +4. **Monitor recall over time** as index grows +5. **Document expected recall ranges** for your use case + +## Additional Resources + +- **Package Documentation**: See `jvector_utils/README.md` +- **Main Script Usage**: See `README.md` + +## Summary + +To verify recall measurement is working: + +```bash +# Step 1: Run unit tests (fast, no dependencies) +python test_recall_measurement.py + +# Step 2: Run integration test (requires OpenSearch) +python test_recall_integration.py + +# Step 3: Use with real workload +python create_and_test_large_index.py --measure-recall --num-vectors 100000 +``` + +If all tests pass, recall measurement is working correctly! 🎉 + diff --git a/scripts/jvector_index_and_search/create_and_test_large_index.py b/scripts/jvector_index_and_search/create_and_test_large_index.py new file mode 100755 index 00000000..6f2f81f6 --- /dev/null +++ b/scripts/jvector_index_and_search/create_and_test_large_index.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 + +""" +Create and test large JVector indices in OpenSearch. + +This script provides a comprehensive testing framework for JVector indices, +including indexing, force merging, search testing, and recall measurement. +""" + +import argparse +import numpy as np +import os + +from jvector_utils.index_operations import create_index, index_vectors, force_merge, get_index_stats +from jvector_utils.search_operations import test_search_with_stats +from jvector_utils.recall_measurement import GroundTruthTracker +from jvector_utils.visualization import plot_merge_times + + +def parse_arguments(): + """Parse command line arguments""" + parser = argparse.ArgumentParser(description="Create and test a large JVector index in OpenSearch") + parser.add_argument("--host", default="localhost:9200", help="OpenSearch host:port") + parser.add_argument("--index", default="large-jvector-index", help="Index name") + parser.add_argument("--dimension", type=int, default=768, help="Vector dimension") + parser.add_argument("--num-vectors", type=int, default=3000000, + help="Number of vectors to index (3M vectors with dim=768 should exceed 2GB)") + parser.add_argument("--batch-size", type=int, default=1000, help="Batch size for indexing") + parser.add_argument("--shards", type=int, default=1, help="Number of shards") + parser.add_argument("--num-searches", type=int, default=5, help="Number of searches to perform for stats testing") + parser.add_argument("--skip-indexing", action="store_true", help="Skip index creation and indexing, only run searches") + parser.add_argument("--force-merge-frequency", type=int, default=0, + help="Force merge after every N documents (0 to disable intermediate merges)") + parser.add_argument("--csv-output", type=str, help="CSV file to save merge time data") + parser.add_argument("--plot", action="store_true", help="Generate plots from CSV data") + parser.add_argument("--min-batch-size-for-quantization", type=int, default=1000000, + help="Minimum batch size for quantization (default: 1M)") + parser.add_argument("--measure-recall", action="store_true", + help="Measure recall by computing ground truth incrementally (memory efficient)") + parser.add_argument("--num-recall-queries", type=int, default=None, + help="Number of query vectors to pre-generate for recall measurement (default: same as --num-searches)") + + args = parser.parse_args() + + # Set default for num_recall_queries if not specified + if args.num_recall_queries is None: + args.num_recall_queries = args.num_searches + + return args + + +def setup_recall_tracking(args): + """Set up recall tracking with ground truth tracker + + Args: + args: Parsed command line arguments + + Returns: + GroundTruthTracker instance or None + """ + if not args.measure_recall: + return None + + print(f"\nRecall measurement enabled - pre-generating {args.num_recall_queries} query vectors") + query_vectors = [np.random.uniform(-1, 1, args.dimension) for _ in range(args.num_recall_queries)] + + # Initialize ground truth tracker + ground_truth_tracker = GroundTruthTracker(query_vectors, k=10, space_type='l2') + + # Calculate memory usage + # Each query vector: dimension * 8 bytes + # Each heap entry: ~(8 bytes for distance + 8 bytes for pointer) * k * num_queries + query_mem = args.num_recall_queries * args.dimension * 8 / (1024 * 1024) + heap_mem = args.num_recall_queries * 10 * 16 / (1024 * 1024) # k=10, ~16 bytes per entry + total_mem = query_mem + heap_mem + + print(f"Memory usage for recall tracking:") + print(f" Query vectors: ~{query_mem:.2f} MB") + print(f" Ground truth heaps: ~{heap_mem:.2f} MB") + print(f" Total: ~{total_mem:.2f} MB") + + return ground_truth_tracker + + +def run_indexing(args, ground_truth_tracker=None): + """Run the indexing workflow + + Args: + args: Parsed command line arguments + ground_truth_tracker: Optional GroundTruthTracker for recall measurement + """ + print(f"Creating large JVector index with {args.num_vectors} vectors of dimension {args.dimension}") + print(f"Estimated size: ~{args.num_vectors * args.dimension * 4 / (1024*1024*1024):.2f} GB (raw vectors only)") + + # Create index + create_index(args.host, args.index, args.dimension, args.shards, args.min_batch_size_for_quantization) + + # Index vectors with ground truth tracker + index_vectors(args.host, args.index, args.num_vectors, args.dimension, args.batch_size, + args.force_merge_frequency, args.csv_output, ground_truth_tracker=ground_truth_tracker) + + if args.measure_recall: + print(f"\nGround truth computed for {args.num_recall_queries} query vectors during indexing") + + # Get stats before final force merge + print("\nIndex stats before final force merge:") + get_index_stats(args.host, args.index) + + # Force merge + force_merge(args.host, args.index) + + # Get stats after force merge + print("\nIndex stats after final force merge:") + get_index_stats(args.host, args.index) + + +def run_search_tests(args, ground_truth_tracker=None): + """Run search tests with statistics + + Args: + args: Parsed command line arguments + ground_truth_tracker: Optional GroundTruthTracker for recall measurement + """ + print("\nTesting search with JVector stats:") + + # Limit num_searches to num_recall_queries if using ground truth tracker + actual_num_searches = args.num_searches + if ground_truth_tracker and args.num_searches > args.num_recall_queries: + print(f"\nWarning: Limiting searches to {args.num_recall_queries} (number of pre-generated recall queries)") + actual_num_searches = args.num_recall_queries + + test_search_with_stats(args.host, args.index, args.dimension, k=10, num_searches=actual_num_searches, + ground_truth_tracker=ground_truth_tracker) + + +def main(): + """Main entry point""" + args = parse_arguments() + + # Handle plot-only mode + if args.plot and args.csv_output: + plot_merge_times(args.csv_output) + return + + # Setup recall tracking if enabled + ground_truth_tracker = None + + if not args.skip_indexing: + # Setup recall tracking before indexing + ground_truth_tracker = setup_recall_tracking(args) + + # Run indexing workflow + run_indexing(args, ground_truth_tracker) + else: + print(f"Skipping index creation and indexing. Using existing index: {args.index}") + # Get current index stats + print("\nCurrent index stats:") + get_index_stats(args.host, args.index) + + if args.measure_recall: + print("\nWarning: Recall measurement requires indexing to track ground truth.") + print("Cannot measure recall when using --skip-indexing flag.") + + # Run search tests + run_search_tests(args, ground_truth_tracker) + + print("\nTest completed successfully!") + + # Generate plots if CSV output was specified + if args.csv_output and os.path.exists(args.csv_output): + print(f"\nCSV data saved to {args.csv_output}") + plot_merge_times(args.csv_output) + + +if __name__ == "__main__": + main() + diff --git a/scripts/jvector_index_and_search/jvector_utils/README.md b/scripts/jvector_index_and_search/jvector_utils/README.md new file mode 100644 index 00000000..21bca829 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/README.md @@ -0,0 +1,203 @@ +# JVector Utils + +A modular Python package for testing and benchmarking JVector indices in OpenSearch. + +## Overview + +This package provides a clean, modular interface for working with JVector indices, including: +- Index creation and management +- Efficient vector indexing with batch operations +- Search testing with detailed statistics +- Memory-efficient recall measurement +- Performance visualization + +## Package Structure + +``` +jvector_utils/ +├── __init__.py # Package initialization +├── index_operations.py # Index creation, indexing, force merge +├── search_operations.py # Search testing and statistics +├── recall_measurement.py # Ground truth tracking and recall calculation +├── stats_utils.py # KNN statistics retrieval and display +└── visualization.py # Performance plotting utilities +``` + +## Modules + +### `index_operations.py` + +Functions for creating and managing JVector indices: + +- **`create_index(host, index_name, dimension, shards, min_batch_size_for_quantization)`** + - Creates a new JVector index with specified configuration + +- **`index_vectors(host, index_name, num_vectors, dimension, batch_size, ...)`** + - Indexes random vectors in batches + - Supports ground truth tracking for recall measurement + - Optional intermediate force merges + +- **`force_merge(host, index_name, max_segments)`** + - Consolidates index segments + - Tracks graph merge and quantization times + +- **`get_index_stats(host, index_name)`** + - Retrieves and displays index size statistics + +### `search_operations.py` + +Functions for testing search performance: + +- **`test_search(host, index_name, dimension, k, query_vector)`** + - Performs a single kNN search + - Returns results, timing, and query vector + +- **`test_search_with_stats(host, index_name, dimension, k, num_searches, ground_truth_tracker)`** + - Runs multiple search iterations + - Collects detailed JVector statistics + - Calculates recall if ground truth tracker is provided + - Reports per-iteration and aggregate metrics + +### `recall_measurement.py` + +Memory-efficient recall measurement: + +- **`GroundTruthTracker` class** + - Tracks exact k-nearest neighbors using min-heaps + - Updates incrementally during indexing + - Memory usage: O(num_queries × k) instead of O(num_vectors × dimension) + - Supports L2 and cosine distance metrics + +- **`calculate_recall(approximate_results, ground_truth)`** + - Computes recall@k metric + - Returns fraction of true neighbors found + +### `stats_utils.py` + +KNN statistics utilities: + +- **`get_knn_stats(host)`** + - Retrieves KNN plugin statistics from OpenSearch + +- **`print_jvector_search_stats(stats)`** + - Displays JVector-specific search metrics + - Aggregates stats across all nodes + +- **`get_knn_stat_value(stats, stat_name)`** + - Extracts a specific metric value + +### `visualization.py` + +Performance visualization: + +- **`plot_merge_times(csv_file)`** + - Generates plots of graph merge time vs. document count + - Plots quantization training time vs. document count + - Saves high-resolution PNG output + +## Usage Examples + +### Basic Index Creation and Search + +```python +from jvector_utils.index_operations import create_index, index_vectors +from jvector_utils.search_operations import test_search + +# Create index +create_index("localhost:9200", "my-index", dimension=768, shards=1) + +# Index vectors +index_vectors("localhost:9200", "my-index", num_vectors=10000, dimension=768) + +# Test search +success, results, query, duration = test_search("localhost:9200", "my-index", dimension=768, k=10) +``` + +### Recall Measurement + +```python +import numpy as np +from jvector_utils.recall_measurement import GroundTruthTracker +from jvector_utils.index_operations import index_vectors +from jvector_utils.search_operations import test_search_with_stats + +# Pre-generate query vectors +query_vectors = [np.random.uniform(-1, 1, 768) for _ in range(10)] + +# Initialize tracker +tracker = GroundTruthTracker(query_vectors, k=10, space_type='l2') + +# Index with ground truth tracking +index_vectors("localhost:9200", "my-index", num_vectors=100000, dimension=768, + ground_truth_tracker=tracker) + +# Test search with recall measurement +test_search_with_stats("localhost:9200", "my-index", dimension=768, k=10, + num_searches=10, ground_truth_tracker=tracker) +``` + +### Performance Visualization + +```python +from jvector_utils.visualization import plot_merge_times + +# Generate plots from CSV data +plot_merge_times("merge_times.csv") +``` + +## Key Features + +### Memory-Efficient Recall Measurement + +The `GroundTruthTracker` uses a heap-based algorithm that: +- Pre-generates query vectors before indexing +- Maintains min-heaps during indexing (one per query) +- Only stores k-nearest neighbors per query +- Achieves ~1000x memory reduction compared to storing all vectors + +**Memory Usage:** +``` +Query vectors: num_queries × dimension × 8 bytes +Ground truth heaps: num_queries × k × 16 bytes +Total: ~6 MB for 1000 queries × 768 dimensions +``` + +Compare to storing all vectors: +``` +All vectors: 1,000,000 × 768 × 8 bytes ≈ 5,859 MB (~6 GB) +``` + +### Detailed Statistics Tracking + +JVector-specific metrics tracked: +- `knn_query_visited_nodes` - Total nodes visited during search +- `knn_query_expanded_nodes` - Nodes expanded during search +- `knn_query_expanded_base_layer_nodes` - Base layer nodes expanded + +Statistics are reported: +- Per search iteration +- As aggregate totals +- As per-search averages + +## Design Principles + +1. **Modularity**: Each module has a single, clear responsibility +2. **Reusability**: Functions can be imported and used independently +3. **Testability**: Small, focused functions are easy to test +4. **Readability**: Clear naming and comprehensive documentation +5. **Efficiency**: Memory-efficient algorithms for large-scale testing + +## Dependencies + +- `requests` - HTTP client for OpenSearch API +- `numpy` - Numerical operations and vector generation +- `matplotlib` - Performance visualization +- Python 3.7+ + +## Integration + +This package is designed to work seamlessly with the `create_and_test_large_index.py` script, +which provides a command-line interface to all functionality. + +See `../README.md` for usage examples of the main script. + diff --git a/scripts/jvector_index_and_search/jvector_utils/__init__.py b/scripts/jvector_index_and_search/jvector_utils/__init__.py new file mode 100644 index 00000000..590a3825 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/__init__.py @@ -0,0 +1,13 @@ +""" +JVector utilities for OpenSearch testing and benchmarking. + +This package provides utilities for: +- Index creation and management +- Vector indexing operations +- Search testing and statistics +- Recall measurement +- Visualization of performance metrics +""" + +__version__ = "1.0.0" + diff --git a/scripts/jvector_index_and_search/jvector_utils/index_operations.py b/scripts/jvector_index_and_search/jvector_utils/index_operations.py new file mode 100644 index 00000000..2d9876e8 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/index_operations.py @@ -0,0 +1,201 @@ +""" +Index operations for JVector in OpenSearch. + +This module provides functions for creating indices, indexing vectors, +and performing force merges with detailed statistics tracking. +""" + +import requests +import json +import time +import numpy as np +import sys +import csv + +from .stats_utils import get_knn_stats, get_knn_stat_value + + +def create_index(host, index_name, dimension, shards=1, min_batch_size_for_quantization=1000000): + """Create a knn index with jvector engine""" + url = f"http://{host}/{index_name}" + + mapping = { + "settings": { + "index": { + "knn": True, + "knn.derived_source.enabled": True, + "number_of_shards": shards, + "number_of_replicas": 0 + } + }, + "mappings": { + "properties": { + "vector_field": { + "type": "knn_vector", + "dimension": dimension, + "method": { + "name": "disk_ann", + "space_type": "l2", + "engine": "jvector", + "parameters": { + "advanced.min_batch_size_for_quantization": min_batch_size_for_quantization + } + } + }, + "id": {"type": "keyword"} + } + } + } + + response = requests.put(url, json=mapping) + if response.status_code != 200: + print(f"Failed to create index: {response.text}") + sys.exit(1) + print(f"Successfully created index {index_name}") + return response.json() + + +def index_vectors(host, index_name, num_vectors, dimension, batch_size=1000, + force_merge_frequency=0, csv_file=None, ground_truth_tracker=None): + """Index vectors in batches + + Args: + host: OpenSearch host:port + index_name: Name of the index + num_vectors: Total number of vectors to index + dimension: Vector dimension + batch_size: Number of vectors per batch + force_merge_frequency: Force merge after every N documents (0 to disable) + csv_file: Optional CSV file to save merge statistics + ground_truth_tracker: Optional GroundTruthTracker to update during indexing + """ + url = f"http://{host}/{index_name}/_bulk" + headers = {"Content-Type": "application/x-ndjson"} + + total_batches = (num_vectors + batch_size - 1) // batch_size + + # Initialize CSV file if provided + csv_writer = None + csv_file_handle = None + if csv_file: + csv_file_handle = open(csv_file, 'w', newline='') + csv_writer = csv.writer(csv_file_handle) + csv_writer.writerow(['num_documents', 'graph_merge_time_ms', 'quantization_training_time_ms', + 'force_merge_duration_sec', 'index_size_bytes']) + + for batch in range(total_batches): + start_idx = batch * batch_size + end_idx = min((batch + 1) * batch_size, num_vectors) + current_batch_size = end_idx - start_idx + + bulk_data = [] + for i in range(current_batch_size): + doc_id = start_idx + i + # Create action line + action = {"index": {"_index": index_name, "_id": str(doc_id)}} + bulk_data.append(json.dumps(action)) + + # Create random vector with values between -1 and 1 + vector = np.random.uniform(-1, 1, dimension).tolist() + document = {"vector_field": vector, "id": str(doc_id)} + bulk_data.append(json.dumps(document)) + + # Update ground truth tracker if provided + if ground_truth_tracker: + ground_truth_tracker.update(str(doc_id), vector) + + bulk_body = "\n".join(bulk_data) + "\n" + + response = requests.post(url, headers=headers, data=bulk_body) + if response.status_code != 200: + print(f"Failed to index batch {batch+1}/{total_batches}: {response.text}") + sys.exit(1) + + print(f"Indexed batch {batch+1}/{total_batches} ({current_batch_size} vectors)") + + # Force merge if frequency is set and we've reached the threshold + if force_merge_frequency > 0 and (end_idx % force_merge_frequency < batch_size): + print(f"\nPerforming intermediate force merge after {end_idx} documents...") + merge_result = force_merge(host, index_name) + print(f"Index stats after intermediate force merge:") + stats = get_index_stats(host, index_name) + + # Write to CSV if enabled + if csv_writer and merge_result: + index_size = stats["indices"][index_name]["total"]["store"]["size_in_bytes"] if stats else 0 + csv_writer.writerow([ + end_idx, + merge_result['graph_merge_time'], + merge_result['quantization_time'], + merge_result['duration'], + index_size + ]) + csv_file_handle.flush() + + if csv_file_handle: + csv_file_handle.close() + + return True + + +def force_merge(host, index_name, max_segments=1): + """Force merge the index to consolidate segments""" + # Get initial KNN stats + initial_stats = get_knn_stats(host) + initial_graph_merge_time = get_knn_stat_value(initial_stats, "knn_graph_merge_time") + initial_quantization_time = get_knn_stat_value(initial_stats, "knn_quantization_training_time") + + # Refresh first to ensure all documents are searchable + refresh_url = f"http://{host}/{index_name}/_refresh" + refresh_response = requests.post(refresh_url) + if refresh_response.status_code != 200: + print(f"Refresh failed: {refresh_response.text}") + + url = f"http://{host}/{index_name}/_forcemerge?max_num_segments={max_segments}&flush=true" + + print(f"Starting force merge to {max_segments} segments...") + start_time = time.time() + response = requests.post(url) + + if response.status_code != 200: + print(f"Force merge failed: {response.text}") + return False + + duration = time.time() - start_time + print(f"Force merge completed in {duration:.2f} seconds") + + # Get final KNN stats + final_stats = get_knn_stats(host) + final_graph_merge_time = get_knn_stat_value(final_stats, "knn_graph_merge_time") + final_quantization_time = get_knn_stat_value(final_stats, "knn_quantization_training_time") + + # Calculate and display the differences + graph_merge_diff = final_graph_merge_time - initial_graph_merge_time + quantization_diff = final_quantization_time - initial_quantization_time + + print(f"KNN Graph Merge Time: +{graph_merge_diff} ms") + print(f"KNN Quantization Training Time: +{quantization_diff} ms") + + return { + 'graph_merge_time': graph_merge_diff, + 'quantization_time': quantization_diff, + 'duration': duration + } + + +def get_index_stats(host, index_name): + """Get index stats including size""" + url = f"http://{host}/{index_name}/_stats" + response = requests.get(url) + + if response.status_code != 200: + print(f"Failed to get index stats: {response.text}") + return None + + stats = response.json() + size_in_bytes = stats["indices"][index_name]["total"]["store"]["size_in_bytes"] + size_in_gb = size_in_bytes / (1024 * 1024 * 1024) + + print(f"Index size: {size_in_bytes} bytes ({size_in_gb:.2f} GB)") + return stats + diff --git a/scripts/jvector_index_and_search/jvector_utils/recall_measurement.py b/scripts/jvector_index_and_search/jvector_utils/recall_measurement.py new file mode 100644 index 00000000..c4ae7b21 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/recall_measurement.py @@ -0,0 +1,109 @@ +""" +Recall measurement utilities for evaluating approximate nearest neighbor search quality. + +This module provides efficient recall calculation using incremental ground truth tracking +with min-heaps, avoiding the need to store all vectors in memory. +""" + +import numpy as np +import heapq + + +class GroundTruthTracker: + """Tracks ground truth for pre-generated query vectors using min-heaps + + This allows computing ground truth incrementally during indexing without + storing all vectors in memory. + """ + + def __init__(self, query_vectors, k, space_type='l2'): + """Initialize tracker with query vectors + + Args: + query_vectors: List of query vectors to track ground truth for + k: Number of nearest neighbors to track + space_type: Distance metric ('l2' or 'cosine') + """ + self.query_vectors = query_vectors + self.k = k + self.space_type = space_type + + # For each query, maintain a max-heap of size k + # We use max-heap so we can efficiently remove the farthest neighbor + # when we find a closer one + # Heap stores tuples of (-distance, doc_id) - negative for max-heap behavior + self.heaps = [[] for _ in query_vectors] + + def update(self, doc_id, vector): + """Update ground truth with a new vector + + Args: + doc_id: Document ID + vector: Vector to add + """ + vector_np = np.array(vector) + + for i, query_vector in enumerate(self.query_vectors): + # Compute distance + if self.space_type == 'l2': + dist = np.linalg.norm(query_vector - vector_np) + elif self.space_type == 'cosine': + query_norm = np.linalg.norm(query_vector) + vector_norm = np.linalg.norm(vector_np) + if query_norm > 0 and vector_norm > 0: + cosine_sim = np.dot(query_vector, vector_np) / (query_norm * vector_norm) + dist = 1 - cosine_sim + else: + dist = 1.0 + else: + raise ValueError(f"Unsupported space_type: {self.space_type}") + + heap = self.heaps[i] + + # If heap is not full, add the item + if len(heap) < self.k: + heapq.heappush(heap, (-dist, doc_id)) + # If this distance is smaller than the largest in heap, replace it + elif dist < -heap[0][0]: + heapq.heapreplace(heap, (-dist, doc_id)) + + def get_ground_truth(self, query_index): + """Get ground truth for a specific query + + Args: + query_index: Index of the query vector + + Returns: + List of doc_ids of the k nearest neighbors + """ + heap = self.heaps[query_index] + # Sort by distance (ascending) and return doc_ids + sorted_results = sorted(heap, key=lambda x: -x[0]) + return [doc_id for _, doc_id in sorted_results] + + def get_query_vector(self, query_index): + """Get the query vector at the given index""" + return self.query_vectors[query_index] + + +def calculate_recall(approximate_results, ground_truth): + """Calculate recall@k + + Args: + approximate_results: List of doc_ids from approximate search + ground_truth: List of doc_ids from exact search + + Returns: + Recall value (0.0 to 1.0) + """ + if not ground_truth: + return 0.0 + + approximate_set = set(approximate_results) + ground_truth_set = set(ground_truth) + + intersection = approximate_set.intersection(ground_truth_set) + recall = len(intersection) / len(ground_truth) + + return recall + diff --git a/scripts/jvector_index_and_search/jvector_utils/search_operations.py b/scripts/jvector_index_and_search/jvector_utils/search_operations.py new file mode 100644 index 00000000..7f720205 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/search_operations.py @@ -0,0 +1,201 @@ +""" +Search operations and testing utilities for JVector. + +This module provides functions for testing kNN search performance, +collecting detailed statistics, and measuring recall. +""" + +import requests +import time +import numpy as np + +from .stats_utils import get_knn_stats, print_jvector_search_stats +from .recall_measurement import calculate_recall + + +def test_search(host, index_name, dimension, k=10, query_vector=None): + """Test kNN search on the index + + Args: + host: OpenSearch host:port + index_name: Name of the index + dimension: Vector dimension + k: Number of nearest neighbors to retrieve + query_vector: Optional pre-generated query vector. If None, a random one is created. + + Returns: + Tuple of (success, results_list, query_vector, duration) + """ + url = f"http://{host}/{index_name}/_search" + + # Create random query vector if not provided + if query_vector is None: + query_vector = np.random.uniform(-1, 1, dimension).tolist() + + query = { + "size": k, + "query": { + "knn": { + "vector_field": { + "vector": query_vector, + "k": k + } + } + } + } + + start_time = time.time() + response = requests.post(url, json=query) + duration = time.time() - start_time + + if response.status_code != 200: + print(f"Search failed: {response.text}") + return False, [], query_vector, duration + + results = response.json() + hits = results["hits"]["hits"] + result_ids = [hit["_id"] for hit in hits] + + print(f"Search completed in {duration:.4f} seconds, found {len(hits)} results") + return True, result_ids, query_vector, duration + + +def test_search_with_stats(host, index_name, dimension, k=10, num_searches=5, ground_truth_tracker=None): + """Test kNN search on the index and report JVector stats for each iteration + + Args: + host: OpenSearch host:port + index_name: Name of the index + dimension: Vector dimension + k: Number of nearest neighbors to retrieve + num_searches: Number of search iterations to perform + ground_truth_tracker: Optional GroundTruthTracker for recall calculation + """ + # Get initial stats + current_stats = get_knn_stats(host) + print("\nInitial JVector Stats:") + print_jvector_search_stats(current_stats) + + # Track recall if tracker is provided + measure_recall = ground_truth_tracker is not None + recall_values = [] + + # JVector-specific metrics we want to track + jvector_metrics = [ + "knn_query_visited_nodes", + "knn_query_expanded_nodes", + "knn_query_expanded_base_layer_nodes" + ] + + # Store initial values + prev_totals = {} + if current_stats and "nodes" in current_stats: + for node_id, node_stats in current_stats["nodes"].items(): + for metric in jvector_metrics: + if metric in node_stats: + if metric not in prev_totals: + prev_totals[metric] = 0 + prev_totals[metric] += int(node_stats[metric]) + + # Perform multiple searches and check stats after each + for i in range(num_searches): + print(f"\n--- Search Iteration {i+1}/{num_searches} ---") + + # If using tracker, use the pre-generated query vector + query_vector = None + if measure_recall: + query_vector = ground_truth_tracker.get_query_vector(i).tolist() + + # Perform search + success, result_ids, query_vector_used, duration = test_search(host, index_name, dimension, k, query_vector) + + # Calculate recall if enabled + if measure_recall and success: + # Use pre-computed ground truth from tracker + ground_truth = ground_truth_tracker.get_ground_truth(i) + recall = calculate_recall(result_ids, ground_truth) + recall_values.append(recall) + print(f"Recall@{k}: {recall:.4f} ({len(set(result_ids).intersection(set(ground_truth)))}/{k} correct)") + + # Get stats after this search + new_stats = get_knn_stats(host) + + # Calculate current totals + current_totals = {} + if new_stats and "nodes" in new_stats: + for node_id, node_stats in new_stats["nodes"].items(): + for metric in jvector_metrics: + if metric in node_stats: + if metric not in current_totals: + current_totals[metric] = 0 + current_totals[metric] += int(node_stats[metric]) + + # Calculate and print differences for this iteration + print("\nJVector Stats for this iteration:") + for metric in jvector_metrics: + prev_val = prev_totals.get(metric, 0) + current_val = current_totals.get(metric, 0) + diff = current_val - prev_val + print(f" {metric}: +{diff}") + + # Update previous totals for next iteration + prev_totals = current_totals + + # Get final stats + final_stats = get_knn_stats(host) + + # Calculate initial and final totals + initial_totals = {} + if current_stats and "nodes" in current_stats: + for node_id, node_stats in current_stats["nodes"].items(): + for metric in jvector_metrics: + if metric in node_stats: + if metric not in initial_totals: + initial_totals[metric] = 0 + initial_totals[metric] += int(node_stats[metric]) + + final_totals = {} + if final_stats and "nodes" in final_stats: + for node_id, node_stats in final_stats["nodes"].items(): + for metric in jvector_metrics: + if metric in node_stats: + if metric not in final_totals: + final_totals[metric] = 0 + final_totals[metric] += int(node_stats[metric]) + + # Print summary + print("\n=== JVector Stats Summary ===") + print("\nInitial Stats:") + for metric, value in initial_totals.items(): + print(f" {metric}: {value}") + + print("\nFinal Stats:") + for metric, value in final_totals.items(): + print(f" {metric}: {value}") + + print("\nTotal Differences (Final - Initial):") + for metric in jvector_metrics: + initial_val = initial_totals.get(metric, 0) + final_val = final_totals.get(metric, 0) + diff = final_val - initial_val + print(f" {metric}: +{diff}") + + # Calculate per-search averages + print("\nAverage per Search:") + for metric in jvector_metrics: + initial_val = initial_totals.get(metric, 0) + final_val = final_totals.get(metric, 0) + diff = final_val - initial_val + avg = diff / num_searches if num_searches > 0 else 0 + print(f" {metric}: {avg:.2f}") + + # Print recall summary if measured + if measure_recall and recall_values: + print("\n=== Recall Summary ===") + print(f"Average Recall@{k}: {np.mean(recall_values):.4f}") + print(f"Min Recall@{k}: {np.min(recall_values):.4f}") + print(f"Max Recall@{k}: {np.max(recall_values):.4f}") + print(f"Std Dev: {np.std(recall_values):.4f}") + + return True + diff --git a/scripts/jvector_index_and_search/jvector_utils/stats_utils.py b/scripts/jvector_index_and_search/jvector_utils/stats_utils.py new file mode 100644 index 00000000..7071c1c9 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/stats_utils.py @@ -0,0 +1,60 @@ +""" +Statistics utilities for JVector KNN operations. + +This module provides functions for retrieving and displaying KNN statistics +from OpenSearch, including JVector-specific search metrics. +""" + +import requests + + +def get_knn_stats(host): + """Get KNN plugin stats including JVector search metrics""" + url = f"http://{host}/_plugins/_knn/stats" + response = requests.get(url) + + if response.status_code != 200: + print(f"Failed to get KNN stats: {response.text}") + return None + + stats = response.json() + return stats + + +def print_jvector_search_stats(stats): + """Extract and print JVector-specific search statistics""" + if not stats or "nodes" not in stats: + print("No KNN stats available") + return + + # JVector-specific metrics we want to track + jvector_metrics = [ + "knn_query_visited_nodes", + "knn_query_expanded_nodes", + "knn_query_expanded_base_layer_nodes" + ] + + # Collect stats from all nodes + all_nodes_stats = {} + for node_id, node_stats in stats["nodes"].items(): + for metric in jvector_metrics: + if metric in node_stats: + if metric not in all_nodes_stats: + all_nodes_stats[metric] = 0 + all_nodes_stats[metric] += int(node_stats[metric]) + + # Print the stats + print("\nJVector Search Statistics:") + for metric, value in all_nodes_stats.items(): + print(f" {metric}: {value}") + + +def get_knn_stat_value(stats, stat_name): + """Extract a specific KNN stat value from all nodes""" + total = 0 + if stats and "nodes" in stats: + for node_id, node_stats in stats["nodes"].items(): + if stat_name in node_stats: + total += int(node_stats[stat_name]) + return total + diff --git a/scripts/jvector_index_and_search/jvector_utils/visualization.py b/scripts/jvector_index_and_search/jvector_utils/visualization.py new file mode 100644 index 00000000..6236a2b1 --- /dev/null +++ b/scripts/jvector_index_and_search/jvector_utils/visualization.py @@ -0,0 +1,61 @@ +""" +Visualization utilities for JVector performance metrics. + +This module provides plotting functions for analyzing merge times, +quantization performance, and other JVector metrics. +""" + +import csv +import os +import matplotlib.pyplot as plt + + +def plot_merge_times(csv_file): + """Plot graph merge time and quantization training time vs number of documents""" + if not os.path.exists(csv_file): + print(f"CSV file {csv_file} not found") + return + + # Read CSV data + num_docs = [] + graph_merge_times = [] + quantization_times = [] + + with open(csv_file, 'r') as f: + reader = csv.DictReader(f) + for row in reader: + num_docs.append(int(row['num_documents'])) + graph_merge_times.append(int(row['graph_merge_time_ms'])) + quantization_times.append(int(row['quantization_training_time_ms'])) + + if not num_docs: + print("No data found in CSV file") + return + + # Create plots + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8)) + + # Graph merge time plot + ax1.plot(num_docs, graph_merge_times, 'b-o', linewidth=2, markersize=6) + ax1.set_xlabel('Number of Documents') + ax1.set_ylabel('Graph Merge Time (ms)') + ax1.set_title('JVector Graph Merge Time vs Number of Documents') + ax1.grid(True, alpha=0.3) + + # Quantization training time plot + ax2.plot(num_docs, quantization_times, 'r-o', linewidth=2, markersize=6) + ax2.set_xlabel('Number of Documents') + ax2.set_ylabel('Quantization Training Time (ms)') + ax2.set_title('JVector Quantization Training Time vs Number of Documents') + ax2.grid(True, alpha=0.3) + + plt.tight_layout() + + # Save plot + plot_file = csv_file.replace('.csv', '_plot.png') + plt.savefig(plot_file, dpi=300, bbox_inches='tight') + print(f"Plot saved to {plot_file}") + + # Show plot + plt.show() + diff --git a/scripts/jvector_index_and_search/test_recall_integration.py b/scripts/jvector_index_and_search/test_recall_integration.py new file mode 100644 index 00000000..97bc6040 --- /dev/null +++ b/scripts/jvector_index_and_search/test_recall_integration.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 + +""" +Integration test for recall measurement with OpenSearch. + +This script creates a small test index, measures recall, and verifies +that the recall measurement is working correctly end-to-end. +""" + +import argparse +import sys +import numpy as np +import requests + +from jvector_utils.index_operations import create_index, index_vectors, force_merge +from jvector_utils.search_operations import test_search_with_stats +from jvector_utils.recall_measurement import GroundTruthTracker + + +def cleanup_index(host, index_name): + """Delete the test index if it exists""" + url = f"http://{host}/{index_name}" + response = requests.delete(url) + if response.status_code in [200, 404]: + return True + print(f"Warning: Failed to delete index: {response.text}") + return False + + +def verify_opensearch_connection(host): + """Verify that OpenSearch is accessible""" + try: + response = requests.get(f"http://{host}") + if response.status_code == 200: + info = response.json() + print(f"✅ Connected to OpenSearch {info.get('version', {}).get('number', 'unknown')}") + return True + else: + print(f"❌ Failed to connect to OpenSearch: {response.status_code}") + return False + except Exception as e: + print(f"❌ Failed to connect to OpenSearch: {e}") + return False + + +def test_recall_with_small_index(host, dimension=128, num_vectors=1000, num_queries=5, k=10): + """ + Test recall measurement with a small index. + + This test: + 1. Creates a small index with known vectors + 2. Tracks ground truth during indexing + 3. Performs searches and measures recall + 4. Verifies that recall is reasonable (should be high for small index) + """ + index_name = "test-recall-index" + + print("\n" + "=" * 70) + print("RECALL INTEGRATION TEST") + print("=" * 70) + print(f"Host: {host}") + print(f"Index: {index_name}") + print(f"Dimension: {dimension}") + print(f"Vectors: {num_vectors}") + print(f"Queries: {num_queries}") + print(f"k: {k}") + print("=" * 70) + + # Step 1: Verify OpenSearch connection + print("\n[1/6] Verifying OpenSearch connection...") + if not verify_opensearch_connection(host): + print("❌ Cannot connect to OpenSearch. Is it running?") + return False + + # Step 2: Cleanup any existing test index + print("\n[2/6] Cleaning up any existing test index...") + cleanup_index(host, index_name) + + # Step 3: Create index + print("\n[3/6] Creating test index...") + try: + create_index(host, index_name, dimension, shards=1, min_batch_size_for_quantization=10000000) + except Exception as e: + print(f"❌ Failed to create index: {e}") + return False + + # Step 4: Setup recall tracking + print(f"\n[4/6] Setting up recall tracking with {num_queries} query vectors...") + query_vectors = [np.random.uniform(-1, 1, dimension) for _ in range(num_queries)] + tracker = GroundTruthTracker(query_vectors, k=k, space_type='l2') + + # Step 5: Index vectors with ground truth tracking + print(f"\n[5/6] Indexing {num_vectors} vectors with ground truth tracking...") + try: + index_vectors(host, index_name, num_vectors, dimension, batch_size=100, + force_merge_frequency=0, csv_file=None, ground_truth_tracker=tracker) + + # Force merge to ensure all vectors are searchable + print("\nForce merging index...") + force_merge(host, index_name) + + except Exception as e: + print(f"❌ Failed to index vectors: {e}") + cleanup_index(host, index_name) + return False + + # Step 6: Test search with recall measurement + print(f"\n[6/6] Testing search with recall measurement...") + try: + test_search_with_stats(host, index_name, dimension, k=k, num_searches=num_queries, + ground_truth_tracker=tracker) + except Exception as e: + print(f"❌ Failed to test search: {e}") + import traceback + traceback.print_exc() + cleanup_index(host, index_name) + return False + + # Cleanup + print("\n[Cleanup] Deleting test index...") + cleanup_index(host, index_name) + + print("\n" + "=" * 70) + print("✅ INTEGRATION TEST COMPLETED SUCCESSFULLY") + print("=" * 70) + print("\nKey observations:") + print("1. Ground truth was tracked incrementally during indexing") + print("2. Recall was measured for each search query") + print("3. Recall values should be high (>0.9) for this small index") + print("4. If recall is low, there may be an issue with the implementation") + print("\nNote: Some variation in recall is expected due to:") + print(" - Approximate nature of JVector search") + print(" - Random vector generation") + print(" - Graph construction parameters") + + return True + + +def main(): + parser = argparse.ArgumentParser(description="Integration test for recall measurement") + parser.add_argument("--host", default="localhost:9200", help="OpenSearch host:port") + parser.add_argument("--dimension", type=int, default=128, help="Vector dimension") + parser.add_argument("--num-vectors", type=int, default=1000, help="Number of vectors to index") + parser.add_argument("--num-queries", type=int, default=5, help="Number of query vectors") + parser.add_argument("--k", type=int, default=10, help="Number of nearest neighbors") + + args = parser.parse_args() + + success = test_recall_with_small_index( + args.host, + dimension=args.dimension, + num_vectors=args.num_vectors, + num_queries=args.num_queries, + k=args.k + ) + + return 0 if success else 1 + + +if __name__ == "__main__": + sys.exit(main()) + diff --git a/scripts/jvector_index_and_search/test_recall_measurement.py b/scripts/jvector_index_and_search/test_recall_measurement.py new file mode 100755 index 00000000..197c9906 --- /dev/null +++ b/scripts/jvector_index_and_search/test_recall_measurement.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 + +""" +Test suite for recall measurement functionality. + +This script tests the GroundTruthTracker and recall calculation to ensure +they work correctly before using them with real OpenSearch indices. +""" + +import numpy as np +import sys +from jvector_utils.recall_measurement import GroundTruthTracker, calculate_recall + + +def test_ground_truth_tracker_basic(): + """Test basic functionality of GroundTruthTracker""" + print("\n=== Test 1: Basic GroundTruthTracker Functionality ===") + + # Create a simple test case with known vectors + dimension = 3 + k = 3 + + # Pre-generate a single query vector + query_vector = np.array([0.0, 0.0, 0.0]) + + # Initialize tracker + tracker = GroundTruthTracker([query_vector], k=k, space_type='l2') + + # Add vectors at known distances from the query + # Vector at distance 1.0 + tracker.update("id_1", [1.0, 0.0, 0.0]) + + # Vector at distance 2.0 + tracker.update("id_2", [0.0, 2.0, 0.0]) + + # Vector at distance 3.0 + tracker.update("id_3", [0.0, 0.0, 3.0]) + + # Vector at distance sqrt(2) ≈ 1.414 + tracker.update("id_4", [1.0, 1.0, 0.0]) + + # Vector at distance sqrt(3) ≈ 1.732 + tracker.update("id_5", [1.0, 1.0, 1.0]) + + # Get ground truth (should be the 3 closest) + ground_truth = tracker.get_ground_truth(0) + + print(f"Query vector: {query_vector}") + print(f"Ground truth (k={k}): {ground_truth}") + + # Expected: id_1 (dist=1.0), id_4 (dist≈1.414), id_5 (dist≈1.732) + expected = ["id_1", "id_4", "id_5"] + + if set(ground_truth) == set(expected): + print("✅ PASS: Ground truth matches expected nearest neighbors") + return True + else: + print(f"❌ FAIL: Expected {expected}, got {ground_truth}") + return False + + +def test_ground_truth_tracker_large(): + """Test GroundTruthTracker with more vectors than k""" + print("\n=== Test 2: GroundTruthTracker with Many Vectors ===") + + dimension = 10 + k = 5 + num_vectors = 100 + + # Create a random query vector + query_vector = np.random.uniform(-1, 1, dimension) + + # Initialize tracker + tracker = GroundTruthTracker([query_vector], k=k, space_type='l2') + + # Add many random vectors and compute ground truth manually + vectors = {} + for i in range(num_vectors): + vec = np.random.uniform(-1, 1, dimension) + doc_id = f"doc_{i}" + vectors[doc_id] = vec + tracker.update(doc_id, vec.tolist()) + + # Get ground truth from tracker + ground_truth = tracker.get_ground_truth(0) + + # Compute ground truth manually using brute force + distances = [] + for doc_id, vec in vectors.items(): + dist = np.sum((query_vector - vec) ** 2) # L2 distance squared + distances.append((dist, doc_id)) + + distances.sort() + expected_ground_truth = [doc_id for _, doc_id in distances[:k]] + + print(f"Number of vectors: {num_vectors}") + print(f"k: {k}") + print(f"Ground truth from tracker: {ground_truth}") + print(f"Expected ground truth: {expected_ground_truth}") + + if set(ground_truth) == set(expected_ground_truth): + print("✅ PASS: Ground truth matches brute-force computation") + return True + else: + print(f"❌ FAIL: Mismatch in ground truth") + print(f" Missing: {set(expected_ground_truth) - set(ground_truth)}") + print(f" Extra: {set(ground_truth) - set(expected_ground_truth)}") + return False + + +def test_calculate_recall(): + """Test the calculate_recall function""" + print("\n=== Test 3: Recall Calculation ===") + + # Test perfect recall + ground_truth = ["1", "2", "3", "4", "5"] + approximate = ["1", "2", "3", "4", "5"] + recall = calculate_recall(approximate, ground_truth) + print(f"Perfect recall: {recall} (expected 1.0)") + + if recall != 1.0: + print("❌ FAIL: Perfect recall should be 1.0") + return False + + # Test 80% recall + approximate = ["1", "2", "3", "4", "6"] # 4 out of 5 correct + recall = calculate_recall(approximate, ground_truth) + print(f"80% recall: {recall} (expected 0.8)") + + if recall != 0.8: + print(f"❌ FAIL: Expected 0.8, got {recall}") + return False + + # Test 0% recall + approximate = ["6", "7", "8", "9", "10"] + recall = calculate_recall(approximate, ground_truth) + print(f"0% recall: {recall} (expected 0.0)") + + if recall != 0.0: + print(f"❌ FAIL: Expected 0.0, got {recall}") + return False + + print("✅ PASS: All recall calculations correct") + return True + + +def test_multiple_queries(): + """Test GroundTruthTracker with multiple query vectors""" + print("\n=== Test 4: Multiple Query Vectors ===") + + dimension = 5 + k = 3 + num_queries = 3 + num_vectors = 50 + + # Create multiple query vectors + query_vectors = [np.random.uniform(-1, 1, dimension) for _ in range(num_queries)] + + # Initialize tracker + tracker = GroundTruthTracker(query_vectors, k=k, space_type='l2') + + # Add vectors + vectors = {} + for i in range(num_vectors): + vec = np.random.uniform(-1, 1, dimension) + doc_id = f"doc_{i}" + vectors[doc_id] = vec + tracker.update(doc_id, vec.tolist()) + + # Verify ground truth for each query + all_pass = True + for query_idx in range(num_queries): + ground_truth = tracker.get_ground_truth(query_idx) + query_vec = query_vectors[query_idx] + + # Compute expected ground truth + distances = [] + for doc_id, vec in vectors.items(): + dist = np.sum((query_vec - vec) ** 2) + distances.append((dist, doc_id)) + + distances.sort() + expected = [doc_id for _, doc_id in distances[:k]] + + if set(ground_truth) != set(expected): + print(f"❌ FAIL: Query {query_idx} ground truth mismatch") + all_pass = False + else: + print(f"✅ Query {query_idx}: Ground truth correct") + + if all_pass: + print("✅ PASS: All queries have correct ground truth") + return True + else: + return False + + +def test_cosine_distance(): + """Test GroundTruthTracker with cosine distance""" + print("\n=== Test 5: Cosine Distance Metric ===") + + dimension = 3 + k = 2 + + # Query vector + query_vector = np.array([1.0, 0.0, 0.0]) + + # Initialize tracker with cosine distance + tracker = GroundTruthTracker([query_vector], k=k, space_type='cosine') + + # Add vectors with known cosine similarities + # Same direction (cosine distance = 0) + tracker.update("id_1", [2.0, 0.0, 0.0]) + + # Orthogonal (cosine distance = 1) + tracker.update("id_2", [0.0, 1.0, 0.0]) + + # 45 degrees (cosine distance ≈ 0.293) + tracker.update("id_3", [1.0, 1.0, 0.0]) + + # Opposite direction (cosine distance = 2) + tracker.update("id_4", [-1.0, 0.0, 0.0]) + + ground_truth = tracker.get_ground_truth(0) + + print(f"Query vector: {query_vector}") + print(f"Ground truth (k={k}): {ground_truth}") + + # Expected: id_1 (same direction), id_3 (45 degrees) + expected = ["id_1", "id_3"] + + if set(ground_truth) == set(expected): + print("✅ PASS: Cosine distance ground truth correct") + return True + else: + print(f"❌ FAIL: Expected {expected}, got {ground_truth}") + return False + + +def test_edge_cases(): + """Test edge cases""" + print("\n=== Test 6: Edge Cases ===") + + all_pass = True + + # Test with k larger than number of vectors + print("\nTest 6a: k larger than number of vectors") + query_vector = np.array([0.0, 0.0]) + tracker = GroundTruthTracker([query_vector], k=10, space_type='l2') + + tracker.update("id_1", [1.0, 0.0]) + tracker.update("id_2", [0.0, 1.0]) + tracker.update("id_3", [1.0, 1.0]) + + ground_truth = tracker.get_ground_truth(0) + + if len(ground_truth) == 3: # Should return all 3 vectors + print("✅ PASS: Returns all vectors when k > num_vectors") + else: + print(f"❌ FAIL: Expected 3 vectors, got {len(ground_truth)}") + all_pass = False + + # Test with identical vectors + print("\nTest 6b: Identical vectors (tie-breaking)") + query_vector = np.array([0.0, 0.0]) + tracker = GroundTruthTracker([query_vector], k=2, space_type='l2') + + # All at same distance + tracker.update("id_1", [1.0, 0.0]) + tracker.update("id_2", [0.0, 1.0]) + tracker.update("id_3", [-1.0, 0.0]) + tracker.update("id_4", [0.0, -1.0]) + + ground_truth = tracker.get_ground_truth(0) + + if len(ground_truth) == 2: + print(f"✅ PASS: Returns k={2} vectors even with ties") + else: + print(f"❌ FAIL: Expected 2 vectors, got {len(ground_truth)}") + all_pass = False + + return all_pass + + +def run_all_tests(): + """Run all tests and report results""" + print("=" * 70) + print("RECALL MEASUREMENT TEST SUITE") + print("=" * 70) + + tests = [ + ("Basic Functionality", test_ground_truth_tracker_basic), + ("Large Vector Set", test_ground_truth_tracker_large), + ("Recall Calculation", test_calculate_recall), + ("Multiple Queries", test_multiple_queries), + ("Cosine Distance", test_cosine_distance), + ("Edge Cases", test_edge_cases), + ] + + results = [] + for test_name, test_func in tests: + try: + passed = test_func() + results.append((test_name, passed)) + except Exception as e: + print(f"\n❌ EXCEPTION in {test_name}: {e}") + import traceback + traceback.print_exc() + results.append((test_name, False)) + + # Print summary + print("\n" + "=" * 70) + print("TEST SUMMARY") + print("=" * 70) + + passed_count = sum(1 for _, passed in results if passed) + total_count = len(results) + + for test_name, passed in results: + status = "✅ PASS" if passed else "❌ FAIL" + print(f"{status}: {test_name}") + + print(f"\nTotal: {passed_count}/{total_count} tests passed") + + if passed_count == total_count: + print("\n🎉 All tests passed! Recall measurement is working correctly.") + return 0 + else: + print(f"\n⚠️ {total_count - passed_count} test(s) failed.") + return 1 + + +if __name__ == "__main__": + sys.exit(run_all_tests()) +