Adding convenience method for NNS (#18)

* adding convenience nns method

* bumping version

* fixing copy-paste errors

* feedback

Author: silvanheller

cottontaildb_client/cottontaildb_client.py

@@ -8,7 +8,8 @@
     EntityDefinition, CreateEntityMessage, InsertMessage, ColumnName, Scan, From, Type, ListSchemaMessage, \
     ListEntityMessage, EntityDetailsMessage, DropEntityMessage, TruncateEntityMessage, OptimizeEntityMessage, \
     IndexName, IndexType, CreateIndexMessage, DropIndexMessage, RebuildIndexMessage, UpdateMessage, \
-    DeleteMessage, Literal, Vector, FloatVector, BatchInsertMessage, Metadata, QueryMessage, Query
+    DeleteMessage, Literal, Vector, FloatVector, BatchInsertMessage, Metadata, QueryMessage, Query, Expression, \
+    FunctionName, Function, Projection, Order
 from .cottontail_pb2_grpc import DDLStub, DMLStub, TXNStub, DQLStub
 
 
@@ -343,6 +344,42 @@ def ping(self):
         """Sends a ping message to the endpoint. If method returns without exception endpoint is connected."""
         self._dql.Ping(Empty())
 
+    def nns(self, schema, entity, query_vector, distance='manhattan', limit=None, vector_col='feature', id_col='id'):
+        """
+        Queries the specified entity with the given vector.
+
+        @param schema: the schema containing the queried entity
+        @param entity: the entity being queried
+        @param query_vector: the query vector. Simple float array.
+        @param distance: the distance to be used
+        @param limit: maximum number of rows to return
+        @param vector_col: column name where the vector is stored
+        @param id_col: column name where the id is stored
+        """
+        schema_name = SchemaName(name=schema)
+        entity_name = EntityName(schema=schema_name, name=entity)
+        nns_col = Expression(column=ColumnName(entity=entity_name, name=vector_col))
+
+        distance_col = ColumnName(name='distance')
+        id_expression = Expression(column=ColumnName(name=id_col))
+
+        nns_expression = Expression(literal=float_vector(*query_vector))
+        fn = FunctionName(name=distance)
+        fun = Function(name=fn, arguments=[nns_col, nns_expression])
+
+        expression = Expression(function=fun)
+
+        projection_element = Projection.ProjectionElement(alias=distance_col,
+                                                          expression=expression)
+        projection = Projection(op=Projection.ProjectionOperation.SELECT,
+                                elements=[projection_element, Projection.ProjectionElement(expression=id_expression)])
+
+        order_component = Order.Component(column=distance_col, direction=Order.Direction.ASCENDING)
+
+        order = Order(components=[order_component])
+
+        return self.query(schema, entity, projection, None, limit=limit, order=order)
+
     def query(self, schema, entity, projection, where, order=None, limit=None, skip=None, from_=None):
         """
         Queries the specified entity where the provided conditions are met and applies the given projection.

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = cottontaildb-client
-version = 0.14.0
+version = 0.14.1
 author = Florian Spiess
 author_email = [email protected]
 description = A Cottontail DB gRPC client.

tests/test_cottontaildb_client.py

@@ -9,9 +9,12 @@
 DB_HOST = 'localhost'
 DB_PORT = 1865
 
-TEST_SCHEMA = 'schema_test'
-TEST_ENTITY = 'entity_test'
-TEST_VECTOR_ENTITY = 'entity_test_vector'
+TEST_SCHEMA_STR = 'schema_test'
+TEST_ENTITY_STR = 'entity_test'
+TEST_VECTOR_ENTITY_STR = 'entity_test_vector'
+TEST_SCHEMA_NAME = SchemaName(name=TEST_SCHEMA_STR)
+TEST_ENTITY_NAME = EntityName(schema=TEST_SCHEMA_NAME, name=TEST_ENTITY_STR)
+TEST_VECTOR_ENTITY_NAME = EntityName(schema=TEST_SCHEMA_NAME, name=TEST_VECTOR_ENTITY_STR)
 TEST_INDEX = 'index_test'
 TEST_COLUMN_ID = 'id'
 TEST_COLUMN_VALUE = 'value'
@@ -32,74 +35,76 @@ def setUp(self):
             self.skipTest(f'error connecting to Cottontail DB at {DB_HOST}:{DB_PORT}')
 
     def tearDown(self):
-        if TEST_SCHEMA in [s.split('.')[-1] for s in self.client.list_schemas()]:
-            self.client.drop_schema(TEST_SCHEMA)
+        if TEST_SCHEMA_STR in [s.split('.')[-1] for s in self.client.list_schemas()]:
+            self.client.drop_schema(TEST_SCHEMA_STR)
         self.client.close()
 
     def test_create_drop_schema(self):
         self._create_schema()
-        self.assert_in(TEST_SCHEMA, self.client.list_schemas(), 'schema was not created')
-        self.client.create_schema(TEST_SCHEMA, exist_ok=True)
-        self.client.drop_schema(TEST_SCHEMA)
-        self.assert_not_in(TEST_SCHEMA, self.client.list_schemas(), 'schema was not dropped')
+        self.assert_in(TEST_SCHEMA_STR, self.client.list_schemas(), 'schema was not created')
+        self.client.create_schema(TEST_SCHEMA_STR, exist_ok=True)
+        self.client.drop_schema(TEST_SCHEMA_STR)
+        self.assert_not_in(TEST_SCHEMA_STR, self.client.list_schemas(), 'schema was not dropped')
 
     def test_create_drop_entity(self):
         self._create_schema()
         self._create_entity()
         self._create_vector_entity()
-        self.assert_in(TEST_ENTITY, self.client.list_entities(TEST_SCHEMA), 'entity was not created')
-        self.assert_in(TEST_VECTOR_ENTITY, self.client.list_entities(TEST_SCHEMA), 'vector entity was not created')
-        self.client.create_entity(TEST_SCHEMA, TEST_ENTITY, [], exist_ok=True)
-        self.client.create_entity(TEST_SCHEMA, TEST_VECTOR_ENTITY, [], exist_ok=True)
-        self.client.drop_entity(TEST_SCHEMA, TEST_ENTITY)
-        self.client.drop_entity(TEST_SCHEMA, TEST_VECTOR_ENTITY)
-        self.assert_not_in(TEST_ENTITY, self.client.list_entities(TEST_SCHEMA), 'entity was not dropped')
-        self.assert_not_in(TEST_VECTOR_ENTITY, self.client.list_entities(TEST_SCHEMA), 'vector entity was not dropped')
+        self.assert_in(TEST_ENTITY_STR, self.client.list_entities(TEST_SCHEMA_STR), 'entity was not created')
+        self.assert_in(TEST_VECTOR_ENTITY_STR, self.client.list_entities(TEST_SCHEMA_STR),
+                       'vector entity was not created')
+        self.client.create_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR, [], exist_ok=True)
+        self.client.create_entity(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR, [], exist_ok=True)
+        self.client.drop_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR)
+        self.client.drop_entity(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR)
+        self.assert_not_in(TEST_ENTITY_STR, self.client.list_entities(TEST_SCHEMA_STR), 'entity was not dropped')
+        self.assert_not_in(TEST_VECTOR_ENTITY_STR, self.client.list_entities(TEST_SCHEMA_STR),
+                           'vector entity was not dropped')
 
     def test_drop_not_exists_entity(self):
         self._create_schema()
-        result = self.client.drop_entity(TEST_SCHEMA, TEST_ENTITY, not_exist_ok=True)
+        result = self.client.drop_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR, not_exist_ok=True)
         self.assertIsNone(result, 'response received after dropping nonexistent entity')
 
     def test_truncate_not_exists_entity(self):
         self._create_schema()
-        result = self.client.truncate_entity(TEST_SCHEMA, TEST_ENTITY, not_exist_ok=True)
+        result = self.client.truncate_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR, not_exist_ok=True)
         self.assertIsNone(result, 'response received after truncating nonexistent entity')
 
     def test_create_truncate_entity(self):
         self._create_schema()
         self._create_entity()
-        self.assert_in(TEST_ENTITY, self.client.list_entities(TEST_SCHEMA), 'entity was not created')
-        self.client.truncate_entity(TEST_SCHEMA, TEST_ENTITY)
-        self.assert_in(TEST_ENTITY, self.client.list_entities(TEST_SCHEMA), 'entity was dropped')
+        self.assert_in(TEST_ENTITY_STR, self.client.list_entities(TEST_SCHEMA_STR), 'entity was not created')
+        self.client.truncate_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR)
+        self.assert_in(TEST_ENTITY_STR, self.client.list_entities(TEST_SCHEMA_STR), 'entity was dropped')
 
     def test_insert(self):
         self._create_schema()
         self._create_entity()
         self._insert()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(details['rows'], 1, 'unexpected number of rows in entity after insert')
 
     def test_vector_insert(self):
         self._create_schema()
         self._create_vector_entity()
         self._insert_vector()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_VECTOR_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR)
         self.assertEqual(details['rows'], 1, 'unexpected number of rows in vector entity after insert')
         print('success')
 
     def test_batch_insert(self):
         self._create_schema()
         self._create_entity()
         self._batch_insert()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(details['rows'], 3, 'unexpected number of rows in entity after batch insert')
 
     def test_batch_insert_vectors(self):
         self._create_schema()
         self._create_vector_entity()
         self._batch_insert_vectors()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_VECTOR_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR)
         self.assertEqual(details['rows'], 3, 'unexpected number of rows in entity after batch insert')
 
     def test_query(self):
@@ -110,6 +115,19 @@ def test_query(self):
         query_result = self._query_value_with_key(query_key)
         self.assertEqual(len(query_result), 1, 'unexpected number of rows returned from query')
 
+    def test_query_vectors(self):
+        self._create_schema()
+        self._create_vector_entity()
+        self._batch_insert_vectors()
+        query = [0.1, 0.2, 0.4]
+
+        results = self.client.nns(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR, query, vector_col='value')
+        self.assertEqual(len(results), 3, 'unexpected number of rows returned from query')
+        # test with limit
+        results = self.client.nns(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR, query, vector_col='value', limit=1)
+        print(results)
+        self.assertEqual(len(results), 1, 'unexpected number of rows returned from query')
+
     def test_update(self):
         self._create_schema()
         self._create_entity()
@@ -124,20 +142,20 @@ def test_optimize(self):
         self._create_schema()
         self._create_entity()
         self._batch_insert()
-        self.client.optimize_entity(TEST_SCHEMA, TEST_ENTITY)
+        self.client.optimize_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR)
 
     def test_create_rebuild_drop_index(self):
         self._create_schema()
         self._create_entity()
         self._batch_insert()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(len(details['indexes']), 0, 'unexpected number of indexes in entity before index creation')
-        self.client.create_index(TEST_SCHEMA, TEST_ENTITY, TEST_INDEX, IndexType.BTREE, [TEST_COLUMN_VALUE])
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        self.client.create_index(TEST_SCHEMA_STR, TEST_ENTITY_STR, TEST_INDEX, IndexType.BTREE, [TEST_COLUMN_VALUE])
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(len(details['indexes']), 1, 'index was not created')
-        self.client.rebuild_index(TEST_SCHEMA, TEST_ENTITY, TEST_INDEX)
-        self.client.drop_index(TEST_SCHEMA, TEST_ENTITY, TEST_INDEX)
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        self.client.rebuild_index(TEST_SCHEMA_STR, TEST_ENTITY_STR, TEST_INDEX)
+        self.client.drop_index(TEST_SCHEMA_STR, TEST_ENTITY_STR, TEST_INDEX)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(len(details['indexes']), 0, 'index was not dropped')
 
     def test_transaction_commit(self):
@@ -146,18 +164,18 @@ def test_transaction_commit(self):
         self.client.start_transaction()
         self._insert()
         self.client.commit_transaction()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(details['rows'], 1, 'unexpected number of rows in entity after committed insert')
 
     def test_transaction_abort(self):
         self._create_schema()
         self._create_entity()
         self.client.start_transaction()
         self._insert()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(details['rows'], 1, 'unexpected number of rows in entity after transaction insert')
         self.client.abort_transaction()
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(details['rows'], 0, 'unexpected number of rows in entity after aborted insert')
 
     def test_delete(self):
@@ -166,8 +184,8 @@ def test_delete(self):
         self._insert()
         where = Where(atomic=AtomicBooleanPredicate(left=ColumnName(name=TEST_COLUMN_VALUE), right=AtomicBooleanOperand(
             expressions=Expressions(expression=[Expression(literal=Literal(intData=0))])), op=ComparisonOperator.EQUAL))
-        self.client.delete(TEST_SCHEMA, TEST_ENTITY, where)
-        details = self.client.get_entity_details(TEST_SCHEMA, TEST_ENTITY)
+        self.client.delete(TEST_SCHEMA_STR, TEST_ENTITY_STR, where)
+        details = self.client.get_entity_details(TEST_SCHEMA_STR, TEST_ENTITY_STR)
         self.assertEqual(details['rows'], 0, 'unexpected number of rows in entity after delete')
 
     def assert_in(self, name, names, message):
@@ -179,51 +197,52 @@ def assert_not_in(self, name, names, message):
         self.assertNotIn(name, [n.split('.')[-1] for n in names], message)
 
     def _create_schema(self):
-        self.client.create_schema(TEST_SCHEMA)
+        self.client.create_schema(TEST_SCHEMA_STR)
 
     def _create_entity(self):
         columns = [
             column_def(TEST_COLUMN_ID, Type.STRING, nullable=False),
             column_def(TEST_COLUMN_VALUE, Type.INTEGER, nullable=False)
         ]
-        self.client.create_entity(TEST_SCHEMA, TEST_ENTITY, columns)
+        self.client.create_entity(TEST_SCHEMA_STR, TEST_ENTITY_STR, columns)
 
     def _create_vector_entity(self):
         columns = [
             column_def(TEST_COLUMN_ID, Type.STRING, nullable=False),
             column_def(TEST_COLUMN_VALUE, Type.FLOAT_VEC, length=3, nullable=False)
         ]
-        self.client.create_entity(TEST_SCHEMA, TEST_VECTOR_ENTITY, columns)
+        self.client.create_entity(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR, columns)
 
     def _insert(self):
-        values = {'id': Literal(stringData='test_0'), 'value': Literal(intData=0)}
-        self.client.insert(TEST_SCHEMA, TEST_ENTITY, values)
+        values = {'id': Literal(stringData='test_0'), TEST_COLUMN_VALUE: Literal(intData=0)}
+        self.client.insert(TEST_SCHEMA_STR, TEST_ENTITY_STR, values)
 
     def _insert_vector(self):
         value_list = [0.2, 0.3, 0.5]
-        values = {'id': Literal(stringData='test_0'), 'value': float_vector(*value_list)}
-        self.client.insert(TEST_SCHEMA, TEST_VECTOR_ENTITY, values)
+        values = {'id': Literal(stringData='test_0'), TEST_COLUMN_VALUE: float_vector(*value_list)}
+        self.client.insert(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR, values)
 
     def _batch_insert(self):
-        columns = ['id', 'value']
+        columns = ['id', TEST_COLUMN_VALUE]
         values = [
             [Literal(stringData='test_1'), Literal(intData=1)],
             [Literal(stringData='test_2'), Literal(intData=2)],
             [Literal(stringData='test_3'), Literal(intData=3)]
         ]
-        self.client.insert_batch(TEST_SCHEMA, TEST_ENTITY, columns, values)
+
+        self.client.insert_batch(TEST_SCHEMA_STR, TEST_ENTITY_STR, columns, values)
 
     def _batch_insert_vectors(self):
-        columns = ['id', 'value']
+        columns = ['id', TEST_COLUMN_VALUE]
         one = [0.1, 0.2, 0.3]
-        two = [0.000001, 0.2, 0.3]
-        three = [0.1, 0.2, 0.3]
+        two = [0.01, 0.02, 0.3]
+        three = [0.9, 0.9, 0.9]
         values = [
             [Literal(stringData='test_1'), float_vector(*one)],
             [Literal(stringData='test_2'), float_vector(*two)],
             [Literal(stringData='test_3'), float_vector(*three)]
         ]
-        self.client.insert_batch(TEST_SCHEMA, TEST_VECTOR_ENTITY, columns, values)
+        self.client.insert_batch(TEST_SCHEMA_STR, TEST_VECTOR_ENTITY_STR, columns, values)
 
     def _update_value_with_key(self, key, value):
         where = Where(atomic=AtomicBooleanPredicate(
@@ -233,12 +252,10 @@ def _update_value_with_key(self, key, value):
             op=ComparisonOperator.EQUAL
         ))
         updates = {TEST_COLUMN_VALUE: Expression(literal=Literal(intData=value))}
-        self.client.update(TEST_SCHEMA, TEST_ENTITY, where, updates)
+        self.client.update(TEST_SCHEMA_STR, TEST_ENTITY_STR, where, updates)
 
     def _query_value_with_key(self, key):
-        schema_name = SchemaName(name=TEST_SCHEMA)
-        entity_name = EntityName(schema=schema_name, name=TEST_ENTITY)
-        expression = Expression(column=ColumnName(entity=entity_name, name=TEST_COLUMN_VALUE))
+        expression = Expression(column=ColumnName(entity=TEST_ENTITY_NAME, name=TEST_COLUMN_VALUE))
         projection_element = Projection.ProjectionElement(expression=expression)
         projection = Projection(op=Projection.ProjectionOperation.SELECT, elements=[projection_element])
         where = Where(atomic=AtomicBooleanPredicate(
@@ -247,4 +264,4 @@ def _query_value_with_key(self, key):
                 expressions=Expressions(expression=[Expression(literal=Literal(stringData=key))])),
             op=ComparisonOperator.EQUAL
         ))
-        return self.client.query(TEST_SCHEMA, TEST_ENTITY, projection, where)
+        return self.client.query(TEST_SCHEMA_STR, TEST_ENTITY_STR, projection, where)