Add PyMongo + Document DB Playground

README.md

@@ -17,9 +17,10 @@ playground's README states what it needs.
 | --------------------------------- | ---------------------- | ------------------------------------------------------------------------- |
 | [mongoose](playgrounds/mongoose/) | Node.js — Mongoose ODM | Express REST API + a CRUD/compatibility test suite using the Mongoose ODM. |
 | [beanie](playgrounds/beanie/)     | Python — Beanie ODM    | FastAPI REST API + a CRUD/compatibility test suite using the Beanie ODM.  |
+| [pymongo](playgrounds/pymongo/)   | Python — PyMongo driver | Flask REST API + a CRUD/compatibility test suite using the raw PyMongo driver. |
 
-More playgrounds are planned (for example **PyMongo** and other MongoDB
-drivers). Contributions are welcome.
+More playgrounds are planned (for example the **MongoDB Node.js native driver**
+and other MongoDB drivers). Contributions are welcome.
 
 ## Getting Started
 
@@ -45,6 +46,14 @@ cd playgrounds/beanie
 ./scripts/run-app.sh    # or run the demo REST API
 ```
 
+To try the PyMongo playground:
+
+```bash
+cd playgrounds/pymongo
+./scripts/run-test.sh   # start DocumentDB locally and run the compatibility suite
+./scripts/run-app.sh    # or run the demo REST API
+```
+
 ## Repository Layout
 
 ```
@@ -53,7 +62,8 @@ documentdb-playground/
 ├── LICENSE
 └── playgrounds/
     ├── mongoose/     # Node.js + Mongoose ODM
-    └── beanie/       # Python + Beanie ODM
+    ├── beanie/       # Python + Beanie ODM
+    └── pymongo/      # Python + PyMongo driver
 ```
 
 ## License

playgrounds/beanie/README.md

@@ -80,7 +80,7 @@ Stop the database when you are done:
 ./scripts/stop-documentdb.sh
 ```
 
-The suite should end with `Passed: 13  Failed: 0`.
+The suite should end with `Passed: 16  Failed: 0`.
 
 ## Trying the API
 
@@ -175,7 +175,7 @@ Verified against `documentdb-local:latest` (release `0.114`):
 | Index creation via `Settings.indexes`   | ✅ Supported  | Built asynchronously by the engine; `createIndexes` returns in ~2s. Avoid `collation`. |
 | Unique indexes                          | ✅ Supported  | Duplicate keys raise `DuplicateKeyError` (code `11000`).              |
 | Aggregation pipelines                   | ✅ Common stages | `$match`, `$group`, `$unwind`, `$sort`, etc. Atlas-only stages differ. |
-| `$vectorSearch`                         | ❌ Not supported | Atlas-only operator.                                                |
+| `$vectorSearch` / vector search         | ✅ Supported  | DocumentDB supports a `cosmosSearch` vector index (e.g. `vector-ivf`) queried via the `$vectorSearch` stage. The test suite creates one and runs a nearest-neighbor query. |
 | Index `collation`                       | ❌ Not supported | `createIndex.collation is not implemented yet`; omit it.           |
 | Change streams / transactions           | ⚠️ Check version | Verify against your DocumentDB version before relying on them.     |
 
@@ -212,9 +212,12 @@ Beanie DocumentDB compatibility test
   ✅ aggregation ($unwind/$group)
   ✅ unique index enforcement (duplicate sku rejected)
   ✅ delete_one
+  ✅ vector index + insert (cosmosSearch vector-ivf)
+  ✅ $vectorSearch returns nearest neighbor
+  ✅ vector cleanup (drop collection)
   ✅ cleanup (drop collection)
 ====================================
-Passed: 13  Failed: 0
+Passed: 16  Failed: 0
 ```
 
 ## What the Scripts Do
@@ -229,7 +232,7 @@ Passed: 13  Failed: 0
 
 ## Verification
 
-- `./scripts/run-test.sh` ends with `Passed: 13  Failed: 0`.
+- `./scripts/run-test.sh` ends with `Passed: 16  Failed: 0`.
 - With `./scripts/run-app.sh` running, `curl http://localhost:3000/health`
   returns `{"status":"healthy","db":"connected"}`, `POST /books` returns `201`
   with the created document, and `GET /stats/genres` returns per-genre counts.

playgrounds/beanie/app/beanie_crud_test.py

@@ -217,6 +217,70 @@ async def _delete_one() -> None:
 
     await step("delete_one", _delete_one)
 
+    # Vector search: DocumentDB supports a `cosmosSearch` vector index and the
+    # `$vectorSearch` aggregation stage. The index is created with the raw
+    # `createIndexes` command since Beanie does not model vector indexes.
+    db_ref = client[DB_NAME]
+    vec_name = f"vectors_{int(time.time() * 1000)}"
+    vectors = db_ref[vec_name]
+
+    async def _vector_index_insert() -> None:
+        await vectors.insert_many(
+            [
+                {"name": "a", "v": [1, 0, 0]},
+                {"name": "b", "v": [0.9, 0.1, 0]},
+                {"name": "c", "v": [0, 0, 1]},
+            ]
+        )
+        res = await db_ref.command(
+            {
+                "createIndexes": vec_name,
+                "indexes": [
+                    {
+                        "name": "v_ivf",
+                        "key": {"v": "cosmosSearch"},
+                        "cosmosSearchOptions": {
+                            "kind": "vector-ivf",
+                            "numLists": 1,
+                            "similarity": "COS",
+                            "dimensions": 3,
+                        },
+                    }
+                ],
+            }
+        )
+        if res.get("ok") != 1:
+            raise RuntimeError("createIndexes did not return ok:1")
+
+    await step("vector index + insert (cosmosSearch vector-ivf)", _vector_index_insert)
+
+    async def _vector_search() -> None:
+        hits = await vectors.aggregate(
+            [
+                {
+                    "$vectorSearch": {
+                        "index": "v_ivf",
+                        "path": "v",
+                        "queryVector": [1, 0, 0],
+                        "numCandidates": 10,
+                        "limit": 2,
+                    }
+                },
+                {"$project": {"name": 1, "_id": 0}},
+            ]
+        ).to_list(length=10)
+        if not hits:
+            raise RuntimeError("vector search returned no results")
+        if hits[0].get("name") != "a":
+            raise RuntimeError(f"expected nearest 'a', got {hits[0].get('name')!r}")
+
+    await step("$vectorSearch returns nearest neighbor", _vector_search)
+
+    async def _vector_cleanup() -> None:
+        await vectors.drop()
+
+    await step("vector cleanup (drop collection)", _vector_cleanup)
+
     async def _drop() -> None:
         await coll.drop()