[sohailshk]

Fixes #41233

The output functions of MLTransform embeddings created with columns= (_dict_output_fn), of the RAG embedding adapters (_add_embedding_fn) and of VertexAIMultiModalEmbeddings (_multimodal_dict_output_fn) assigned the embeddings onto the input elements. The programming guide warns that "Changes to input elements may be visible to other DoFns executing within the same bundle, and may cause correctness issues", and that happens here: another consumer of the same PCollection (for example beam.combiners.ToList()) received the embedded values. The returned element has to be a different object from the input, so each function now returns a shallow copy (copy.copy for the dict, dataclasses.replace for the items and the nested Vertex image/video/text objects); the values themselves are not copied.

Tests: one per function (the Vertex one uses a stand-in response, no GCP); each fails on master. pytest apache_beam/ml/transforms/base_test.py apache_beam/ml/rag/embeddings/base_test.py apache_beam/ml/transforms/embeddings/vertex_ai_test.py: 40 passed, 12 skipped (TFT tests not installed). Two-branch pipeline (input into MLTransform with a test handler that reverses the text, and into ToList()): on master the ToList branch received cba/def for the inputs abc/def on the FnApiRunner and cba/fed on the default runner (Prism); which elements are hit depends on how the runner fuses the stages. After the change it receives abc/def on both.

[github-actions]

📋 Review Summary

This PR resolves #41233 by preventing in-place mutation of input elements in _add_embedding_fn, _dict_output_fn, and _multimodal_dict_output_fn. Each output function now shallow-copies input dictionaries (copy.copy) and dataclass instances (dataclasses.replace) before attaching embeddings, upholding Apache Beam's element immutability contract across sibling transforms.

🔍 General Feedback

• The use of copy.copy (preserving dict subclasses) and dataclasses.replace on EmbeddableItem, VertexImage, and VertexVideo avoids unnecessary deep copies of underlying payloads while cleanly isolating output mutations.
• Unit tests thoroughly cover all three output functions and verify that input references and nested objects remain unchanged.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

Assigning reviewers:

R: @tvalentyn for label python.
Backup reviewer: shunping (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).