feat: Add vector_search_by_key method to sync and async clients vec-330

src/aerospike_vector_search/aio/client.py

@@ -489,6 +489,95 @@ async def is_indexed(
             raise types.AVSServerError(rpc_error=e)
         return self._respond_is_indexed(response)
 
+    async def vector_search_by_key(
+        self,
+        *,
+        search_namespace: str,
+        index_name: str,
+        key: Union[int, str, bytes, bytearray],
+        key_namespace: str,
+        vector_field: str,
+        limit: int,
+        key_set: Optional[str] = None,
+        search_params: Optional[types.HnswSearchParams] = None,
+        include_fields: Optional[list[str]] = None,
+        exclude_fields: Optional[list[str]] = None,
+        timeout: Optional[int] = None,
+    ) -> list[types.Neighbor]:
+        """
+        Perform a Hierarchical Navigable Small World (HNSW) vector search in Aerospike Vector Search by primary record key.
+
+        :param search_namespace: The namespace that stores the records to be searched.
+        :type search_namespace: str
+
+        :param index_name: The name of the index to use in the search.
+        :type index_name: str
+
+        :param key: The primary key of the record that stores the vector to use in the search.
+        :type key: Union[int, str, bytes, bytearray]
+
+        :param key_namespace: The namespace that stores the record.
+        :type key_namespace: str
+
+        :param vector_field: The name of the field containing vector data.
+        :type vector_field: str
+
+        :param limit: The maximum number of neighbors to return. K value.
+        :type limit: int
+
+        :param key_set: The name of the set from which to read the record to search by. Defaults to None.
+        :type key_set: Optional[str]
+
+        :param search_params: Parameters for the HNSW algorithm.
+            If None, the default parameters for the index are used. Defaults to None.
+        :type search_params: Optional[types_pb2.HnswSearchParams]
+
+        :param include_fields: A list of field names to retrieve from the results.
+            When used, fields that are not included are not sent by the server,
+            saving on network traffic.
+            If a field is listed in both include_fields and exclude_fields,
+            exclude_fields takes priority, and the field is not returned.
+            If None, all fields are retrieved. Defaults to None.
+        :type include_fields: Optional[list[str]]
+
+        :param exclude_fields: A list of field names to exclude from the results.
+            When used, the excluded fields are not sent by the server,
+            saving on network traffic.
+            If None, all fields are retrieved. Defaults to None.
+        :type exclude_fields: Optional[list[str]]
+
+        :param timeout: Time in seconds this operation will wait before raising an :class:`AVSServerError <aerospike_vector_search.types.AVSServerError>`. Defaults to None.
+        :type timeout: int
+
+        Returns:
+            list[types.Neighbor]: A list of neighbors records found by the search.
+
+        Raises:
+            AVSServerError: Raised if an error occurs during the RPC communication with the server while attempting to vector search.
+            This error could occur due to various reasons such as network issues, server-side failures, or invalid request parameters.
+        """
+        rec_and_key = await self.get(
+            namespace=key_namespace,
+            key=key,
+            set_name=key_set,
+            timeout=timeout,
+        )
+
+        vector = rec_and_key.fields[vector_field]
+
+        neighbors = await self.vector_search(
+            namespace=search_namespace,
+            index_name=index_name,
+            query=vector,
+            limit=limit,
+            search_params=search_params,
+            include_fields=include_fields,
+            exclude_fields=exclude_fields,
+            timeout=timeout,
+        )
+
+        return neighbors
+
     async def vector_search(
         self,
         *,

src/aerospike_vector_search/client.py

@@ -468,6 +468,100 @@ def is_indexed(
             raise types.AVSServerError(rpc_error=e)
         return self._respond_is_indexed(response)
 
+    def vector_search_by_key(
+        self,
+        *,
+        search_namespace: str,
+        index_name: str,
+        key: Union[int, str, bytes, bytearray],
+        key_namespace: str,
+        vector_field: str,
+        limit: int,
+        key_set: Optional[str] = None,
+        search_params: Optional[types.HnswSearchParams] = None,
+        include_fields: Optional[list[str]] = None,
+        exclude_fields: Optional[list[str]] = None,
+        timeout: Optional[int] = None,
+    ) -> list[types.Neighbor]:
+        """
+        Perform a Hierarchical Navigable Small World (HNSW) vector search in Aerospike Vector Search by primary record key.
+
+        :param search_namespace: The namespace that stores the records to be searched.
+        :type search_namespace: str
+
+        :param index_name: The name of the index to use in the search.
+        :type index_name: str
+
+        :param key: The primary key of the record that stores the vector to use in the search.
+        :type key: Union[int, str, bytes, bytearray]
+
+        :param key_namespace: The namespace that stores the record.
+        :type key_namespace: str
+
+        :param vector_field: The name of the field containing vector data.
+        :type vector_field: str
+
+        :param limit: The maximum number of neighbors to return. K value.
+        :type limit: int
+
+        :param key_set: The name of the set from which to read the record to search by. Defaults to None.
+        :type key_set: Optional[str]
+
+        :param search_params: Parameters for the HNSW algorithm.
+            If None, the default parameters for the index are used. Defaults to None.
+        :type search_params: Optional[types_pb2.HnswSearchParams]
+
+        :param include_fields: A list of field names to retrieve from the results.
+            When used, fields that are not included are not sent by the server,
+            saving on network traffic.
+            If a field is listed in both include_fields and exclude_fields,
+            exclude_fields takes priority, and the field is not returned.
+            If None, all fields are retrieved. Defaults to None.
+        :type include_fields: Optional[list[str]]
+
+        :param exclude_fields: A list of field names to exclude from the results.
+            When used, the excluded fields are not sent by the server,
+            saving on network traffic.
+            If None, all fields are retrieved. Defaults to None.
+        :type exclude_fields: Optional[list[str]]
+
+        :param timeout: Time in seconds this operation will wait before raising an :class:`AVSServerError <aerospike_vector_search.types.AVSServerError>`. Defaults to None.
+        :type timeout: int
+
+        :param field_names: Deprecated, use include_fields instead.
+        :type field_names: Optional[list[str]]
+
+        Returns:
+            list[types.Neighbor]: A list of neighbors records found by the search.
+
+        Raises:
+            AVSServerError: Raised if an error occurs during the RPC communication with the server while attempting to vector search.
+            This error could occur due to various reasons such as network issues, server-side failures, or invalid request parameters.
+        """
+        rec_and_key = self.get(
+            namespace=key_namespace,
+            key=key,
+            set_name=key_set,
+            timeout=timeout,
+        )
+
+        vector = rec_and_key.fields[vector_field]
+
+        neighbors = self.vector_search(
+            namespace=search_namespace,
+            index_name=index_name,
+            query=vector,
+            limit=limit,
+            search_params=search_params,
+            include_fields=include_fields,
+            exclude_fields=exclude_fields,
+            timeout=timeout,
+        )
+
+        return neighbors
+
+
+
     def vector_search(
         self,
         *,

src/aerospike_vector_search/types.py

@@ -43,6 +43,13 @@ def __init__(self, *, namespace: str, set: str, key: Any) -> None:
         self.set = set
         self.key = key
 
+    def __repr__(self) -> str:
+        return (
+            f"Key(namespace={self.namespace}, "
+            f"set={self.set}, "
+            f"key={self.key})"
+        )
+
     def __str__(self):
         """
         Returns a string representation of the key.
@@ -128,6 +135,13 @@ def __init__(self, *, key: Key, fields: dict[str, Any], distance: float) -> None
         self.fields = fields
         self.distance = distance
 
+    def __repr__(self) -> str:
+        return (
+            f"Neighbor(key={self.key}, "
+            f"fields={self.fields}, "
+            f"distance={self.distance})"
+        )
+
     def __str__(self):
         """
         Returns a string representation of the neighboring record.
@@ -153,15 +167,13 @@ def __str__(self):
     def __eq__(self, other) -> bool:
         if not isinstance(other, Neighbor):
             return NotImplemented
-
         return (
             self.distance == other.distance
             and self.key == other.key
             and self.fields == other.fields
         )
 
 
-
 class VectorDistanceMetric(enum.Enum):
     """
     Enumeration of vector distance metrics used for comparing vectors.

tests/aerospike.conf

@@ -57,7 +57,7 @@ network {
 namespace avs-meta {
   replication-factor 2
 storage-engine memory {
-    data-size 2G
+    data-size 1G
 }
 nsup-period 100
 }

tests/assets/aerospike.conf

@@ -57,7 +57,7 @@ network {
 namespace avs-meta {
   replication-factor 2
 storage-engine memory {
-    data-size 2G
+    data-size 1G
 }
 nsup-period 100
 }

tests/standard/aio/test_extensive_vector_search.py

@@ -205,8 +205,12 @@ async def test_vector_search_with_set_same_as_index(
     query_numpy,
     session_vector_client,
     session_admin_client,
+    extensive_vector_search,
 ):
 
+    if not extensive_vector_search:
+        pytest.skip("Extensive vector tests disabled")
+
     await session_admin_client.index_create(
         namespace="test",
         name="demo2",