diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index af29730..e66d10c 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -17,9 +17,11 @@ jobs: name: Run Lint and Tests runs-on: ubuntu-latest strategy: + fail-fast: false matrix: python-version: [ '3.9' ] - + pymilvus-version: [ '2.5.*', '2.6.*' ] + steps: - name: Checkout code uses: actions/checkout@v3 @@ -35,6 +37,9 @@ jobs: - name: Lint run: hatch run lint:all + - name: Pin pymilvus ${{ matrix.pymilvus-version }} + run: hatch run pip install "pymilvus==${{ matrix.pymilvus-version }}" + - name: Run Milvus run: | wget https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh diff --git a/src/milvus_haystack/document_store.py b/src/milvus_haystack/document_store.py index 47a0f27..b4a781b 100644 --- a/src/milvus_haystack/document_store.py +++ b/src/milvus_haystack/document_store.py @@ -18,7 +18,7 @@ MilvusClient, MilvusException, RRFRanker, - utility, + connections, ) from pymilvus.client.abstract import BaseRanker from pymilvus.client.types import LoadState @@ -194,33 +194,28 @@ def __init__( self._check_function() # Create the connection to the server - if connection_args is None: - self.connection_args = DEFAULT_MILVUS_CONNECTION - self._milvus_client = MilvusClient( - **self.connection_args, - ) + resolved_args = connection_args if connection_args is not None else DEFAULT_MILVUS_CONNECTION + self.connection_args = resolved_args + self._milvus_client = MilvusClient(**resolved_args) self.alias = self.client._using - self.col: Optional[Collection] = None - - # Grab the existing collection if it exists - if utility.has_collection(self.collection_name, using=self.alias): - self.col = Collection( - self.collection_name, - using=self.alias, - ) + # Since pymilvus 2.6.x, MilvusClient no longer registers its connection + # with the ORM `connections` registry; register its handler so the ORM + # `Collection` exposed by `self.col` can find the connection. + if hasattr(connections, "_alias_handlers"): + connections._alias_handlers.setdefault(self.alias, self.client._get_connection()) + self._col_cache: Optional[Collection] = None + self._cache_key: Optional[str] = None + + # Apply properties to the existing collection if it exists + if self.client.has_collection(self.collection_name): if self.collection_properties is not None: - self.col.set_properties(self.collection_properties) - # If need to drop old, drop it - if drop_old and isinstance(self.col, Collection): - self.col.drop() - self.col = None + self.client.alter_collection_properties(self.collection_name, self.collection_properties) + # If need to drop old, drop it + if drop_old: + self._drop_collection() # Initialize the vector store - self._init( - partition_names=partition_names, - replica_number=replica_number, - timeout=timeout, - ) + self._init(timeout=timeout) self._dummy_value = 999.0 def _check_function(self): @@ -250,6 +245,33 @@ def client(self) -> MilvusClient: """Get client.""" return self._milvus_client + @property + def col(self) -> Optional[Collection]: + """The ORM Collection object, or None if the collection does not exist. + + Kept for backward compatibility with code that accesses `store.col`; + internal operations use the `MilvusClient` API via `self.client`. + """ + current_key = f"{self.collection_name}:{self.alias}" + if self._cache_key == current_key and self._col_cache is not None: + return self._col_cache + if self.client.has_collection(self.collection_name): + self._col_cache = Collection(self.collection_name, using=self.alias) + self._cache_key = current_key + return self._col_cache + return None + + def _drop_collection(self) -> None: + self.client.drop_collection(self.collection_name) + # pymilvus keeps the dropped collection's schema in a connection-level + # cache, so recreating a collection with the same name can read back + # stale schema. See https://github.com/milvus-io/pymilvus/issues/3058 + conn = self.client._get_connection() + if hasattr(conn, "schema_cache"): + conn.schema_cache.pop(self.collection_name, None) + self._col_cache = None + self._cache_key = None + def count_documents(self) -> int: """ Returns how many documents are present in the document store. @@ -260,9 +282,11 @@ def count_documents(self) -> int: logger.debug("No existing collection to count.") return 0 count_expr = "count(*)" - res = self.col.query( - expr="", + res = self.client.query( + self.collection_name, + filter="", output_fields=[count_expr], + partition_names=self.partition_names, ) doc_num = res[0][count_expr] return doc_num @@ -346,10 +370,12 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc # Perform the Query. try: - res = self.col.query( - expr=expr, + res = self.client.query( + self.collection_name, + filter=expr, output_fields=output_fields, limit=MAX_LIMIT_SIZE, + partition_names=self.partition_names, ) except MilvusException as err: logger.error("Failed to query documents with filters expr: %s", expr) @@ -432,16 +458,8 @@ def write_documents(self, documents: List[Document], policy: DuplicatePolicy = D return 0 # If the collection hasn't been initialized yet, perform all steps to do so - kwargs: Dict[str, Any] = {} - if not isinstance(self.col, Collection): - kwargs = {"embeddings": embeddings, "metas": metas} - if self.partition_names: - kwargs["partition_names"] = self.partition_names - if self.replica_number: - kwargs["replica_number"] = self.replica_number - if self.timeout: - kwargs["timeout"] = self.timeout - self._init(**kwargs) + if self.col is None: + self._init(embeddings=embeddings, metas=metas, timeout=self.timeout) insert_list: list[dict] = [] for i in range(len(ids)): @@ -468,7 +486,7 @@ def write_documents(self, documents: List[Document], policy: DuplicatePolicy = D total_count = len(insert_list) batch_size = 1000 wrote_ids = [] - if not isinstance(self.col, Collection): + if self.col is None: raise MilvusException(message="Collection is not initialized") for i in range(0, total_count, batch_size): # Grab end index @@ -476,9 +494,8 @@ def write_documents(self, documents: List[Document], policy: DuplicatePolicy = D batch_insert_list = insert_list[i:end] # Insert into the collection. try: - # res: Collection - res = self.col.insert(batch_insert_list, timeout=None, **kwargs) - wrote_ids.extend(res.primary_keys) + res = self.client.insert(self.collection_name, batch_insert_list, timeout=None) + wrote_ids.extend(res["ids"]) except MilvusException as err: logger.error("Failed to insert batch starting at entity: %s/%s", i, total_count) raise err @@ -495,7 +512,7 @@ def delete_documents(self, document_ids: List[str]) -> None: return None expr = "id in ['" + "','".join(document_ids) + "']" logger.info(expr) - self.col.delete(expr) + self.client.delete(self.collection_name, filter=expr) def to_dict(self) -> Dict[str, Any]: """ @@ -567,8 +584,6 @@ def _init( self, embeddings: Optional[List] = None, metas: Optional[List[Dict]] = None, - partition_names: Optional[List] = None, - replica_number: int = 1, timeout: Optional[float] = None, ) -> None: if embeddings is not None: @@ -576,11 +591,7 @@ def _init( self._extract_fields() self._create_index() self._create_search_params() - self._load( - partition_names=partition_names, - replica_number=replica_number, - timeout=timeout, - ) + self._load(timeout=timeout) def _create_collection(self, embeddings: list, metas: Optional[List[Dict]] = None) -> None: # Determine embedding dim @@ -625,29 +636,28 @@ def _create_collection(self, embeddings: list, metas: Optional[List[Dict]] = Non # Create the collection try: - self.col = Collection( - name=self.collection_name, + self.client.create_collection( + collection_name=self.collection_name, schema=schema, consistency_level=self.consistency_level, - using=self.alias, ) # Set the collection properties if they exist if self.collection_properties is not None: - self.col.set_properties(self.collection_properties) + self.client.alter_collection_properties(self.collection_name, self.collection_properties) except MilvusException as err: logger.error("Failed to create collection: %s error: %s", self.collection_name, err) raise err def _extract_fields(self) -> None: """Grab the existing fields from the Collection""" - if isinstance(self.col, Collection): - schema = self.col.schema - for x in schema.fields: - self.fields.append(x.name) + if self.col is not None: + schema_info = self.client.describe_collection(self.collection_name) + for x in schema_info["fields"]: + self.fields.append(x["name"]) def _create_index(self) -> None: """Create an index on the collection""" - if isinstance(self.col, Collection) and self._get_index() is None: + if self.col is not None and self._get_index() is None: try: # If no index params, use a default AUTOINDEX based one if self.index_params is None: @@ -658,11 +668,9 @@ def _create_index(self) -> None: } try: - self.col.create_index( - self._vector_field, - index_params=self.index_params, - using=self.alias, - ) + index_params = self.client.prepare_index_params() + index_params.add_index(field_name=self._vector_field, **self.index_params) + self.client.create_index(self.collection_name, index_params) # If default did not work, most likely on Zilliz Cloud except MilvusException: @@ -672,11 +680,10 @@ def _create_index(self) -> None: "index_type": "AUTOINDEX", "params": {}, } - self.col.create_index( - self._vector_field, - index_params=self.index_params, - using=self.alias, - ) + index_params = self.client.prepare_index_params() + index_params.add_index(field_name=self._vector_field, **self.index_params) + self.client.create_index(self.collection_name, index_params) + if self._sparse_vector_field: if self.sparse_index_params is None: if self._sparse_mode == EmbeddingMode.EMBEDDING_MODEL: @@ -690,11 +697,9 @@ def _create_index(self) -> None: "metric_type": "BM25", "params": {}, } - self.col.create_index( - self._sparse_vector_field, - index_params=self.sparse_index_params, - using=self.alias, - ) + sparse_index_params = self.client.prepare_index_params() + sparse_index_params.add_index(field_name=self._sparse_vector_field, **self.sparse_index_params) + self.client.create_index(self.collection_name, sparse_index_params) logger.debug( "Successfully created an index on collection: %s", @@ -707,7 +712,7 @@ def _create_index(self) -> None: def _create_search_params(self) -> None: """Generate search params based on the current index type""" - if isinstance(self.col, Collection) and self.search_params is None: + if self.col is not None and self.search_params is None: index = self._get_index() if index is not None: index_type: str = index["index_param"]["index_type"] @@ -717,29 +722,36 @@ def _create_search_params(self) -> None: def _get_index(self) -> Optional[Dict[str, Any]]: """Return the vector index information if it exists""" - if isinstance(self.col, Collection): - for x in self.col.indexes: + # `MilvusClient.describe_index` raises when the field has no index yet, + # so keep the ORM index iteration here: it simply yields nothing on the + # first-time index creation path. + col = self.col + if col is not None: + for x in col.indexes: if x.field_name == self._vector_field: return x.to_dict() return None - def _load( - self, - partition_names: Optional[list] = None, - replica_number: int = 1, - timeout: Optional[float] = None, - ) -> None: + def _load(self, timeout: Optional[float] = None) -> None: """Load the collection if available.""" if ( - isinstance(self.col, Collection) + self.col is not None and self._get_index() is not None - and utility.load_state(self.collection_name, using=self.alias) == LoadState.NotLoad + and self.client.get_load_state(self.collection_name)["state"] == LoadState.NotLoad ): - self.col.load( - partition_names=partition_names, - replica_number=replica_number, - timeout=timeout, - ) + if self.partition_names: + self.client.load_partitions( + self.collection_name, + partition_names=self.partition_names, + replica_number=self.replica_number, + timeout=timeout, + ) + else: + self.client.load_collection( + self.collection_name, + replica_number=self.replica_number, + timeout=timeout, + ) def _resolve_value(self, secret: Union[str, Secret]): if isinstance(secret, Secret): @@ -770,7 +782,7 @@ def _embedding_retrieval( # Build expr. if not filters: - expr = None + expr = "" else: expr = parse_filters(filters) @@ -778,13 +790,15 @@ def _embedding_retrieval( search_data = self._prepare_search_data( query_text=query_text, query_embedding=query_embedding, field=self._vector_field ) - res = self.col.search( + res = self.client.search( + self.collection_name, data=[search_data], anns_field=self._vector_field, - param=self.search_params, + search_params=self.search_params, limit=top_k, - expr=expr, + filter=expr, output_fields=output_fields, + partition_names=self.partition_names, timeout=None, ) distance_to_score_fn = self._select_score_fn() @@ -820,7 +834,7 @@ def _sparse_embedding_retrieval( # Build expr. if not filters: - expr = None + expr = "" else: expr = parse_filters(filters) @@ -829,13 +843,15 @@ def _sparse_embedding_retrieval( query_text=query_text, query_embedding=query_sparse_embedding, field=self._sparse_vector_field ) - res = self.col.search( + res = self.client.search( + self.collection_name, data=[search_data], anns_field=self._sparse_vector_field, - param=self.sparse_search_params, + search_params=self.sparse_search_params, limit=top_k, - expr=expr, + filter=expr, output_fields=output_fields, + partition_names=self.partition_names, timeout=None, ) docs = self._parse_search_result(res) @@ -875,7 +891,7 @@ def _hybrid_retrieval( # Build expr. if not filters: - expr = None + expr = "" else: expr = parse_filters(filters) @@ -897,14 +913,21 @@ def _hybrid_retrieval( ) # Search topK docs based on dense and sparse vectors and rerank. - res = self.col.hybrid_search([dense_req, sparse_req], rerank=reranker, limit=top_k, output_fields=output_fields) + res = self.client.hybrid_search( + self.collection_name, + [dense_req, sparse_req], + ranker=reranker, + limit=top_k, + output_fields=output_fields, + partition_names=self.partition_names, + ) docs = self._parse_search_result(res) return docs def _parse_search_result(self, result, distance_to_score_fn=lambda x: x) -> List[Document]: docs = [] for res in result[0]: - data = {x: res.entity.get(x) for x in res.entity.fields} + data = dict(res.fields) doc = self._parse_document(data) doc.score = distance_to_score_fn(res.distance) docs.append(doc)