deepset-ai / deepset-ai/haystack

InMemoryDocumentStore.embedding_retrieval() ignores the store-level return_embedding default

Open Beginner friendly
#12,803 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
26.6k
Forks
3.2k
Avg merge
1d 3h
Merged PRs (30d)
194

Description

Describe the bug

On a default-constructed InMemoryDocumentStore, embedding_retrieval() returns documents with their
embedding stripped, even though the store is documented to return embeddings by default, and even
though the two sibling read paths on the very same instance (filter_documents() and
bm25_retrieval()) do return them.

InMemoryDocumentStore.__init__ takes return_embedding: bool = True and documents it as "Whether to
return the embedding of the retrieved Documents. Default is True." (document_store.py:74, :98).
embedding_retrieval documents the same fallback again at :809-811: "If not provided, the value of
the return_embedding parameter set at component initialization will be used."

That second promise never holds. The parameter is declared
return_embedding: bool | None = False (document_store.py:800) but the fallback is written as
resolved_return_embedding = self.return_embedding if return_embedding is None else return_embedding
(:846). Omitting the argument yields False, not None, so the is None branch is unreachable
through normal use: the store-level value is consulted only if the caller explicitly passes None.
embedding_retrieval_async is a second instance of the same wiring: it is annotated plain
return_embedding: bool = False (:1077, no | None) and forwards that concrete False to the sync
method (:1093-1099), so the async path cannot reach the fallback at all.

The bug is invisible in pipelines, which is presumably why it has survived: InMemoryEmbeddingRetriever.run()
resolves the sentinel itself and always passes a concrete bool (embedding_retriever.py:171-180). It is
visible to anyone calling the store method directly, and to integrations that do.

Error message

No exception. Silent data loss: Document.embedding is None in the returned documents.

Expected behavior

embedding_retrieval() / embedding_retrieval_async() should agree with the store-level
return_embedding value and with filter_documents() / bm25_retrieval() when the caller does not
specify an override.

To Reproduce

from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore

QUERY_EMB = [1.0, 0.0, 0.0]

def build_store():
    store = InMemoryDocumentStore(embedding_similarity_function="cosine")
    store.write_documents([
        Document(content="x", embedding=[1.0, 0.0, 0.0]),
        Document(content="y", embedding=[0.0, 1.0, 0.0]),
    ])
    return store

store = build_store()
print(store.return_embedding)                                        # True (documented default)
print([d.embedding for d in store.filter_documents()])                # embeddings present
print([d.embedding for d in store.bm25_retrieval(query="x")])         # embeddings present
print([d.embedding for d in store.embedding_retrieval(QUERY_EMB)])    # [None, None]  <-- unexpected
print([d.embedding for d in store.embedding_retrieval(QUERY_EMB, return_embedding=None)])  # present

Actual output on python3 -u repro.py:

store.return_embedding (documented default) = True
filter_documents()          -> embeddings: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
bm25_retrieval('x')         -> embeddings: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
embedding_retrieval(qe)     -> embeddings: [None, None]
embedding_retrieval(qe, None) -> embeddings: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
embedding_retrieval_async(qe) -> embeddings: [None, None]

No API keys, no network, no document store server: this is the in-memory store only.

Possible solutions

Two coherent resolutions, and they are not equivalent, so I would rather get a call than guess:

  1. Make the sentinel real: default return_embedding to None in embedding_retrieval (:800) and
    embedding_retrieval_async (:1077, with the bool | None annotation), so :846 does what both
    docstrings say. Verified locally: this makes the two methods agree with filter_documents() and
    bm25_retrieval(). It does change behavior for a caller who constructs the store with
    return_embedding=True and omits the argument — they start receiving embeddings, which is what
    __init__ advertises. InMemoryEmbeddingRetriever and any other caller that passes the argument
    explicitly are unaffected, and test/document_stores/test_in_memory.py is unchanged by it
    (344 passed / 8 skipped both before and after).
  2. If stripping on omission was the intent, then the is None fallback at :846 is dead code and the
    two docstring sentences that promise it should go, along with the | None in the signature.

Option 1 looks like the original intent: #9622 (59403de1f) introduced the store attribute, the
bool | None annotation, the docstring sentence and the is None resolution in one commit, and only
the default value contradicts the rest.

Additional context

Reproduced against current main (b717d0065); haystack/document_stores/in_memory/document_store.py
is byte-identical to the commit I first ran this on, and none of the intervening upstream commits touch
that file. Searched open and closed issues for return_embedding and embedding_retrieval: the only
matches are the older non-InMemory reports (#7037 is MongoDBAtlasDocumentStore under 1.x) and the
feature PR #9622 itself; no open pull request changes either method.

Related but distinct, and not what this issue is about: #12653 (metadata aliasing when
return_embedding=False in filter_documents).

Happy to send the PR for option 1, with a sync/async regression test pair, if that direction works.

FAQ Check

  • I have checked the FAQ and my question is not answered there.

System:

  • OS: macOS (arm64)
  • GPU/CPU: CPU only, no model involved
  • Haystack version: 3.2.0-rc0 (source checkout, editable install; the affected file is identical to main at b717d0065)
  • Python: 3.12.8
  • hatch not used locally; the repro and the tests were run with pytest directly against the source tree.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in haystack/document_stores/in_memory/document_store.py, inspecting the sync and async embedding_retrieval signatures and their return_embedding resolution. Run the focused tests in test/document_stores/test_in_memory.py, then add sync and async regression coverage for omitted and explicit return_embedding values. Done means both methods honor the store default while explicit overrides still work.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
databases
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
82/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.