Refactor/shard key helpers (#19765)

* Refactor shard key creation into separate methods for improved readability and maintainability

* Add tests for shard key and payload index creation in QdrantVectorStore

* Bump version to 0.8.3 in pyproject.toml

* chore: Import`UnexpectedResponse` on module level instead of inside multiple functions

Author: sumitaryal

llama-index-integrations/vector_stores/llama-index-vector-stores-qdrant/llama_index/vector_stores/qdrant/base.py

@@ -52,6 +52,7 @@
     IsEmptyCondition,
 )
 from qdrant_client.qdrant_fastembed import IDF_EMBEDDING_MODELS
+from qdrant_client.http.exceptions import UnexpectedResponse
 
 logger = logging.getLogger(__name__)
 import_err_msg = (
@@ -575,8 +576,6 @@ async def async_add(
             ValueError: If trying to using async methods without aclient
 
         """
-        from qdrant_client.http.exceptions import UnexpectedResponse
-
         self._ensure_async_client()
 
         collection_initialized = await self._acollection_exists(self.collection_name)
@@ -799,9 +798,6 @@ def client(self) -> Any:
 
     def _create_collection(self, collection_name: str, vector_size: int) -> None:
         """Create a Qdrant collection."""
-        from qdrant_client.http import models as rest
-        from qdrant_client.http.exceptions import UnexpectedResponse
-
         dense_config = self._dense_config or rest.VectorParams(
             size=vector_size,
             distance=rest.Distance.COSINE,
@@ -843,11 +839,7 @@ def _create_collection(self, collection_name: str, vector_size: int) -> None:
                 )
 
             if self._shard_keys:
-                for shard_key in self._shard_keys:
-                    self._client.create_shard_key(
-                        collection_name=collection_name,
-                        shard_key=shard_key,
-                    )
+                self._create_shard_keys()
 
             # To improve search performance Qdrant recommends setting up
             # a payload index for fields used in filters.
@@ -870,30 +862,15 @@ def _create_collection(self, collection_name: str, vector_size: int) -> None:
             )
 
             if self._shard_keys:
-                for shard_key in self._shard_keys:
-                    try:
-                        self._client.create_shard_key(
-                            collection_name=collection_name,
-                            shard_key=shard_key,
-                        )
-                    except (RpcError, ValueError, UnexpectedResponse) as exc:
-                        if "already exists" not in str(exc):
-                            raise exc  # noqa: TRY201
-                        logger.warning(
-                            "Shard key %s already exists, skipping creation.",
-                            shard_key,
-                        )
-                        continue
+                self._create_shard_keys()
+
             if self._payload_indexes:
                 self._create_payload_indexes()
 
         self._collection_initialized = True
 
     async def _acreate_collection(self, collection_name: str, vector_size: int) -> None:
         """Asynchronous method to create a Qdrant collection."""
-        from qdrant_client.http import models as rest
-        from qdrant_client.http.exceptions import UnexpectedResponse
-
         dense_config = self._dense_config or rest.VectorParams(
             size=vector_size,
             distance=rest.Distance.COSINE,
@@ -932,11 +909,7 @@ async def _acreate_collection(self, collection_name: str, vector_size: int) -> N
                 )
 
             if self._shard_keys:
-                for shard_key in self._shard_keys:
-                    await self._aclient.create_shard_key(
-                        collection_name=collection_name,
-                        shard_key=shard_key,
-                    )
+                await self._acreate_shard_keys()
 
             # To improve search performance Qdrant recommends setting up
             # a payload index for fields used in filters.
@@ -959,20 +932,8 @@ async def _acreate_collection(self, collection_name: str, vector_size: int) -> N
             )
 
             if self._shard_keys:
-                for shard_key in self._shard_keys:
-                    try:
-                        await self._client.create_shard_key(
-                            collection_name=collection_name,
-                            shard_key=shard_key,
-                        )
-                    except (RpcError, ValueError, UnexpectedResponse) as exc:
-                        if "already exists" not in str(exc):
-                            raise exc  # noqa: TRY201
-                        logger.warning(
-                            "Shard key %s already exists, skipping creation.",
-                            shard_key,
-                        )
-                        continue
+                await self._acreate_shard_keys()
+
             if self._payload_indexes:
                 await self._acreate_payload_indexes()
 
@@ -986,6 +947,46 @@ async def _acollection_exists(self, collection_name: str) -> bool:
         """Asynchronous method to check if a collection exists."""
         return await self._aclient.collection_exists(collection_name)
 
+    def _create_shard_keys(self) -> None:
+        """Create shard keys in Qdrant collection."""
+        if not self._shard_keys:
+            return
+
+        for shard_key in self._shard_keys:
+            try:
+                self._client.create_shard_key(
+                    collection_name=self.collection_name,
+                    shard_key=shard_key,
+                )
+            except (RpcError, ValueError, UnexpectedResponse) as exc:
+                if "already exists" not in str(exc):
+                    raise exc  # noqa: TRY201
+                logger.warning(
+                    "Shard key %s already exists, skipping creation.",
+                    shard_key,
+                )
+                continue
+
+    async def _acreate_shard_keys(self) -> None:
+        """Asynchronous method to create shard keys in Qdrant collection."""
+        if not self._shard_keys:
+            return
+
+        for shard_key in self._shard_keys:
+            try:
+                await self._aclient.create_shard_key(
+                    collection_name=self.collection_name,
+                    shard_key=shard_key,
+                )
+            except (RpcError, ValueError, UnexpectedResponse) as exc:
+                if "already exists" not in str(exc):
+                    raise exc  # noqa: TRY201
+                logger.warning(
+                    "Shard key %s already exists, skipping creation.",
+                    shard_key,
+                )
+                continue
+
     def _create_payload_indexes(self) -> None:
         """Create payload indexes in Qdrant collection."""
         if not self._payload_indexes:

llama-index-integrations/vector_stores/llama-index-vector-stores-qdrant/pyproject.toml

@@ -28,7 +28,7 @@ dev = [
 
 [project]
 name = "llama-index-vector-stores-qdrant"
-version = "0.8.2"
+version = "0.8.3"
 description = "llama-index vector_stores qdrant integration"
 authors = [{name = "Your Name", email = "[email protected]"}]
 requires-python = ">=3.9,<3.14"

llama-index-integrations/vector_stores/llama-index-vector-stores-qdrant/tests/test_vector_stores_qdrant.py

@@ -653,3 +653,44 @@ def test_payload_indexes_created_on_existing_collection_sync(
         missing = schema_keys - payload_keys
         extra = payload_keys - schema_keys
         assert not missing and not extra, f"Missing: {missing}, Extra: {extra}"
+
+
+@requires_qdrant_cluster
+@pytest.mark.asyncio
+async def test_arecreate_shard_keys(
+    shard_vector_store: QdrantVectorStore,
+):
+    await shard_vector_store._acreate_shard_keys()
+
+
+@requires_qdrant_cluster
+def test_recreate_shard_keys(
+    shard_vector_store: QdrantVectorStore,
+):
+    shard_vector_store._create_shard_keys()
+
+
+@pytest.mark.asyncio
+async def test_acreate_shard_keys_returns_early_when_no_shard_keys(
+    vector_store: QdrantVectorStore,
+):
+    await vector_store._acreate_shard_keys()
+
+
+def test_create_shard_keys_returns_early_when_no_shard_keys(
+    vector_store: QdrantVectorStore,
+):
+    vector_store._create_shard_keys()
+
+
+@pytest.mark.asyncio
+async def test_acreate_payload_indexes_returns_early_when_no_payload_indexes(
+    vector_store: QdrantVectorStore,
+):
+    await vector_store._acreate_payload_indexes()
+
+
+def test_create_payload_indexes_returns_early_when_no_payload_indexes(
+    vector_store: QdrantVectorStore,
+):
+    vector_store._create_payload_indexes()