Skip to content

Commit 2a8ecc8

Browse files
pierrecouchpierrecouch
committed
Add sparse and dense type for explain
Change-Id: I89a7870426c19ea62f318affd84c87a390e96faa Reviewed-on: https://review.couchbase.org/c/testrunner/+/237736 Reviewed-by: Ajay Bhullar <[email protected]> Tested-by: Pierre Regazzoni <[email protected]>
1 parent 20022dc commit 2a8ecc8

File tree

1 file changed

+26
-18
lines changed

1 file changed

+26
-18
lines changed

‎pytests/tuqquery/tuq_vectorsearch.py‎

Lines changed: 26 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -535,12 +535,16 @@ def test_advise_ann(self):
535535
similarity = self.distance.lower()
536536
if similarity in ['l2', 'euclidean']:
537537
similarity += '_squared'
538-
expected_index1 = f"CREATE INDEX adv_brand_size_vecVECTOR_id ON `default`(`brand`,`size`,`vec` VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
539-
expected_index2 = f"CREATE INDEX adv_size_brand_vecVECTOR_id ON `default`(`size`,`brand`,`vec` VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
540-
expected_bhive_index1 = f"CREATE VECTOR INDEX adv_VECTOR_vecVECTOR_INCLUDE_brand_size_id ON `default`(`vec` VECTOR) INCLUDE (`brand`,`size`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
541-
expected_bhive_index2 = f"CREATE VECTOR INDEX adv_VECTOR_vecVECTOR_INCLUDE_size_brand_id ON `default`(`vec` VECTOR) INCLUDE (`size`,`brand`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
538+
vector_type = self.vector_type.upper()
539+
expected_index1 = f"CREATE INDEX adv_brand_size_vec{vector_type}VECTOR_id ON `default`(`brand`,`size`,`vec` {vector_type} VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
540+
expected_index2 = f"CREATE INDEX adv_size_brand_vec{vector_type}VECTOR_id ON `default`(`size`,`brand`,`vec` {vector_type} VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
541+
expected_bhive_index1 = f"CREATE VECTOR INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_brand_size_id ON `default`(`vec` {vector_type} VECTOR) INCLUDE (`brand`,`size`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
542+
expected_bhive_index2 = f"CREATE VECTOR INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_size_brand_id ON `default`(`vec` {vector_type} VECTOR) INCLUDE (`size`,`brand`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
542543
expected_property = "ORDER pushdown, LIMIT pushdown"
543-
advise_ann_query = f'ADVISE SELECT id, size, brand FROM default WHERE size = 6 AND brand = "Puma" ORDER BY ANN_DISTANCE(vec, {self.xq[1].tolist()}, "{self.distance}") LIMIT 100'
544+
if self.vector_type == 'sparse':
545+
advise_ann_query = f'ADVISE SELECT id, size, brand FROM default WHERE size = 6 AND brand = "Puma" ORDER BY SPARSE_VECTOR_DISTANCE(vec, {self.xq[1].tolist()}, {self.nprobes}) LIMIT 100'
546+
else:
547+
advise_ann_query = f'ADVISE SELECT id, size, brand FROM default WHERE size = 6 AND brand = "Puma" ORDER BY ANN_DISTANCE(vec, {self.xq[1].tolist()}, "{self.distance}", {self.nprobes}) LIMIT 100'
544548
advise = self.run_cbq_query(advise_ann_query)
545549
self.log.info(advise['results'])
546550
adviseinfo = advise['results'][0]['advice']['adviseinfo']
@@ -557,20 +561,24 @@ def test_advise_ann(self):
557561
self.run_cbq_query(index_statement)
558562
self.run_cbq_query(bhive_index_statement)
559563
finally:
560-
self.run_cbq_query(f'DROP INDEX adv_brand_size_vecVECTOR_id IF EXISTS on default')
561-
self.run_cbq_query(f'DROP INDEX adv_size_brand_vecVECTOR_id IF EXISTS on default')
562-
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vecVECTOR_INCLUDE_brand_size_id IF EXISTS on default')
563-
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vecVECTOR_INCLUDE_size_brand_id IF EXISTS on default')
564+
self.run_cbq_query(f'DROP INDEX adv_brand_size_vec{vector_type}VECTOR_id IF EXISTS on default')
565+
self.run_cbq_query(f'DROP INDEX adv_size_brand_vec{vector_type}VECTOR_id IF EXISTS on default')
566+
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_brand_size_id IF EXISTS on default')
567+
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_size_brand_id IF EXISTS on default')
564568

565569
def test_advise_ann_no_limit(self):
566570
similarity = self.distance.lower()
567571
if similarity in ['l2', 'euclidean']:
568572
similarity += '_squared'
569-
expected_index1 = f"CREATE INDEX adv_brand_size_vecVECTOR_id ON `default`(`brand`,`size`,`vec` VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
570-
expected_index2 = f"CREATE INDEX adv_size_brand_vecVECTOR_id ON `default`(`size`,`brand`,`vec` VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
571-
expected_bhive_index1 = f"CREATE VECTOR INDEX adv_VECTOR_vecVECTOR_INCLUDE_brand_size_id ON `default`(`vec` VECTOR) INCLUDE (`brand`,`size`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
572-
expected_bhive_index2 = f"CREATE VECTOR INDEX adv_VECTOR_vecVECTOR_INCLUDE_size_brand_id ON `default`(`vec` VECTOR) INCLUDE (`size`,`brand`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
573-
advise_ann_query = f'ADVISE SELECT id, size, brand FROM default WHERE size = 6 AND brand = "Puma" ORDER BY ANN_DISTANCE(vec, {self.xq[1].tolist()}, "{self.distance}")'
573+
vector_type = self.vector_type.upper()
574+
expected_index1 = f"CREATE INDEX adv_brand_size_vec{vector_type}VECTOR_id ON `default`(`brand`,`size`,`vec` {vector_type} VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
575+
expected_index2 = f"CREATE INDEX adv_size_brand_vec{vector_type}VECTOR_id ON `default`(`size`,`brand`,`vec` {vector_type} VECTOR,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
576+
expected_bhive_index1 = f"CREATE VECTOR INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_brand_size_id ON `default`(`vec` {vector_type} VECTOR) INCLUDE (`brand`,`size`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
577+
expected_bhive_index2 = f"CREATE VECTOR INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_size_brand_id ON `default`(`vec` {vector_type} VECTOR) INCLUDE (`size`,`brand`,`id`) WITH {{ 'dimension': 128, 'similarity': '{similarity}', 'description': 'IVF,SQ8' }}"
578+
if self.vector_type == 'sparse':
579+
advise_ann_query = f'ADVISE SELECT id, size, brand FROM default WHERE size = 6 AND brand = "Puma" ORDER BY SPARSE_VECTOR_DISTANCE(vec, {self.xq[1].tolist()}, {self.nprobes})'
580+
else:
581+
advise_ann_query = f'ADVISE SELECT id, size, brand FROM default WHERE size = 6 AND brand = "Puma" ORDER BY ANN_DISTANCE(vec, {self.xq[1].tolist()}, "{self.distance}", {self.nprobes})'
574582
advise = self.run_cbq_query(advise_ann_query)
575583
self.log.info(advise['results'])
576584
adviseinfo = advise['results'][0]['advice']['adviseinfo']
@@ -583,10 +591,10 @@ def test_advise_ann_no_limit(self):
583591
self.run_cbq_query(index_statement)
584592
self.run_cbq_query(bhive_index_statement)
585593
finally:
586-
self.run_cbq_query(f'DROP INDEX adv_brand_size_vecVECTOR_id IF EXISTS on default')
587-
self.run_cbq_query(f'DROP INDEX adv_size_brand_vecVECTOR_id IF EXISTS on default')
588-
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vecVECTOR_INCLUDE_brand_size_id IF EXISTS on default')
589-
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vecVECTOR_INCLUDE_size_brand_id IF EXISTS on default')
594+
self.run_cbq_query(f'DROP INDEX adv_brand_size_vec{vector_type}VECTOR_id IF EXISTS on default')
595+
self.run_cbq_query(f'DROP INDEX adv_size_brand_vec{vector_type}VECTOR_id IF EXISTS on default')
596+
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_brand_size_id IF EXISTS on default')
597+
self.run_cbq_query(f'DROP INDEX adv_VECTOR_vec{vector_type}VECTOR_INCLUDE_size_brand_id IF EXISTS on default')
590598

591599
def test_advise_knn(self):
592600
expected_index1 = f"CREATE INDEX adv_brand_size ON `default`(`brand`,`size`)"

0 commit comments

Comments
 (0)