Repository navigation
feat: add Amazon S3 Vectors document store - #3918
oleksii-udovychenko wants to merge 15 commits into
Conversation
|
Hi @oleksii-udovychenko, thanks for your interest in contributing to Haystack! 🙏 This is an automated message to help us keep the review queue healthy. |
|
This isn't a duplicate. It's an adoption of #3149, which has been stalled since May. @dotKokott's 12 commits are preserved with authorship intact and rebased onto I can't push to #3149: GitHub doesn't allow retargeting a PR's head branch, and I'd need @dotKokott to add me as a collaborator on his fork. Happy to close this if he'd rather finish it himself. |
|
@oleksii-udovychenko yes feel free to continue this! I will not get around to it |
|
@anakin87, just a gentle nudge to review this as it has been sitting open for 2 weeks. Happy to address any feedback asap. |
anakin87
left a comment
There was a problem hiding this comment.
Thank you for your work!
I left a few comments.
Feel also free to refresh this PR based on new S3 Vectors capabilities and align it a bit to this repo style.
| PYTHONUNBUFFERED: "1" | ||
| FORCE_COLOR: "1" | ||
| TEST_MATRIX_OS: '["ubuntu-latest", "windows-latest", "macos-latest"]' | ||
| TEST_MATRIX_PYTHON: '["3.10", "3.14"]' |
There was a problem hiding this comment.
We now use 3.11-3.15 please align
There was a problem hiding this comment.
Done, the workflow now matches the current template and tests on 3.11 and 3.15.
| @@ -0,0 +1,108 @@ | |||
| # Amazon S3 Vectors — Haystack Integration Example | |||
There was a problem hiding this comment.
we recently prefer to put examples in documentation, so please remove this file
There was a problem hiding this comment.
Removed.
| dynamic = ["version"] | ||
| description = "Haystack integration for amazon_s3_vectors" | ||
| readme = "README.md" | ||
| requires-python = ">=3.10" |
There was a problem hiding this comment.
Done: requires-python = ">=3.11", classifiers from 3.11 to 3.15.
| "Programming Language :: Python :: Implementation :: CPython", | ||
| "Programming Language :: Python :: Implementation :: PyPy", | ||
| ] | ||
| dependencies = ["haystack-ai>=2.26.1", "boto3>=1.42.0"] |
There was a problem hiding this comment.
let's use at least haystack-ai>=3.0.0
There was a problem hiding this comment.
Done, haystack-ai>=3.0.0.
| # S3 Vectors allows up to 500 vectors per put_vectors call | ||
| _WRITE_BATCH_SIZE = 500 | ||
|
|
||
| # S3 Vectors allows up to 500 keys per delete_vectors call | ||
| _DELETE_BATCH_SIZE = 500 | ||
|
|
||
| # Maximum number of results from query_vectors | ||
| _MAX_TOP_K = 100 | ||
|
|
||
| # S3 Vectors metadata limits | ||
| _MAX_TOTAL_METADATA_BYTES = 40 * 1024 # 40 KB total |
There was a problem hiding this comment.
Could you please verify whether these limits still hold? For example, I see that max top_k is 10,000 now
In general, please align this PR to the most recent S3 Vectors behaviors
There was a problem hiding this comment.
I checked them against the current S3 Vectors docs and botocore. topK now goes up to 10,000 and results come back in pages of up to 100, so _MAX_TOP_K is 10,000 and _embedding_retrieval follows nextToken.
Pagination needs boto3 1.43.31, so I raised the floor. The other limits are unchanged: 500 vectors per PutVectors/DeleteVectors, 40 KB metadata per vector, 2 KB filterable. I left the new metadata prefiltering (queryMode, ENHANCED indexes) out of this PR.
| """ | ||
| Return the number of documents in the document store. | ||
|
|
||
| .. note:: |
There was a problem hiding this comment.
We don't use ".. note::"/".. warning::" syntax. Please align to the repository style
There was a problem hiding this comment.
Replaced them with plain Note: lines.
| if filters: | ||
| _validate_filters(filters) |
There was a problem hiding this comment.
let's validate filters at the beginning
There was a problem hiding this comment.
Done, filters are validated before anything is listed.
| Return documents matching the provided filters. | ||
|
|
||
| .. warning:: | ||
|
|
There was a problem hiding this comment.
- Add a warning that filter_documents() filters locally using Haystack, while the retriever filters through S3. The same filter can match different documents, especially with array metadata.
- For this reason, and until S3 does not provide a filtering API, I recommend removing all optional methods that use
filter_documents:delete_by_filter,update_by_filter,count_documents_by_filter,count_unique_metadata_by_filter,get_metadata_fields_info,get_metadata_field_min_max, andget_metadata_field_unique_values
There was a problem hiding this comment.
Removed all seven methods, with their helpers and tests, and added the note to filter_documents. The S3 Vectors docs confirm the array case: $eq on an array field matches when any element matches, while Haystack compares the whole list.
| filters=filters, | ||
| top_k=top_k, | ||
| ) | ||
| return {"documents": docs} |
There was a problem hiding this comment.
let's add a close method here that invokes document_store.close
| WRITE_SLEEP_SECONDS = 5 | ||
| DELETE_SLEEP_SECONDS = 5 |
There was a problem hiding this comment.
I am not sure that these are still needed. Could you please verify?
There was a problem hiding this comment.
They aren't needed anymore. I put a vector and then listed, got and queried it, and deleted it and then listed and queried again, 20 times each on a new index. All 100 reads saw the new state immediately, so I removed the sleeps, and the integration suite now runs in 81 s. That index was ENHANCED, since the bucket was new; I didn't test a CLASSIC one.
Implements issue deepset-ai#2110 - Amazon S3 Vectors document store integration with: - S3VectorsDocumentStore: full DocumentStore protocol (count, write, filter, delete) - S3VectorsEmbeddingRetriever: embedding-based retrieval with metadata filtering - Filter conversion from Haystack format to S3 Vectors filter syntax - Auto-creation of vector buckets and indexes - AWS credential support via Secret (or default credential chain) - 49 unit tests covering store, retriever, filters, and serialization - README with usage examples and known limitations
…rkflow - boto3 lower bound set to 1.42.0 (when s3vectors service was added) - pydoc filename changed to amazon_s3_vectors.md (underscores, matching folder name) - Quote $GITHUB_OUTPUT in workflow to fix shellcheck SC2086
- Flatten test classes into standalone functions (matching pinecone/qdrant pattern) - Assert full serialized dict structure in to_dict/from_dict tests - Use Mock(spec=...) for retriever tests instead of MagicMock+patch - Verify _embedding_retrieval call args match exactly - Add test_from_dict_no_filter_policy (backward compat) - Add test_init_is_lazy
Remove tests that just verify mock plumbing (count, write, delete calling the mock client). Keep tests that verify our actual logic: - Serialization roundtrip (full dict structure) - Score conversion (cosine + euclidean) - Filter conversion (pure function with real logic) - Duplicate policy batch checks (SKIP/NONE) - Document <-> S3 vector conversion - Input validation Before: 49 unit tests (many testing mock behavior) After: 26 unit tests (all testing our code) + 12 integration tests
- Class docstring: top_k cap, dimension limit, metadata limits, float32 only - write_documents: embedding required, 40KB metadata limit - _embedding_retrieval: top_k=100 cap, no embeddings in response - Retriever run: top_k=100, server-side filters, no embeddings returned
…ity, deduplicate retrieval logic - Replace hand-rolled _apply_filters_in_memory/_document_matches/_compare with haystack.utils.filters.document_matches_filter (same utility used by InMemoryDocumentStore). Gains NOT operator, nested dotted field paths, and date comparison support for free. (-65 lines) - Deduplicate blob/content reconstruction in _embedding_retrieval() by reusing _s3_vector_to_document() + dataclasses.replace() (-20 lines) - Make filter_documents() warning conditional on filters actually being provided (no warning when listing all documents)
Matches the pattern used by the amazon_bedrock workflow: - top-level id-token: write permission - AWS_REGION env var - configure-aws-credentials step (skipped on fork PRs and dependabot) - integration tests gated on successful auth
Matches the repo convention used across other integrations.
Trim the README down to badges + integration page / changelog links + the AWS auth note, matching the pattern used by amazon_bedrock, anthropic, qdrant, and other integrations in this repo. The rich usage / limitations content is preserved for the haystack.deepset.ai integration page. Addresses deepset-ai#3149 (comment)
…ments Move type / embedding validation out of the per-batch loop and into a single pass before any `put_vectors` call, so a bad input never leaves the store with a partial write. Also align with Haystack's duplicate-policy convention: `DuplicatePolicy.FAIL` is now supported and, together with `DuplicatePolicy.NONE`, raises `DuplicateDocumentError` (was `DocumentStoreError` for NONE) when any target id already exists. Unit tests updated accordingly. Addresses deepset-ai#3149 (comment)
…ion tests Replace the handcrafted integration suite with Haystack's standard DocumentStore test contract, following the pgvector / pinecone pattern: * tests/conftest.py: shared `document_store` fixture. One vector bucket per session, one index per test for isolation. Wraps `write_documents` and `delete_documents` to (a) inject a default zero embedding for any Document missing one, since S3 Vectors requires embeddings, and (b) sleep briefly afterwards to absorb eventual consistency. * test_document_store.py: appends `TestDocumentStore` integration class inheriting CountDocumentsTest, WriteDocumentsTest, DeleteDocumentsTest, FilterableDocsFixtureMixin. Overrides `assert_documents_are_equal` for float32 round-trip tolerance. * test_filters.py: appends `TestFilters` inheriting FilterDocumentsTest. `filter_documents` already routes matching through haystack's document_matches_filter, so no operators are skipped. * test_integration.py: trimmed to the retriever-specific tests not covered by the base mixins (embedding retrieval, retriever component, to_dict/from_dict roundtrip on a live store). Coverage goes from 12 handcrafted integration tests to 53 (10 store + 39 filter + 4 retriever). Addresses deepset-ai#3149 (comment)
- `write_documents` now only supports `DuplicatePolicy.OVERWRITE`. `put_vectors` is an
upsert and S3 Vectors has no cheap existence check — `get_vectors` is capped at 100
keys and the only alternative is a full `ListVectors` scan — so the per-batch
`get_vectors` pass is dropped rather than paying for a read on every write. Other
policies are ignored with a warning.
- Documents without an embedding are stored with a non-zero placeholder vector, as
`PineconeDocumentStore` does, and reported in one warning before any write happens
instead of raising. Haystack's `filterable_docs` fixture ships embedding-less
documents on purpose, so refusing them makes `FilterDocumentsTest` unsatisfiable.
The caller's `Document` is not mutated.
- The test fixture no longer back-fills embeddings, which removes the
`warnings.simplefilter("ignore")` that was hiding the store's own warnings. It now
only absorbs S3 Vectors' eventual consistency.
- Add `delete_all_documents()`, `delete_by_filter()` and `update_by_filter()`, and mix
in `DeleteAllTest`, `DeleteByFilterTest` and `UpdateByFilterTest` so the integration
suite covers everything in `DocumentStoreBaseExtendedTests`. `FilterDocumentsTest`
stays in `test_filters.py`, as in pgvector.
- Factor the `list_vectors` pagination into `_iter_vectors()` instead of writing a
third copy of the loop.
- Move the retriever integration tests into `test_embedding_retriever.py` and delete
`test_integration.py`; its `to_dict`/`from_dict` roundtrip is already covered by the
unit tests.
- Drop the two decorative comment banners in the test files.
Verified against a real S3 Vectors bucket in us-east-1: 59 passed, 4 skipped — the two
unsupported duplicate policies, plus the two `recreate_index` variants that self-skip.
Unit tests, `fmt-check` and mypy are green.
Implement the whole optional DocumentStore extension surface Pinecone has, and mix in every matching test from haystack.testing.document_store: - count_documents_by_filter, count_unique_metadata_by_filter, get_metadata_fields_info, get_metadata_field_min_max and get_metadata_field_unique_values, with CountDocumentsByFilterTest, CountUniqueMetadataByFilterTest, GetMetadataFieldsInfoTest, GetMetadataFieldMinMaxTest and GetMetadataFieldUniqueValuesTest. - close(), which releases the boto3 client's connection pool. The store stays usable: the next call lazily builds a new client. - The store-specific Pinecone tests too (close-and-reopen, sentence-window retriever, metadata-introspection edge cases). Two deliberate divergences from Pinecone: no TOP_K_LIMIT truncation warnings, because list_vectors paginates over the whole index and these counts are exact; and get_metadata_field_unique_values keeps differently typed values distinct, so the base mixin's distinct-types test passes unmodified rather than needing Pinecone's override. Fix two metadata bugs that the new sentence-window test surfaced. S3 Vectors accepts a metadata value only if it is a string, a boolean, a finite number, or a non-empty array of either all strings or all non-boolean numbers; declaring a key non-filterable does not relax this. Anything else is a ValidationException from PutVectors, so: - Unstorable values are now dropped with a warning instead of failing the write, as PineconeDocumentStore does. This was not theoretical: DocumentSplitter with split_overlap adds a _split_overlap list of dicts to every chunk, so DocumentSplitter -> S3VectorsDocumentStore died with a ParamValidationError. Tuples are converted to lists, which the service accepts but botocore's document-type validator does not. - doc.blob.meta is an arbitrary dict, so writing any document with blob metadata failed. It is now JSON-encoded on write and decoded on read.
- Test on Python 3.11 and 3.15, require Python >= 3.11 and haystack-ai >= 3.0.0 - Raise the top_k cap to 10,000 and follow QueryVectors pagination, which needs boto3 >= 1.43.31 - Stop storing Document.blob - Validate filters before listing in filter_documents and note that local filtering can differ from S3 Vectors filtering on arrays - Remove the optional methods built on filter_documents - Add S3VectorsEmbeddingRetriever.close - Remove the example file and the eventual-consistency sleeps in the test fixtures - Fix the retriever usage example for Haystack 3, where the Sentence Transformers embedders live in sentence-transformers-haystack - Note in the class docstring that search is approximate, so a top_k close to the number of stored or matching documents returns fewer
4079306 to
51d92c3
Compare
oleksii-udovychenko
left a comment
There was a problem hiding this comment.
Thanks for the review @anakin87. I rebased on main and addressed all comments, with replies in each thread. A few changes go beyond the comments:
boto3>=1.43.31, the first release withQueryVectorspagination, which the 10,000top_kneeds.- The usage example in the retriever docstring still imported the Sentence Transformers embedders from Haystack core, so it now uses
sentence-transformers-haystack. - The class docstring no longer uses
**bold**. - Writing a document with a
bloblogs a warning, since the blob is no longer stored.
The integration tests pass locally against AWS: 55 passed, 4 skipped (the duplicate-policy and recreate-index tests that don't apply). Fork PRs can't use the CI role, so they're skipped in CI.
I also ran the store and the retriever end to end in us-east-1, on a new bucket (so the index is ENHANCED), with 10,000 documents carrying random 768-dimension embeddings and a category field with four values. In every query the query document came first with a score of 1.0, scores were in descending order, and no document appeared twice. Asking for nearly the whole index, or for every match of a filter, returns about 80% of it. Raw boto3 query_vectors returns the same counts (8,053 of 10,000, and 2,005 of 2,500 with the filter), so that comes from the approximate search, not from the pagination:
| Operation | Result |
|---|---|
write_documents, 10,000 docs |
19.9 s, 20 put_vectors calls |
count_documents |
10,000, 3.4 s |
retriever, top_k=10 and 100 |
10 and 100 docs, 1 page each, 0.3 s |
retriever, top_k=1000 |
1,000 docs, 10 pages, 1.4 s |
retriever, top_k=10000 |
8,024 docs, 81 pages, 10.5 s |
retriever, top_k=10000, category == "paper" (2,500 matches) |
2,022 docs, all paper, 21 pages, 2.8 s |
filter_documents, rank < 100 |
100 docs, 12.6 s |
delete_documents, 5,000 ids |
5,000 left |
delete_all_documents |
0 left |
| PYTHONUNBUFFERED: "1" | ||
| FORCE_COLOR: "1" | ||
| TEST_MATRIX_OS: '["ubuntu-latest", "windows-latest", "macos-latest"]' | ||
| TEST_MATRIX_PYTHON: '["3.10", "3.14"]' |
There was a problem hiding this comment.
Done, the workflow now matches the current template and tests on 3.11 and 3.15.
| @@ -0,0 +1,108 @@ | |||
| # Amazon S3 Vectors — Haystack Integration Example | |||
There was a problem hiding this comment.
Removed.
| dynamic = ["version"] | ||
| description = "Haystack integration for amazon_s3_vectors" | ||
| readme = "README.md" | ||
| requires-python = ">=3.10" |
There was a problem hiding this comment.
Done: requires-python = ">=3.11", classifiers from 3.11 to 3.15.
| "Programming Language :: Python :: Implementation :: CPython", | ||
| "Programming Language :: Python :: Implementation :: PyPy", | ||
| ] | ||
| dependencies = ["haystack-ai>=2.26.1", "boto3>=1.42.0"] |
There was a problem hiding this comment.
Done, haystack-ai>=3.0.0.
| # S3 Vectors allows up to 500 vectors per put_vectors call | ||
| _WRITE_BATCH_SIZE = 500 | ||
|
|
||
| # S3 Vectors allows up to 500 keys per delete_vectors call | ||
| _DELETE_BATCH_SIZE = 500 | ||
|
|
||
| # Maximum number of results from query_vectors | ||
| _MAX_TOP_K = 100 | ||
|
|
||
| # S3 Vectors metadata limits | ||
| _MAX_TOTAL_METADATA_BYTES = 40 * 1024 # 40 KB total |
There was a problem hiding this comment.
I checked them against the current S3 Vectors docs and botocore. topK now goes up to 10,000 and results come back in pages of up to 100, so _MAX_TOP_K is 10,000 and _embedding_retrieval follows nextToken.
Pagination needs boto3 1.43.31, so I raised the floor. The other limits are unchanged: 500 vectors per PutVectors/DeleteVectors, 40 KB metadata per vector, 2 KB filterable. I left the new metadata prefiltering (queryMode, ENHANCED indexes) out of this PR.
| """ | ||
| Return the number of documents in the document store. | ||
|
|
||
| .. note:: |
There was a problem hiding this comment.
Replaced them with plain Note: lines.
| if filters: | ||
| _validate_filters(filters) |
There was a problem hiding this comment.
Done, filters are validated before anything is listed.
| Return documents matching the provided filters. | ||
|
|
||
| .. warning:: | ||
|
|
There was a problem hiding this comment.
Removed all seven methods, with their helpers and tests, and added the note to filter_documents. The S3 Vectors docs confirm the array case: $eq on an array field matches when any element matches, while Haystack compares the whole list.
| filters=filters, | ||
| top_k=top_k, | ||
| ) | ||
| return {"documents": docs} |
| WRITE_SLEEP_SECONDS = 5 | ||
| DELETE_SLEEP_SECONDS = 5 |
There was a problem hiding this comment.
They aren't needed anymore. I put a vector and then listed, got and queried it, and deleted it and then listed and queried again, 20 times each on a new index. All 100 reads saw the new state immediately, so I removed the sleeps, and the integration suite now runs in 81 s. That index was ENHANCED, since the bucket was new; I didn't test a CLASSIC one.
|
@anakin87 friendly ping: I addressed your review comments and pushed the changes on October 15. Could you take another look when you have time? |
Supersedes #3149. Fixes #2110.
This continues @dotKokott's work on the Amazon S3 Vectors document store. His 12 commits are preserved unchanged (rebased onto current
main, authorship intact); the commits on top address the outstanding review feedback from @anakin87.@dotKokott, happy to hand this back or close it if you'd rather finish it yourself.
Review feedback addressed
DuplicatePolicy.OVERWRITEas in PineconeGetVectors(100 keys/call) or a fullListVectorsscan, and no conditional put. So theget_vectorspass is gone and onlyOVERWRITEis supported; other policies are ignored with a warningsimplefilter("ignore")DocumentStoreBaseExtendedTestsdelete_all_documents(),delete_by_filter()andupdate_by_filter(), and mixed inDeleteAllTest,DeleteByFilterTest,UpdateByFilterTest. Also the rest of the optional surface, see belowtest_embedding_retriever.pytest_integration.pydeletedsuggestionblocks deleting comment bannersBrought up to the same coverage as
PineconeDocumentStoreRather than stop at
DocumentStoreBaseExtendedTests, the store now implements the whole optional extension surface Pinecone has, and mixes in every matching test fromhaystack.testing.document_store:count_documents_by_filter(),count_unique_metadata_by_filter(),get_metadata_fields_info(),get_metadata_field_min_max(),get_metadata_field_unique_values(), covered byCountDocumentsByFilterTest,CountUniqueMetadataByFilterTest,GetMetadataFieldsInfoTest,GetMetadataFieldMinMaxTestandGetMetadataFieldUniqueValuesTest.close(), which releases the boto3 client's connection pool. The store stays usable, since the next call lazily builds a new client.test_close_and_reopen,test_sentence_window_retriever, and the metadata-introspection edge cases (consistent types, boolean/string min-max, list-valued fields).Two deliberate divergences from Pinecone:
TOP_K_LIMITtruncation warnings.list_vectorspaginates over the whole index, so these counts are exact rather than capped.get_metadata_field_unique_values()keeps values of different types distinct (the int1, the float1.0, the str"1"andTrue), so the base mixin'stest_get_metadata_field_unique_values_distinct_typespasses unmodified. Pinecone has to override it because it coerces numeric metadata tofloat.asyncis not implemented. Pinecone's 14*_asyncmethods would needaiobotocoreas a new runtime dependency; that felt like a separate PR rather than something to slip into this one.Behaviour changes
1.
write_documentsno longer raises on a document without an embedding. It stores a non-zero placeholder vector and warns, the same thingPineconeDocumentStoredoes with its_dummy_vector. Two reasons:filterable_docsfixture deliberately includes documents with no embedding, so a store that rejects them can't satisfyFilterDocumentsTestwithout overriding most of it.cosinedistance rejects zero-norm vectors outright (ValidationException: cosine distance does not support vectors with zero norm).filter_documents()strips the placeholder again, so such a document reads back withembedding=Nonerather than with a vector the caller never supplied. The caller'sDocumentis never mutated, since the substitution happens in_document_to_s3_vector.2. Unstorable metadata values are dropped with a warning instead of failing the write. S3 Vectors accepts a metadata value only if it is a string, a boolean, a finite number, or a non-empty array of either all strings or all non-boolean numbers. I probed the live API to pin this down exactly. All of the following are a
ValidationExceptionfromPutVectors:NoneMetadata values must be strings, numbers, booleans, or arrays{"a": 1}[]Empty arrays are not allowed in metadata[1, "x"]Metadata array values must be strings or numbers[True, False][{"doc_id": ..., "range": [0, 24]}]NaN/infDeclaring a key non-filterable does not relax any of this. I checked that separately.
This is not a theoretical concern.
DocumentSplitterwithsplit_overlap > 0adds a_split_overlaplist-of-dicts to every chunk, so before this change a plainDocumentSplitter → S3VectorsDocumentStorepipeline died with aParamValidationError. It's whattest_sentence_window_retrievercaught. Pinecone has the same class of problem and solves it the same way (_discard_invalid_meta).Two related fixes fell out of the same probe:
doc.blob.metais an arbitrary dict, so writing any document with blob metadata failed. It is now JSON-encoded on the way out and decoded on the way back in.list.Testing
The integration job is skipped on fork PRs, so I ran everything locally against a real S3 Vectors bucket in
us-east-1:test_write_documents_duplicate_fail/_skip(unsupported policies) and the tworecreate_indexvariants oftest_delete_all_documents, which self-skip when the store'sdelete_all_documentshas no such parameter.hatch run fmt-checkandhatch run test:typesgreen as well.Pipelines, real Bedrock Titan embeddings and no mocks:DocumentSplitter → AmazonBedrockDocumentEmbedder → DocumentWriter,AmazonBedrockTextEmbedder → S3VectorsEmbeddingRetriever, in-pipeline metadata filtering, and aPipeline.dumps()/loads()round-trip. Semantic ranking came out right (0.35 for the on-topic document vs 0.01 and −0.04 for the others).PutVectorsAPI to establish exactly which metadata values it accepts, which is where the two behaviour changes above come from.One thing I can't do myself: the workflow reuses the repo's existing
secrets.AWS_CI_ROLE_ARN(the same secretamazon_bedrock.ymluses), but that role needs S3 Vectors permissions before the integration job can pass in CI. Could a maintainer adds3vectors:CreateVectorBucket,CreateIndex,PutVectors,ListVectors,QueryVectors,GetVectors,DeleteVectors,DeleteIndex,DeleteVectorBucket,ListIndexes,GetIndexandGetVectorBucket? The suite creates one vector bucket per session and tears it down afterwards.