#832

Committed 12 Oct 2023 08:31AM UTC coverage: 0.0% (-68.5%) from 68.493%

Build # #832

Build Type

push

circle-ci

Committed by

xzdandy

Commit Message

Minor

Run Details

0 of 12331 relevant lines covered (0.0%)

0.0 hits per line

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

0.0

/evadb/executor/create_index_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 pathlib import Path

import pandas as pd

from evadb.catalog.catalog_type import VectorStoreType
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 ExecutorError, handle_vector_store_params
from evadb.expression.function_expression import FunctionExpression
from evadb.models.storage.batch import Batch
from evadb.plan_nodes.create_index_plan import CreateIndexPlan
from evadb.third_party.databases.interface import get_database_handler
from evadb.third_party.vector_stores.types import FeaturePayload
from evadb.third_party.vector_stores.utils import VectorStoreFactory
from evadb.utils.logging_manager import logger


class CreateIndexExecutor(AbstractExecutor):
    def __init__(self, db: EvaDBDatabase, node: CreateIndexPlan):
        super().__init__(db, node)

        self.name = self.node.name
        self.if_not_exists = self.node.if_not_exists
        self.table_ref = self.node.table_ref
        self.col_list = self.node.col_list
        self.vector_store_type = self.node.vector_store_type
        self.project_expr_list = self.node.project_expr_list
        self.index_def = self.node.index_def

    def exec(self, *args, **kwargs):
        # Vector type specific creation.
        if self.vector_store_type == VectorStoreType.PGVECTOR:
            self._create_native_index()
        else:
            self._create_evadb_index()

        yield Batch(
            pd.DataFrame([f"Index {self.name} successfully added to the database."])
        )

    # Create index through the native storage engine.
    def _create_native_index(self):
        table = self.table_ref.table
        db_catalog_entry = self.catalog().get_database_catalog_entry(
            table.database_name
        )

        with get_database_handler(
            db_catalog_entry.engine, **db_catalog_entry.params
        ) as handler:
            # As other libraries, we default to HNSW and L2 distance.
            resp = handler.execute_native_query(
                f"""CREATE INDEX {self.name} ON {table.table_name}
                    USING hnsw ({self.col_list[0].name} vector_l2_ops)"""
            )
            if resp.error is not None:
                raise ExecutorError(
                    f"Native engine create index encounters error: {resp.error}"
                )

    # On-disk saving path for EvaDB index.
    def _get_evadb_index_save_path(self) -> Path:
        index_dir = Path(self.config.get_value("storage", "index_dir"))
        if not index_dir.exists():
            index_dir.mkdir(parents=True, exist_ok=True)
        return str(
            index_dir / Path("{}_{}.index".format(self.vector_store_type, self.name))
        )

    # Create EvaDB index.
    def _create_evadb_index(self):
        # Find function expression.
        function_expression, function_expression_signature = None, None
        for project_expr in self.project_expr_list:
            if isinstance(project_expr, FunctionExpression):
                function_expression = project_expr
                function_expression_signature = project_expr.signature()

        # Get feature tables.
        feat_tb_catalog_entry = self.table_ref.table.table_obj

        # Get feature column.
        feat_col_name = self.col_list[0].name
        feat_col_catalog_entry = [
            col for col in feat_tb_catalog_entry.columns if col.name == feat_col_name
        ][0]

        if function_expression is not None:
            feat_col_name = function_expression.output_objs[0].name

        index_catalog_entry = self.catalog().get_index_catalog_entry_by_name(self.name)
        index_path = self._get_evadb_index_save_path()

        if index_catalog_entry is not None:
            msg = f"Index {self.name} already exists."
            if self.if_not_exists:
                if (
                    index_catalog_entry.feat_column == feat_col_catalog_entry
                    and index_catalog_entry.function_signature
                    == function_expression_signature
                    and index_catalog_entry.type == self.node.vector_store_type
                ):
                    # Only update index if everything matches.
                    logger.warn(msg + " It will be updated on existing table.")
                    index = VectorStoreFactory.init_vector_store(
                        self.vector_store_type,
                        self.name,
                        **handle_vector_store_params(
                            self.vector_store_type, index_path
                        ),
                    )
                else:
                    # Skip index update if CREATE INDEX runs on a different index.
                    logger.warn(msg)
                    return
            else:
                logger.error(msg)
                raise ExecutorError(msg)
        else:
            index = None

        try:
            # Add features to index.
            for input_batch in self.children[0].exec():
                input_batch.drop_column_alias()
                feat = input_batch.column_as_numpy_array(feat_col_name)
                row_num = input_batch.column_as_numpy_array(ROW_NUM_COLUMN)

                for i in range(len(input_batch)):
                    row_feat = feat[i].reshape(1, -1)

                    # Create new index if not exists.
                    if index is None:
                        input_dim = row_feat.shape[1]
                        index = VectorStoreFactory.init_vector_store(
                            self.vector_store_type,
                            self.name,
                            **handle_vector_store_params(
                                self.vector_store_type, index_path
                            ),
                        )
                        index.create(input_dim)

                    # Row ID for mapping back to the row.
                    index.add([FeaturePayload(row_num[i], row_feat)])

            # Persist index.
            index.persist()

            # Save to catalog.
            if index_catalog_entry is None:
                self.catalog().insert_index_catalog_entry(
                    self.name,
                    index_path,
                    self.vector_store_type,
                    feat_col_catalog_entry,
                    function_expression_signature,
                    self.index_def,
                )
        except Exception as e:
            # Delete index.
            if index:
                index.delete()

            # Throw exception back to user.
            raise ExecutorError(str(e))

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 pathlib import Path	×
16
17	import pandas as pd	×
18
19	from evadb.catalog.catalog_type import VectorStoreType	×
20	from evadb.catalog.sql_config import ROW_NUM_COLUMN	×
21	from evadb.database import EvaDBDatabase	×
22	from evadb.executor.abstract_executor import AbstractExecutor	×
23	from evadb.executor.executor_utils import ExecutorError, handle_vector_store_params	×
24	from evadb.expression.function_expression import FunctionExpression	×
25	from evadb.models.storage.batch import Batch	×
26	from evadb.plan_nodes.create_index_plan import CreateIndexPlan	×
27	from evadb.third_party.databases.interface import get_database_handler	×
28	from evadb.third_party.vector_stores.types import FeaturePayload	×
29	from evadb.third_party.vector_stores.utils import VectorStoreFactory	×
30	from evadb.utils.logging_manager import logger	×
31
32
33	class CreateIndexExecutor(AbstractExecutor):	×
34	def __init__(self, db: EvaDBDatabase, node: CreateIndexPlan):	×
35	super().__init__(db, node)	×
36
37	self.name = self.node.name	×
38	self.if_not_exists = self.node.if_not_exists	×
39	self.table_ref = self.node.table_ref	×
40	self.col_list = self.node.col_list	×
41	self.vector_store_type = self.node.vector_store_type	×
42	self.project_expr_list = self.node.project_expr_list	×
43	self.index_def = self.node.index_def	×
44
45	def exec(self, args, *kwargs):	×
46	# Vector type specific creation.
47	if self.vector_store_type == VectorStoreType.PGVECTOR:	×
48	self._create_native_index()	×
49	else:
50	self._create_evadb_index()	×
51
52	yield Batch(	×
53	pd.DataFrame([f"Index {self.name} successfully added to the database."])
54	)
55
56	# Create index through the native storage engine.
57	def _create_native_index(self):	×
58	table = self.table_ref.table	×
59	db_catalog_entry = self.catalog().get_database_catalog_entry(	×
60	table.database_name
61	)
62
63	with get_database_handler(	×
64	db_catalog_entry.engine, **db_catalog_entry.params
65	) as handler:
66	# As other libraries, we default to HNSW and L2 distance.
67	resp = handler.execute_native_query(	×
68	f"""CREATE INDEX {self.name} ON {table.table_name}
69	USING hnsw ({self.col_list[0].name} vector_l2_ops)"""
70	)
71	if resp.error is not None:	×
72	raise ExecutorError(
73	f"Native engine create index encounters error: {resp.error}"
74	)
75
76	# On-disk saving path for EvaDB index.
77	def _get_evadb_index_save_path(self) -> Path:	×
78	index_dir = Path(self.config.get_value("storage", "index_dir"))	×
79	if not index_dir.exists():	×
80	index_dir.mkdir(parents=True, exist_ok=True)	×
81	return str(	×
82	index_dir / Path("{}_{}.index".format(self.vector_store_type, self.name))
83	)
84
85	# Create EvaDB index.
86	def _create_evadb_index(self):	×
87	# Find function expression.
88	function_expression, function_expression_signature = None, None	×
89	for project_expr in self.project_expr_list:	×
90	if isinstance(project_expr, FunctionExpression):	×
91	function_expression = project_expr	×
92	function_expression_signature = project_expr.signature()	×
93
94	# Get feature tables.
95	feat_tb_catalog_entry = self.table_ref.table.table_obj	×
96
97	# Get feature column.
98	feat_col_name = self.col_list[0].name	×
99	feat_col_catalog_entry = [	×
100	col for col in feat_tb_catalog_entry.columns if col.name == feat_col_name
101	][0]
102
103	if function_expression is not None:	×
104	feat_col_name = function_expression.output_objs[0].name	×
105
106	index_catalog_entry = self.catalog().get_index_catalog_entry_by_name(self.name)	×
107	index_path = self._get_evadb_index_save_path()	×
108
109	if index_catalog_entry is not None:	×
110	msg = f"Index {self.name} already exists."	×
111	if self.if_not_exists:	×
112	if (	×
113	index_catalog_entry.feat_column == feat_col_catalog_entry
114	and index_catalog_entry.function_signature
115	== function_expression_signature
116	and index_catalog_entry.type == self.node.vector_store_type
117	):
118	# Only update index if everything matches.
119	logger.warn(msg + " It will be updated on existing table.")	×
120	index = VectorStoreFactory.init_vector_store(	×
121	self.vector_store_type,
122	self.name,
123	**handle_vector_store_params(
124	self.vector_store_type, index_path
125	),
126	)
127	else:
128	# Skip index update if CREATE INDEX runs on a different index.
129	logger.warn(msg)	×
130	return	×
131	else:
132	logger.error(msg)	×
133	raise ExecutorError(msg)
134	else:
135	index = None	×
136
137	try:	×
138	# Add features to index.
139	for input_batch in self.children[0].exec():	×
140	input_batch.drop_column_alias()	×
141	feat = input_batch.column_as_numpy_array(feat_col_name)	×
142	row_num = input_batch.column_as_numpy_array(ROW_NUM_COLUMN)	×
143
144	for i in range(len(input_batch)):	×
145	row_feat = feat[i].reshape(1, -1)	×
146
147	# Create new index if not exists.
148	if index is None:	×
149	input_dim = row_feat.shape[1]	×
150	index = VectorStoreFactory.init_vector_store(	×
151	self.vector_store_type,
152	self.name,
153	**handle_vector_store_params(
154	self.vector_store_type, index_path
155	),
156	)
157	index.create(input_dim)	×
158
159	# Row ID for mapping back to the row.
160	index.add([FeaturePayload(row_num[i], row_feat)])	×
161
162	# Persist index.
163	index.persist()	×
164
165	# Save to catalog.
166	if index_catalog_entry is None:	×
167	self.catalog().insert_index_catalog_entry(	×
168	self.name,
169	index_path,
170	self.vector_store_type,
171	feat_col_catalog_entry,
172	function_expression_signature,
173	self.index_def,
174	)
175	except Exception as e:
176	# Delete index.
177	if index:
178	index.delete()
179
180	# Throw exception back to user.
181	raise ExecutorError(str(e))

georgia-tech-db / eva / #832

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