53827065-0ae8-4037-abb8-c5bca28f1b89

Committed 19 Sep 2023 01:35AM UTC coverage: 74.296% (-0.08%) from 74.375%

Build # 53827065-0ae8-4037-abb8-c5bca28f1b89

Build Type

push

circle-ci

Committed by

jiashenC

Commit Message

add index scan for native pgvector

Run Details

40 of 40 new or added lines in 5 files covered. (100.0%)

8897 of 11975 relevant lines covered (74.3%)

0.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

36.36

/evadb/executor/vector_index_scan_executor.py

# coding=utf-8
# Copyright 2018-2023 EvaDB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Iterator

import pandas as pd

from evadb.catalog.sql_config import ROW_NUM_COLUMN
from evadb.database import EvaDBDatabase
from evadb.executor.abstract_executor import AbstractExecutor
from evadb.executor.executor_utils import handle_vector_store_params
from evadb.models.storage.batch import Batch
from evadb.plan_nodes.vector_index_scan_plan import VectorIndexScanPlan
from evadb.plan_nodes.storage_plan import StoragePlan
from evadb.third_party.vector_stores.types import VectorIndexQuery
from evadb.third_party.vector_stores.utils import VectorStoreFactory
from evadb.third_party.databases.interface import get_database_handler
from evadb.executor.executor_utils import ExecutorError
from evadb.utils.logging_manager import logger
from evadb.catalog.models.utils import VectorStoreType


# Helper function for getting row_num column alias.
def get_row_num_column_alias(column_list):
    for column in column_list:
        alias, col_name = column.split(".")
        if col_name == ROW_NUM_COLUMN:
            return alias


class VectorIndexScanExecutor(AbstractExecutor):
    def __init__(self, db: EvaDBDatabase, node: VectorIndexScanPlan):
        super().__init__(db, node)

        self.index_name = node.index.name
        self.vector_store_type = node.index.type
        self.feat_column = node.index.feat_column
        self.limit_count = node.limit_count
        self.search_query_expr = node.search_query_expr

    def exec(self, *args, **kwargs) -> Iterator[Batch]:
        if self.vector_store_type == VectorStoreType.PGVECTOR:
            return self._native_vector_index_scan()
        else:
            return self._evadb_vector_index_scan(*args, **kwargs)


    def _get_search_query_results(self):
        # Get the query feature vector. Create a dummy
        # batch to retreat a single file path.
        dummy_batch = Batch(
            frames=pd.DataFrame(
                {"0": [0]},
            )
        )
        search_batch = self.search_query_expr.evaluate(dummy_batch)

        # Scan index. The search batch comes from the Open call.
        feature_col_name = self.search_query_expr.output_objs[0].name
        search_batch.drop_column_alias()
        search_feat = search_batch.column_as_numpy_array(feature_col_name)[0]
        search_feat = search_feat.reshape(1, -1)
        return search_feat


    def _native_vector_index_scan(self):
        search_feat = self._get_search_query_results()
        search_feat = search_feat.reshape(-1).tolist()

        tb_catalog_entry = list(self.node.find_all(StoragePlan))[0].table
        db_catalog_entry = self.db.catalog().get_database_catalog_entry(tb_catalog_entry.database_name)
        with get_database_handler(db_catalog_entry.engine, **db_catalog_entry.params) as handler:
            resp = handler.execute_native_query(f"""SELECT * FROM {tb_catalog_entry.name} 
                                                ORDER BY {self.feat_column.name} <-> '{search_feat}' 
                                                LIMIT {self.limit_count}""")
            if resp.error is not None:
                raise ExecutorError(f"Native index can encounters {resp.error}")
            res = Batch(frames=resp.data)
            res.modify_column_alias(tb_catalog_entry.name)
            yield res

    def _evadb_vector_index_scan(self, *args, **kwargs):
        # Fetch the index from disk.
        index_catalog_entry = self.catalog().get_index_catalog_entry_by_name(
            self.index_name
        )
        self.index_path = index_catalog_entry.save_file_path
        self.index = VectorStoreFactory.init_vector_store(
            self.vector_store_type,
            self.index_name,
            **handle_vector_store_params(self.vector_store_type, self.index_path),
        )

        search_feat = self._get_search_query_results()
        index_result = self.index.query(
            VectorIndexQuery(search_feat, self.limit_count.value)
        )
        # todo support queries over distance as well
        # distance_list = index_result.similarities
        row_num_np = index_result.ids

        # Load projected columns from disk and join with search results.
        row_num_col_name = None

        # handle the case where the index_results are less than self.limit_count.value
        num_required_results = self.limit_count.value
        if len(index_result.ids) < self.limit_count.value:
            num_required_results = len(index_result.ids)
            logger.warning(
                f"The index {self.index_name} returned only {num_required_results} results, which is fewer than the required {self.limit_count.value}."
            )

        res_row_list = [None for _ in range(num_required_results)]
        for batch in self.children[0].exec(**kwargs):
            column_list = batch.columns
            if not row_num_col_name:
                row_num_alias = get_row_num_column_alias(column_list)
                row_num_col_name = "{}.{}".format(row_num_alias, ROW_NUM_COLUMN)

            # Nested join.
            for _, row in batch.frames.iterrows():
                for idx, row_num in enumerate(row_num_np):
                    if row_num == row[row_num_col_name]:
                        res_row = dict()
                        for col_name in column_list:
                            res_row[col_name] = row[col_name]
                        res_row_list[idx] = res_row

        yield Batch(pd.DataFrame(res_row_list))

1	# coding=utf-8
2	# Copyright 2018-2023 EvaDB
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	# http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15	from typing import Iterator	1✔
16
17	import pandas as pd	1✔
18
19	from evadb.catalog.sql_config import ROW_NUM_COLUMN	1✔
20	from evadb.database import EvaDBDatabase	1✔
21	from evadb.executor.abstract_executor import AbstractExecutor	1✔
22	from evadb.executor.executor_utils import handle_vector_store_params	1✔
23	from evadb.models.storage.batch import Batch	1✔
24	from evadb.plan_nodes.vector_index_scan_plan import VectorIndexScanPlan	1✔
25	from evadb.plan_nodes.storage_plan import StoragePlan	1✔
26	from evadb.third_party.vector_stores.types import VectorIndexQuery	1✔
27	from evadb.third_party.vector_stores.utils import VectorStoreFactory	1✔
28	from evadb.third_party.databases.interface import get_database_handler	1✔
29	from evadb.executor.executor_utils import ExecutorError	1✔
30	from evadb.utils.logging_manager import logger	1✔
31	from evadb.catalog.models.utils import VectorStoreType	1✔
32
33
34	# Helper function for getting row_num column alias.
35	def get_row_num_column_alias(column_list):	1✔
36	for column in column_list:	×
37	alias, col_name = column.split(".")	×
38	if col_name == ROW_NUM_COLUMN:	×
39	return alias	×
40
41
42	class VectorIndexScanExecutor(AbstractExecutor):	1✔
43	def __init__(self, db: EvaDBDatabase, node: VectorIndexScanPlan):	1✔
44	super().__init__(db, node)	1✔
45
46	self.index_name = node.index.name	1✔
47	self.vector_store_type = node.index.type	1✔
48	self.feat_column = node.index.feat_column	1✔
49	self.limit_count = node.limit_count	1✔
50	self.search_query_expr = node.search_query_expr	1✔
51
52	def exec(self, args, *kwargs) -> Iterator[Batch]:	1✔
53	if self.vector_store_type == VectorStoreType.PGVECTOR:	×
54	return self._native_vector_index_scan()	×
55	else:
56	return self._evadb_vector_index_scan(args, *kwargs)	×
57
58
59	def _get_search_query_results(self):	1✔
60	# Get the query feature vector. Create a dummy
61	# batch to retreat a single file path.
62	dummy_batch = Batch(	×
63	frames=pd.DataFrame(
64	{"0": [0]},
65	)
66	)
67	search_batch = self.search_query_expr.evaluate(dummy_batch)	×
68
69	# Scan index. The search batch comes from the Open call.
70	feature_col_name = self.search_query_expr.output_objs[0].name	×
71	search_batch.drop_column_alias()	×
72	search_feat = search_batch.column_as_numpy_array(feature_col_name)[0]	×
73	search_feat = search_feat.reshape(1, -1)	×
74	return search_feat	×
75
76
77	def _native_vector_index_scan(self):	1✔
78	search_feat = self._get_search_query_results()	×
79	search_feat = search_feat.reshape(-1).tolist()	×
80
81	tb_catalog_entry = list(self.node.find_all(StoragePlan))[0].table	×
82	db_catalog_entry = self.db.catalog().get_database_catalog_entry(tb_catalog_entry.database_name)	×
83	with get_database_handler(db_catalog_entry.engine, **db_catalog_entry.params) as handler:	×
84	resp = handler.execute_native_query(f"""SELECT * FROM {tb_catalog_entry.name}	×
85	ORDER BY {self.feat_column.name} <-> '{search_feat}'
86	LIMIT {self.limit_count}""")
87	if resp.error is not None:	×
88	raise ExecutorError(f"Native index can encounters {resp.error}")
89	res = Batch(frames=resp.data)	×
90	res.modify_column_alias(tb_catalog_entry.name)	×
91	yield res	×
92
93	def _evadb_vector_index_scan(self, args, *kwargs):	1✔
94	# Fetch the index from disk.
95	index_catalog_entry = self.catalog().get_index_catalog_entry_by_name(	×
96	self.index_name
97	)
98	self.index_path = index_catalog_entry.save_file_path	×
99	self.index = VectorStoreFactory.init_vector_store(	×
100	self.vector_store_type,
101	self.index_name,
102	**handle_vector_store_params(self.vector_store_type, self.index_path),
103	)
104
105	search_feat = self._get_search_query_results()	×
106	index_result = self.index.query(	×
107	VectorIndexQuery(search_feat, self.limit_count.value)
108	)
109	# todo support queries over distance as well
110	# distance_list = index_result.similarities
111	row_num_np = index_result.ids	×
112
113	# Load projected columns from disk and join with search results.
114	row_num_col_name = None	×
115
116	# handle the case where the index_results are less than self.limit_count.value
117	num_required_results = self.limit_count.value	×
118	if len(index_result.ids) < self.limit_count.value:	×
119	num_required_results = len(index_result.ids)	×
120	logger.warning(	×
121	f"The index {self.index_name} returned only {num_required_results} results, which is fewer than the required {self.limit_count.value}."
122	)
123
124	res_row_list = [None for _ in range(num_required_results)]	×
125	for batch in self.children[0].exec(**kwargs):	×
126	column_list = batch.columns	×
127	if not row_num_col_name:	×
128	row_num_alias = get_row_num_column_alias(column_list)	×
129	row_num_col_name = "{}.{}".format(row_num_alias, ROW_NUM_COLUMN)	×
130
131	# Nested join.
132	for _, row in batch.frames.iterrows():	×
133	for idx, row_num in enumerate(row_num_np):	×
134	if row_num == row[row_num_col_name]:	×
135	res_row = dict()	×
136	for col_name in column_list:	×
137	res_row[col_name] = row[col_name]	×
138	res_row_list[idx] = res_row	×
139
140	yield Batch(pd.DataFrame(res_row_list))	×

georgia-tech-db / eva / 53827065-0ae8-4037-abb8-c5bca28f1b89

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous