Skip to content

Commit 74f5930

Browse files
andrewmathew1Andrew Mathew
andrewmathew1
and
Andrew Mathew
authored
added diskANN testing back to test_query_vector_similarity (Azure#39654)
Co-authored-by: Andrew Mathew <[email protected]>
1 parent 96b2d6e commit 74f5930

File tree

2 files changed

+38
-37
lines changed

2 files changed

+38
-37
lines changed

‎sdk/cosmos/azure-cosmos/test/test_query_vector_similarity.py‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,14 @@ def setUpClass(cls):
6363
vector_embedding_policy=test_config.get_vector_embedding_policy(data_type="float32",
6464
distance_function="euclidean",
6565
dimensions=128))
66-
# cls.created_diskANN_dotproduct_container = cls.test_db.create_container(
67-
# id="diskANN" + cls.TEST_CONTAINER_ID,
68-
# partition_key=PartitionKey(path="/pk"),
69-
# offer_throughput=test_config.TestConfig.THROUGHPUT_FOR_5_PARTITIONS,
70-
# indexing_policy=test_config.get_vector_indexing_policy(embedding_type="diskANN"),
71-
# vector_embedding_policy=test_config.get_vector_embedding_policy(data_type="float32",
72-
# distance_function="dotproduct",
73-
# dimensions=128))
66+
cls.created_diskANN_dotproduct_container = cls.test_db.create_container(
67+
id="diskANN" + cls.TEST_CONTAINER_ID,
68+
partition_key=PartitionKey(path="/pk"),
69+
offer_throughput=test_config.TestConfig.THROUGHPUT_FOR_5_PARTITIONS,
70+
indexing_policy=test_config.get_vector_indexing_policy(embedding_type="diskANN"),
71+
vector_embedding_policy=test_config.get_vector_embedding_policy(data_type="float32",
72+
distance_function="dotproduct",
73+
dimensions=128))
7474
cls.created_large_container = cls.test_db.create_container(
7575
id="large_container" + cls.TEST_CONTAINER_ID,
7676
partition_key=PartitionKey(path="/pk"),
@@ -82,14 +82,14 @@ def setUpClass(cls):
8282
for item in vector_test_data.get_vector_items():
8383
cls.created_quantized_cosine_container.create_item(item)
8484
cls.created_flat_euclidean_container.create_item(item)
85-
# cls.created_diskANN_dotproduct_container.create_item(item)
85+
cls.created_diskANN_dotproduct_container.create_item(item)
8686

8787
@classmethod
8888
def tearDownClass(cls):
8989
try:
9090
cls.test_db.delete_container("quantized" + cls.TEST_CONTAINER_ID)
9191
cls.test_db.delete_container("flat" + cls.TEST_CONTAINER_ID)
92-
# cls.test_db.delete_container("diskANN" + cls.TEST_CONTAINER_ID)
92+
cls.test_db.delete_container("diskANN" + cls.TEST_CONTAINER_ID)
9393
cls.test_db.delete_container("large_container" + cls.TEST_CONTAINER_ID)
9494
cls.client.delete_database(cls.test_db.id)
9595
except exceptions.CosmosHttpResponseError:
@@ -145,10 +145,10 @@ def test_ordering_distances(self):
145145
enable_cross_partition_query=True))
146146
verify_ordering(quantized_list, "euclidean")
147147

148-
# disk_ann_list = list(
149-
# self.created_diskANN_dotproduct_container.query_items(query=specs_query,
150-
# enable_cross_partition_query=True))
151-
# verify_ordering(disk_ann_list, "euclidean")
148+
disk_ann_list = list(
149+
self.created_diskANN_dotproduct_container.query_items(query=specs_query,
150+
enable_cross_partition_query=True))
151+
verify_ordering(disk_ann_list, "euclidean")
152152
# test cosine distance
153153
for i in range(1, 11):
154154
vanilla_query = "SELECT TOP {} c.text, VectorDistance(c.embedding, [{}]) AS " \
@@ -168,10 +168,10 @@ def test_ordering_distances(self):
168168
enable_cross_partition_query=True))
169169
verify_ordering(quantized_list, "cosine")
170170

171-
# disk_ann_list = list(
172-
# self.created_diskANN_dotproduct_container.query_items(query=specs_query,
173-
# enable_cross_partition_query=True))
174-
# verify_ordering(disk_ann_list, "cosine")
171+
disk_ann_list = list(
172+
self.created_diskANN_dotproduct_container.query_items(query=specs_query,
173+
enable_cross_partition_query=True))
174+
verify_ordering(disk_ann_list, "cosine")
175175
# test dot product distance
176176
for i in range(1, 11):
177177
vanilla_query = "SELECT TOP {} c.text, VectorDistance(c.embedding, [{}]) AS " \
@@ -191,10 +191,10 @@ def test_ordering_distances(self):
191191
enable_cross_partition_query=True))
192192
verify_ordering(quantized_list, "dotproduct")
193193

194-
# disk_ann_list = list(
195-
# self.created_diskANN_dotproduct_container.query_items(query=vanilla_query,
196-
# enable_cross_partition_query=True))
197-
# verify_ordering(disk_ann_list, "dotproduct")
194+
disk_ann_list = list(
195+
self.created_diskANN_dotproduct_container.query_items(query=vanilla_query,
196+
enable_cross_partition_query=True))
197+
verify_ordering(disk_ann_list, "dotproduct")
198198

199199
def test_vector_query_pagination(self):
200200
# load up previously calculated embedding for the given string

‎sdk/cosmos/azure-cosmos/test/test_query_vector_similarity_async.py‎

Lines changed: 16 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -63,14 +63,14 @@ def setUpClass(cls):
6363
vector_embedding_policy=test_config.get_vector_embedding_policy(data_type="float32",
6464
distance_function="euclidean",
6565
dimensions=128))
66-
# cls.created_diskANN_dotproduct_container = cls.test_db.create_container(
67-
# id="diskANN" + cls.TEST_CONTAINER_ID,
68-
# partition_key=PartitionKey(path="/pk"),
69-
# offer_throughput=test_config.TestConfig.THROUGHPUT_FOR_5_PARTITIONS,
70-
# indexing_policy=test_config.get_vector_indexing_policy(embedding_type="diskANN"),
71-
# vector_embedding_policy=test_config.get_vector_embedding_policy(data_type="float32",
72-
# distance_function="dotproduct",
73-
# dimensions=128))
66+
cls.created_diskANN_dotproduct_container = cls.test_db.create_container(
67+
id="diskANN" + cls.TEST_CONTAINER_ID,
68+
partition_key=PartitionKey(path="/pk"),
69+
offer_throughput=test_config.TestConfig.THROUGHPUT_FOR_5_PARTITIONS,
70+
indexing_policy=test_config.get_vector_indexing_policy(embedding_type="diskANN"),
71+
vector_embedding_policy=test_config.get_vector_embedding_policy(data_type="float32",
72+
distance_function="dotproduct",
73+
dimensions=128))
7474
cls.created_large_container = cls.test_db.create_container(
7575
id="large_container" + cls.TEST_CONTAINER_ID,
7676
partition_key=PartitionKey(path="/pk"),
@@ -82,7 +82,7 @@ def setUpClass(cls):
8282
for item in vector_test_data.get_vector_items():
8383
cls.created_quantized_cosine_container.create_item(item)
8484
cls.created_flat_euclidean_container.create_item(item)
85-
# await cls.created_diskANN_dotproduct_container.create_item(item)
85+
cls.created_diskANN_dotproduct_container.create_item(item)
8686

8787
@classmethod
8888
def tearDownClass(cls):
@@ -96,6 +96,7 @@ async def asyncSetUp(self):
9696
self.test_db = self.client.get_database_client(self.test_db.id)
9797
self.created_flat_euclidean_container = self.test_db.get_container_client(self.created_flat_euclidean_container.id)
9898
self.created_quantized_cosine_container = self.test_db.get_container_client(self.created_quantized_cosine_container.id)
99+
self.created_diskANN_dotproduct_container = self.test_db.get_container_client(self.created_diskANN_dotproduct_container.id)
99100
self.created_large_container = self.test_db.get_container_client(self.created_large_container.id)
100101

101102
async def asyncTearDown(self):
@@ -148,8 +149,8 @@ async def test_ordering_distances_async(self):
148149
quantized_list = [item async for item in self.created_quantized_cosine_container.query_items(query=specs_query)]
149150
verify_ordering(quantized_list, "euclidean")
150151

151-
# disk_ann_list = [item async for item in self.created_diskANN_dotproduct_container.query_items(query=specs_query)]
152-
# verify_ordering(disk_ann_list, "euclidean")
152+
disk_ann_list = [item async for item in self.created_diskANN_dotproduct_container.query_items(query=specs_query)]
153+
verify_ordering(disk_ann_list, "euclidean")
153154
# test cosine distance
154155
for i in range(1, 11):
155156
vanilla_query = "SELECT TOP {} c.text, VectorDistance(c.embedding, [{}]) AS " \
@@ -166,8 +167,8 @@ async def test_ordering_distances_async(self):
166167
quantized_list = [item async for item in self.created_quantized_cosine_container.query_items(query=vanilla_query)]
167168
verify_ordering(quantized_list, "cosine")
168169

169-
# disk_ann_list = [item async for item in self.created_diskANN_dotproduct_container.query_items(query=specs_query)]
170-
# verify_ordering(disk_ann_list, "cosine")
170+
disk_ann_list = [item async for item in self.created_diskANN_dotproduct_container.query_items(query=specs_query)]
171+
verify_ordering(disk_ann_list, "cosine")
171172
# test dot product distance
172173
for i in range(1, 11):
173174
vanilla_query = "SELECT TOP {} c.text, VectorDistance(c.embedding, [{}]) AS " \
@@ -184,8 +185,8 @@ async def test_ordering_distances_async(self):
184185
quantized_list = [item async for item in self.created_quantized_cosine_container.query_items(query=specs_query)]
185186
verify_ordering(quantized_list, "dotproduct")
186187

187-
# disk_ann_list = [item async for item in self.created_diskANN_dotproduct_container.query_items(query=vanilla_query)]
188-
# verify_ordering(disk_ann_list, "dotproduct")
188+
disk_ann_list = [item async for item in self.created_diskANN_dotproduct_container.query_items(query=vanilla_query)]
189+
verify_ordering(disk_ann_list, "dotproduct")
189190

190191
async def test_vector_query_pagination_async(self):
191192
# load up previously calculated embedding for the given string

0 commit comments

Comments
 (0)