#793

Committed 19 Sep 2023 03:07PM UTC coverage: 73.096% (-1.0%) from 74.113%

Build # #793

Build Type

push

circle-ci

Committed by Jiashen Cao

Commit Message

fix lint issue

Run Details

8751 of 11972 relevant lines covered (73.1%)

0.73 hits per line

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

31.82

/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.models.storage.batch import Batch
from evadb.plan_nodes.create_index_plan import CreateIndexPlan
from evadb.storage.storage_engine import StorageEngine
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)

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

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

    # Create index through the native storage engine.
    def _create_native_index(self):
        table = self.node.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.node.name} ON {table.table_name}
                    USING hnsw ({self.node.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.node.vector_store_type, self.node.name))
        )

    # Create EvaDB index.
    def _create_evadb_index(self):
        if self.catalog().get_index_catalog_entry_by_name(self.node.name):
            msg = f"Index {self.node.name} already exists."
            if self.node.if_not_exists:
                logger.warn(msg)
                return
            else:
                logger.error(msg)
                raise ExecutorError(msg)

        index = None
        index_path = self._get_evadb_index_save_path()

        try:
            # Get feature tables.
            feat_catalog_entry = self.node.table_ref.table.table_obj

            # Get feature column.
            feat_col_name = self.node.col_list[0].name
            feat_column = [
                col for col in feat_catalog_entry.columns if col.name == feat_col_name
            ][0]

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

georgia-tech-db / eva / #793

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