wehub-resource-sync
337 lines
12 KiB
Python
337 lines
12 KiB
Python
"""
|
|
Test DiskANN graph partitioning functionality.
|
|
|
|
Tests the automatic graph partitioning feature that was implemented to save
|
|
storage space by partitioning large DiskANN indices and safely deleting
|
|
redundant files while maintaining search functionality.
|
|
"""
|
|
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("CI") == "true",
|
|
reason="Skip DiskANN partition tests in CI - requires specific hardware and large memory",
|
|
)
|
|
def test_diskann_without_partition():
|
|
"""Test DiskANN index building without partition (baseline)."""
|
|
from leann.api import LeannBuilder, LeannSearcher
|
|
|
|
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
|
|
index_path = str(Path(temp_dir) / "test_no_partition.leann")
|
|
|
|
texts = [
|
|
f"Document {i} discusses topic {i % 10} with detailed analysis of subject {i // 10}."
|
|
for i in range(500)
|
|
]
|
|
|
|
builder = LeannBuilder(
|
|
backend_name="diskann",
|
|
embedding_model="facebook/contriever",
|
|
embedding_mode="sentence-transformers",
|
|
num_neighbors=32,
|
|
search_list_size=50,
|
|
is_recompute=False,
|
|
)
|
|
|
|
for text in texts:
|
|
builder.add_text(text)
|
|
|
|
builder.build_index(index_path)
|
|
|
|
index_dir = Path(index_path).parent
|
|
assert index_dir.exists()
|
|
|
|
index_prefix = Path(index_path).stem
|
|
required_files = [
|
|
f"{index_prefix}_disk.index",
|
|
f"{index_prefix}_pq_compressed.bin",
|
|
f"{index_prefix}_pq_pivots.bin",
|
|
]
|
|
|
|
generated_files = [f.name for f in index_dir.glob(f"{index_prefix}*")]
|
|
print(f"Generated files: {generated_files}")
|
|
|
|
for required_file in required_files:
|
|
file_path = index_dir / required_file
|
|
assert file_path.exists(), f"Required file {required_file} not found"
|
|
|
|
partition_files = [f"{index_prefix}_disk_graph.index", f"{index_prefix}_partition.bin"]
|
|
|
|
for partition_file in partition_files:
|
|
file_path = index_dir / partition_file
|
|
assert not file_path.exists(), (
|
|
f"Partition file {partition_file} should not exist in non-partition mode"
|
|
)
|
|
|
|
with LeannSearcher(index_path) as searcher:
|
|
results = searcher.search("topic 3 analysis", top_k=3)
|
|
|
|
assert len(results) > 0
|
|
assert all(
|
|
result.score is not None and result.score != float("-inf") for result in results
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("CI") == "true",
|
|
reason="Skip DiskANN partition tests in CI - requires specific hardware and large memory",
|
|
)
|
|
def test_diskann_with_partition():
|
|
"""Test DiskANN index building with automatic graph partitioning."""
|
|
from leann.api import LeannBuilder
|
|
|
|
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
|
|
index_path = str(Path(temp_dir) / "test_with_partition.leann")
|
|
|
|
texts = [
|
|
f"Document {i} explores subject {i % 15} with comprehensive coverage of area {i // 15}."
|
|
for i in range(500)
|
|
]
|
|
|
|
builder = LeannBuilder(
|
|
backend_name="diskann",
|
|
embedding_model="facebook/contriever",
|
|
embedding_mode="sentence-transformers",
|
|
num_neighbors=32,
|
|
search_list_size=50,
|
|
is_recompute=True,
|
|
)
|
|
|
|
for text in texts:
|
|
builder.add_text(text)
|
|
|
|
builder.build_index(index_path)
|
|
|
|
index_dir = Path(index_path).parent
|
|
assert index_dir.exists()
|
|
|
|
index_prefix = Path(index_path).stem
|
|
partition_files = [
|
|
f"{index_prefix}_disk_graph.index",
|
|
f"{index_prefix}_partition.bin",
|
|
f"{index_prefix}_pq_compressed.bin",
|
|
f"{index_prefix}_pq_pivots.bin",
|
|
]
|
|
|
|
for partition_file in partition_files:
|
|
file_path = index_dir / partition_file
|
|
assert file_path.exists(), f"Expected partition file {partition_file} not found"
|
|
|
|
large_files = [f"{index_prefix}_disk.index", f"{index_prefix}_disk_beam_search.index"]
|
|
|
|
for large_file in large_files:
|
|
file_path = index_dir / large_file
|
|
assert not file_path.exists(), (
|
|
f"Large file {large_file} should have been deleted for storage saving"
|
|
)
|
|
|
|
required_files = [
|
|
f"{index_prefix}_disk.index_medoids.bin",
|
|
f"{index_prefix}_disk.index_max_base_norm.bin",
|
|
]
|
|
|
|
for req_file in required_files:
|
|
file_path = index_dir / req_file
|
|
assert file_path.exists(), (
|
|
f"Required auxiliary file {req_file} missing for partition mode"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("CI") == "true",
|
|
reason="Skip DiskANN partition tests in CI - requires specific hardware and large memory",
|
|
)
|
|
def test_diskann_partition_search_functionality():
|
|
"""Test that search works correctly with partitioned indices."""
|
|
from leann.api import LeannBuilder, LeannSearcher
|
|
|
|
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
|
|
index_path = str(Path(temp_dir) / "test_partition_search.leann")
|
|
|
|
texts = [
|
|
"LEANN is a storage-efficient approximate nearest neighbor search system.",
|
|
"Graph partitioning helps reduce memory usage in large scale vector search.",
|
|
"DiskANN provides high-performance disk-based approximate nearest neighbor search.",
|
|
"Vector embeddings enable semantic search over unstructured text data.",
|
|
"Approximate nearest neighbor algorithms trade accuracy for speed and storage.",
|
|
] * 100
|
|
|
|
builder = LeannBuilder(
|
|
backend_name="diskann",
|
|
embedding_model="facebook/contriever",
|
|
embedding_mode="sentence-transformers",
|
|
is_recompute=True,
|
|
)
|
|
|
|
for text in texts:
|
|
builder.add_text(text)
|
|
|
|
builder.build_index(index_path)
|
|
|
|
with LeannSearcher(index_path) as searcher:
|
|
test_queries = [
|
|
("vector search algorithms", 5),
|
|
("LEANN storage efficiency", 3),
|
|
("graph partitioning memory", 4),
|
|
("approximate nearest neighbor", 7),
|
|
]
|
|
|
|
for query, top_k in test_queries:
|
|
results = searcher.search(query, top_k=top_k)
|
|
|
|
assert len(results) == top_k, f"Expected {top_k} results for query '{query}'"
|
|
assert all(result.score is not None for result in results), (
|
|
"All results should have scores"
|
|
)
|
|
assert all(result.score != float("-inf") for result in results), (
|
|
"No result should have -inf score"
|
|
)
|
|
assert all(result.text is not None for result in results), (
|
|
"All results should have text"
|
|
)
|
|
|
|
scores = [result.score for result in results]
|
|
assert scores == sorted(scores, reverse=True), (
|
|
"Results should be sorted by score descending"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("CI") == "true",
|
|
reason="Skip DiskANN partition tests in CI - requires specific hardware and large memory",
|
|
)
|
|
def test_diskann_medoid_and_norm_files():
|
|
"""Test that medoid and max_base_norm files are correctly generated and used."""
|
|
import struct
|
|
|
|
from leann.api import LeannBuilder, LeannSearcher
|
|
|
|
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
|
|
index_path = str(Path(temp_dir) / "test_medoid_norm.leann")
|
|
|
|
texts = [f"Test document {i} with content about subject {i % 10}." for i in range(200)]
|
|
|
|
builder = LeannBuilder(
|
|
backend_name="diskann",
|
|
embedding_model="facebook/contriever",
|
|
embedding_mode="sentence-transformers",
|
|
is_recompute=True,
|
|
)
|
|
|
|
for text in texts:
|
|
builder.add_text(text)
|
|
|
|
builder.build_index(index_path)
|
|
|
|
index_dir = Path(index_path).parent
|
|
index_prefix = Path(index_path).stem
|
|
|
|
medoids_file = index_dir / f"{index_prefix}_disk.index_medoids.bin"
|
|
assert medoids_file.exists(), "Medoids file should be generated"
|
|
|
|
with open(medoids_file, "rb") as f:
|
|
nshards = struct.unpack("<I", f.read(4))[0]
|
|
one_val = struct.unpack("<I", f.read(4))[0]
|
|
medoid_id = struct.unpack("<I", f.read(4))[0]
|
|
|
|
assert nshards == 1, "Single-shot build should have 1 shard"
|
|
assert one_val == 1, "Expected value should be 1"
|
|
assert medoid_id >= 0, "Medoid ID should be valid (not hardcoded 0)"
|
|
|
|
norm_file = index_dir / f"{index_prefix}_disk.index_max_base_norm.bin"
|
|
assert norm_file.exists(), "Max base norm file should be generated"
|
|
|
|
with open(norm_file, "rb") as f:
|
|
npts = struct.unpack("<I", f.read(4))[0]
|
|
ndims = struct.unpack("<I", f.read(4))[0]
|
|
norm_val = struct.unpack("<f", f.read(4))[0]
|
|
|
|
assert npts == 1, "Should have 1 norm point"
|
|
assert ndims == 1, "Should have 1 dimension"
|
|
assert norm_val > 0, "Norm value should be positive"
|
|
assert norm_val != float("inf"), "Norm value should be finite"
|
|
|
|
with LeannSearcher(index_path) as searcher:
|
|
results = searcher.search("test subject", top_k=3)
|
|
|
|
assert len(results) > 0
|
|
assert all(result.score != float("-inf") for result in results), (
|
|
"Scores should not be -inf when norm file is correct"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
os.environ.get("CI") == "true",
|
|
reason="Skip performance comparison in CI - requires significant compute time",
|
|
)
|
|
def test_diskann_vs_hnsw_performance():
|
|
"""Compare DiskANN (with partition) vs HNSW performance."""
|
|
import time
|
|
|
|
from leann.api import LeannBuilder, LeannSearcher
|
|
|
|
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
|
|
texts = [
|
|
f"Performance test document {i} covering topic {i % 20} in detail." for i in range(1000)
|
|
]
|
|
query = "performance topic test"
|
|
|
|
diskann_path = str(Path(temp_dir) / "perf_diskann.leann")
|
|
diskann_builder = LeannBuilder(
|
|
backend_name="diskann",
|
|
embedding_model="facebook/contriever",
|
|
embedding_mode="sentence-transformers",
|
|
is_recompute=True,
|
|
)
|
|
|
|
for text in texts:
|
|
diskann_builder.add_text(text)
|
|
|
|
start_time = time.time()
|
|
diskann_builder.build_index(diskann_path)
|
|
|
|
hnsw_path = str(Path(temp_dir) / "perf_hnsw.leann")
|
|
hnsw_builder = LeannBuilder(
|
|
backend_name="hnsw",
|
|
embedding_model="facebook/contriever",
|
|
embedding_mode="sentence-transformers",
|
|
is_recompute=True,
|
|
)
|
|
|
|
for text in texts:
|
|
hnsw_builder.add_text(text)
|
|
|
|
start_time = time.time()
|
|
hnsw_builder.build_index(hnsw_path)
|
|
|
|
with (
|
|
LeannSearcher(diskann_path) as diskann_searcher,
|
|
LeannSearcher(hnsw_path) as hnsw_searcher,
|
|
):
|
|
diskann_searcher.search(query, top_k=5)
|
|
hnsw_searcher.search(query, top_k=5)
|
|
|
|
start_time = time.time()
|
|
diskann_results = diskann_searcher.search(query, top_k=10)
|
|
diskann_search_time = time.time() - start_time
|
|
|
|
start_time = time.time()
|
|
hnsw_results = hnsw_searcher.search(query, top_k=10)
|
|
hnsw_search_time = time.time() - start_time
|
|
|
|
assert len(diskann_results) == 10
|
|
assert len(hnsw_results) == 10
|
|
assert all(r.score != float("-inf") for r in diskann_results)
|
|
assert all(r.score != float("-inf") for r in hnsw_results)
|
|
|
|
if hnsw_search_time > 0:
|
|
speed_ratio = hnsw_search_time / diskann_search_time
|
|
print(f"DiskANN search time: {diskann_search_time:.4f}s")
|
|
print(f"HNSW search time: {hnsw_search_time:.4f}s")
|
|
print(f"DiskANN is {speed_ratio:.2f}x faster than HNSW")
|