PYTHON-4208 [Vector Search GA] Add support for types in search index creation (#1544)
This commit is contained in:
Noah Stapp
GitHub
parent
944bea3892
commit
17c64a6dd8
@ -2410,7 +2410,7 @@ class Collection(common.BaseObject, Generic[_DocumentType]):
|
||||
.. versionadded:: 4.5
|
||||
"""
|
||||
if not isinstance(model, SearchIndexModel):
|
||||
model = SearchIndexModel(model["definition"], model.get("name"))
|
||||
model = SearchIndexModel(model["definition"], model.get("name"), model.get("type"))
|
||||
return self.create_search_indexes([model], session, comment, **kwargs)[0]
|
||||
|
||||
def create_search_indexes(
|
||||
|
||||
@ -588,7 +588,12 @@ class SearchIndexModel:
|
||||
|
||||
__slots__ = ("__document",)
|
||||
|
||||
def __init__(self, definition: Mapping[str, Any], name: Optional[str] = None) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
definition: Mapping[str, Any],
|
||||
name: Optional[str] = None,
|
||||
type: Optional[str] = "search",
|
||||
) -> None:
|
||||
"""Create a Search Index instance.
|
||||
|
||||
For use with :meth:`~pymongo.collection.Collection.create_search_index` and :meth:`~pymongo.collection.Collection.create_search_indexes`.
|
||||
@ -604,6 +609,7 @@ class SearchIndexModel:
|
||||
self.__document = dict(name=name, definition=definition)
|
||||
else:
|
||||
self.__document = dict(definition=definition)
|
||||
self.__document["type"] = type # type: ignore[assignment]
|
||||
|
||||
@property
|
||||
def document(self) -> Mapping[str, Any]:
|
||||
|
||||
@ -40,52 +40,6 @@
|
||||
"tests": [
|
||||
{
|
||||
"description": "no name provided for an index definition",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndex",
|
||||
"object": "collection0",
|
||||
"arguments": {
|
||||
"model": {
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"expectError": {
|
||||
"isError": true,
|
||||
"errorContains": "Atlas"
|
||||
}
|
||||
}
|
||||
],
|
||||
"expectEvents": [
|
||||
{
|
||||
"client": "client0",
|
||||
"events": [
|
||||
{
|
||||
"commandStartedEvent": {
|
||||
"command": {
|
||||
"createSearchIndexes": "collection0",
|
||||
"indexes": [
|
||||
{
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "name provided for an index definition",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndex",
|
||||
@ -97,7 +51,7 @@
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index"
|
||||
"type": "search"
|
||||
}
|
||||
},
|
||||
"expectError": {
|
||||
@ -121,7 +75,117 @@
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index"
|
||||
"type": "search"
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "name provided for an index definition",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndex",
|
||||
"object": "collection0",
|
||||
"arguments": {
|
||||
"model": {
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "search"
|
||||
}
|
||||
},
|
||||
"expectError": {
|
||||
"isError": true,
|
||||
"errorContains": "Atlas"
|
||||
}
|
||||
}
|
||||
],
|
||||
"expectEvents": [
|
||||
{
|
||||
"client": "client0",
|
||||
"events": [
|
||||
{
|
||||
"commandStartedEvent": {
|
||||
"command": {
|
||||
"createSearchIndexes": "collection0",
|
||||
"indexes": [
|
||||
{
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "search"
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "create a vector search index",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndex",
|
||||
"object": "collection0",
|
||||
"arguments": {
|
||||
"model": {
|
||||
"definition": {
|
||||
"fields": [
|
||||
{
|
||||
"type": "vector",
|
||||
"path": "plot_embedding",
|
||||
"numDimensions": 1536,
|
||||
"similarity": "euclidean"
|
||||
}
|
||||
]
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "vectorSearch"
|
||||
}
|
||||
},
|
||||
"expectError": {
|
||||
"isError": true,
|
||||
"errorContains": "Atlas"
|
||||
}
|
||||
}
|
||||
],
|
||||
"expectEvents": [
|
||||
{
|
||||
"client": "client0",
|
||||
"events": [
|
||||
{
|
||||
"commandStartedEvent": {
|
||||
"command": {
|
||||
"createSearchIndexes": "collection0",
|
||||
"indexes": [
|
||||
{
|
||||
"definition": {
|
||||
"fields": [
|
||||
{
|
||||
"type": "vector",
|
||||
"path": "plot_embedding",
|
||||
"numDimensions": 1536,
|
||||
"similarity": "euclidean"
|
||||
}
|
||||
]
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "vectorSearch"
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
|
||||
@ -72,54 +72,6 @@
|
||||
},
|
||||
{
|
||||
"description": "no name provided for an index definition",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndexes",
|
||||
"object": "collection0",
|
||||
"arguments": {
|
||||
"models": [
|
||||
{
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"expectError": {
|
||||
"isError": true,
|
||||
"errorContains": "Atlas"
|
||||
}
|
||||
}
|
||||
],
|
||||
"expectEvents": [
|
||||
{
|
||||
"client": "client0",
|
||||
"events": [
|
||||
{
|
||||
"commandStartedEvent": {
|
||||
"command": {
|
||||
"createSearchIndexes": "collection0",
|
||||
"indexes": [
|
||||
{
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "name provided for an index definition",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndexes",
|
||||
@ -132,7 +84,7 @@
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index"
|
||||
"type": "search"
|
||||
}
|
||||
]
|
||||
},
|
||||
@ -157,7 +109,121 @@
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index"
|
||||
"type": "search"
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "name provided for an index definition",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndexes",
|
||||
"object": "collection0",
|
||||
"arguments": {
|
||||
"models": [
|
||||
{
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "search"
|
||||
}
|
||||
]
|
||||
},
|
||||
"expectError": {
|
||||
"isError": true,
|
||||
"errorContains": "Atlas"
|
||||
}
|
||||
}
|
||||
],
|
||||
"expectEvents": [
|
||||
{
|
||||
"client": "client0",
|
||||
"events": [
|
||||
{
|
||||
"commandStartedEvent": {
|
||||
"command": {
|
||||
"createSearchIndexes": "collection0",
|
||||
"indexes": [
|
||||
{
|
||||
"definition": {
|
||||
"mappings": {
|
||||
"dynamic": true
|
||||
}
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "search"
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"description": "create a vector search index",
|
||||
"operations": [
|
||||
{
|
||||
"name": "createSearchIndexes",
|
||||
"object": "collection0",
|
||||
"arguments": {
|
||||
"models": [
|
||||
{
|
||||
"definition": {
|
||||
"fields": [
|
||||
{
|
||||
"type": "vector",
|
||||
"path": "plot_embedding",
|
||||
"numDimensions": 1536,
|
||||
"similarity": "euclidean"
|
||||
}
|
||||
]
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "vectorSearch"
|
||||
}
|
||||
]
|
||||
},
|
||||
"expectError": {
|
||||
"isError": true,
|
||||
"errorContains": "Atlas"
|
||||
}
|
||||
}
|
||||
],
|
||||
"expectEvents": [
|
||||
{
|
||||
"client": "client0",
|
||||
"events": [
|
||||
{
|
||||
"commandStartedEvent": {
|
||||
"command": {
|
||||
"createSearchIndexes": "collection0",
|
||||
"indexes": [
|
||||
{
|
||||
"definition": {
|
||||
"fields": [
|
||||
{
|
||||
"type": "vector",
|
||||
"path": "plot_embedding",
|
||||
"numDimensions": 1536,
|
||||
"similarity": "euclidean"
|
||||
}
|
||||
]
|
||||
},
|
||||
"name": "test index",
|
||||
"type": "vectorSearch"
|
||||
}
|
||||
],
|
||||
"$db": "database0"
|
||||
|
||||
@ -217,6 +217,70 @@ class TestSearchIndexProse(unittest.TestCase):
|
||||
# Run a ``dropSearchIndex`` command and assert that no error is thrown.
|
||||
coll0.drop_search_index("foo")
|
||||
|
||||
def test_case_7(self):
|
||||
"""Driver handles index types."""
|
||||
|
||||
# Create a collection with the "create" command using a randomly generated name (referred to as ``coll0``).
|
||||
coll0 = self.db[f"col{uuid.uuid4()}"]
|
||||
coll0.insert_one({})
|
||||
|
||||
# Use these search and vector search definitions for indexes.
|
||||
search_definition = {"mappings": {"dynamic": False}}
|
||||
vector_search_definition = {
|
||||
"fields": [
|
||||
{
|
||||
"type": "vector",
|
||||
"path": "plot_embedding",
|
||||
"numDimensions": 1536,
|
||||
"similarity": "euclidean",
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
# Create a new search index on ``coll0`` that implicitly passes its type.
|
||||
implicit_search_resp = coll0.create_search_index(
|
||||
model={"name": _NAME + "-implicit", "definition": search_definition}
|
||||
)
|
||||
|
||||
# Get the index definition.
|
||||
resp = coll0.list_search_indexes(name=implicit_search_resp).next()
|
||||
|
||||
# Assert that the index model contains the correct index type: ``"search"``.
|
||||
self.assertEqual(resp["type"], "search")
|
||||
|
||||
# Create a new search index on ``coll0`` that explicitly passes its type.
|
||||
explicit_search_resp = coll0.create_search_index(
|
||||
model={"name": _NAME + "-explicit", "type": "search", "definition": search_definition}
|
||||
)
|
||||
|
||||
# Get the index definition.
|
||||
resp = coll0.list_search_indexes(name=explicit_search_resp).next()
|
||||
|
||||
# Assert that the index model contains the correct index type: ``"search"``.
|
||||
self.assertEqual(resp["type"], "search")
|
||||
|
||||
# Create a new vector search index on ``coll0`` that explicitly passes its type.
|
||||
explicit_vector_resp = coll0.create_search_index(
|
||||
model={
|
||||
"name": _NAME + "-vector",
|
||||
"type": "vectorSearch",
|
||||
"definition": vector_search_definition,
|
||||
}
|
||||
)
|
||||
|
||||
# Get the index definition.
|
||||
resp = coll0.list_search_indexes(name=explicit_vector_resp).next()
|
||||
|
||||
# Assert that the index model contains the correct index type: ``"vectorSearch"``.
|
||||
self.assertEqual(resp["type"], "vectorSearch")
|
||||
|
||||
# Catch the error raised when trying to create a vector search index without specifying the type
|
||||
with self.assertRaises(OperationFailure) as e:
|
||||
coll0.create_search_index(
|
||||
model={"name": _NAME + "-error", "definition": vector_search_definition}
|
||||
)
|
||||
self.assertIn("Attribute mappings missing.", e.exception.details["errmsg"])
|
||||
|
||||
|
||||
globals().update(
|
||||
generate_test_classes(
|
||||
|
||||
Loading…
Reference in New Issue
Block a user