Skip to content

Commit

Permalink
Implement $vectorSearch stage builder
Browse files Browse the repository at this point in the history
  • Loading branch information
GromNaN committed Nov 26, 2024
1 parent df9daca commit e433932
Show file tree
Hide file tree
Showing 8 changed files with 504 additions and 2 deletions.
2 changes: 1 addition & 1 deletion generator/config/search/queryString.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -32,4 +32,4 @@ tests:
-
$project:
_id: 0
title: 1
title: 1
2 changes: 1 addition & 1 deletion generator/config/search/wildcard.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -57,4 +57,4 @@ tests:
-
$project:
_id: 0
title: 1
title: 1
115 changes: 115 additions & 0 deletions generator/config/stage/vectorSearch.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
# $schema: ../schema.json
name: $vectorSearch
link: 'https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/'
type:
- stage
encode: object
description: |
Skips the first n documents where n is the specified skip number and passes the remaining documents unmodified to the pipeline. For each input document, outputs either zero documents (for the first n documents) or one document (if after the first n documents).
arguments:
-
name: index
type:
- string
-
name: limit
type:
- int
-
name: path
type:
- searchPath
-
name: queryVector
type:
- array # of numbers
-
name: exact
optional: true
type:
- bool
-
name: filter
optional: true
type:
- query
-
name: numCandidates
optional: true
type:
- int
tests:
-
name: 'ANN Basic'
link: 'https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/#ann-examples'
pipeline:
-
$vectorSearch:
index: 'vector_index'
path: 'plot_embedding'
queryVector:
- -0.0016261312
- -0.028070757
- -0.011342932
# skip other numbers, not relevant to the test
numCandidates: 150
limit: 10
-
$project:
_id: 0
plot: 1
title: 1
score:
$meta: 'vectorSearchScore'

-
name: 'ANN Filter'
link: 'https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/#ann-examples'
pipeline:
-
$vectorSearch:
index: 'vector_index'
path: 'plot_embedding'
filter:
$and:
-
year:
$lt: 1975
queryVector:
- 0.02421053
- -0.022372592
- -0.006231137
# skip other numbers, not relevant to the test
numCandidates: 150
limit: 10
-
$project:
_id: 0
title: 1
plot: 1
year: 1
score:
$meta: 'vectorSearchScore'

-
name: 'ENN'
link: 'https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/#enn-examples'
pipeline:
-
$vectorSearch:
index: 'vector_index'
path: 'plot_embedding'
queryVector:
- -0.006954097
- -0.009932499
- -0.001311474
# skip other numbers, not relevant to the test
exact: true
limit: 10
-
$project:
_id: 0
plot: 1
title: 1
score:
$meta: 'vectorSearchScore'
24 changes: 24 additions & 0 deletions src/Builder/Stage/FactoryTrait.php

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

26 changes: 26 additions & 0 deletions src/Builder/Stage/FluentFactoryTrait.php

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

93 changes: 93 additions & 0 deletions src/Builder/Stage/VectorSearchStage.php

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit e433932

Please sign in to comment.